skinning/utils/visualize.py

# Modified from https://github.com/PeizhuoLi/neural-blend-shapes/blob/main/blender_scripts/vertex_color.py
import bpy
import numpy as np
import random
import os
import matplotlib.pyplot as plt

def generate_color_table(num_colors, base_color_path):
    base_colors = np.load(base_color_path)
    idx = list(range(base_colors.shape[0]))
    random.seed(5)
    random.shuffle(idx)
    pt = 0
    res = []
    for i in range(num_colors):
        res.append(base_colors[idx[pt]])
        pt += 1
        pt %= base_colors.shape[0]
    return np.array(res)

def weight2color(weight, colors):
    res = np.matmul(weight, colors)
    res = np.clip(res, 0, 1)
    return res

def export_obj_with_vertex_colors(obj_object, vertex_colors, filepath):
    mesh = obj_object.data

    with open(filepath, 'w') as f:
        f.write("# Exported OBJ with vertex colors\n")
        for v, color in zip(mesh.vertices, vertex_colors):
            f.write("v {} {} {} {} {} {}\n".format(
                v.co.x, v.co.y, v.co.z,
                color[0], color[1], color[2]
            ))
        
        f.write("g {}\n".format(obj_object.name))
        for poly in mesh.polygons:
            vertices = [v + 1 for v in poly.vertices]
            face = "f " + " ".join([str(v) for v in vertices]) + "\n"
            f.write(face)

def compute_error_map(weight_method, weight_gt):
    error = np.abs(weight_method - weight_gt)  # Shape: (n_vertex, n_color)
    error_map = np.max(error, axis=1)  # Taking the maximum error across all colors for each vertex
    # error_map = np.mean(error, axis=1) 
    # Normalize error_map to [0,1]
    error_map = np.clip(error_map, 0, 1)
    return error_map

def error_map_to_color(error_map, colormap='Reds'):
    cmap = plt.get_cmap(colormap)
    colors = cmap(error_map)[:, :3]  # Discard alpha channel
    return colors

def generate_color_bar(filepath, colormap='Reds'):
    fig, ax = plt.subplots(figsize=(8, 1.5))
    fig.subplots_adjust(bottom=0.5)

    cmap = plt.get_cmap(colormap)
    norm = plt.Normalize(vmin=0, vmax=1)
    cb1 = plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap),
                      cax=ax, orientation='horizontal')
    cb1.ax.tick_params(labelsize=26) 
    plt.savefig(filepath, bbox_inches='tight')
    plt.close()

def main():
    base_color_path = "utils/tableau_color.npy"

    ### xlv2_test
    obj_path = "xlv2_test_mesh" # save from npz files, check skeleton/data_utils/convert_npz_to_mesh_rig.py
    weight_path1 = "gt_onxlv2_npy"
    weight_path2 = "ours_onxlv2_npy"
    # weight_path3 = "rignet_onxlv2_npy"
    # weight_path4 = "gvb_onxlv2_npy"
    save_path = "obj_w_skincolor_errormap_onxlv2_test"
    
    color_bar_path = "error_color_bar.png" # Path to save the color bar image

    if not os.path.exists(save_path):
        os.makedirs(save_path)
        
    generate_color_bar(color_bar_path, colormap='Reds')

    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete(use_global=False)

    for obj_file in os.listdir(obj_path):
        obj_file_path = os.path.join(obj_path, obj_file)
        wt_gt_file = os.path.join(weight_path1, obj_file.replace(".obj", "_skin.npy"))
        wt_ours_file = os.path.join(weight_path2, obj_file.replace(".obj", "_skin.npy"))
        # wt_rignet_file = os.path.join(weight_path3, obj_file.replace(".obj", "_skin.npy"))
        # wt_gvb_file = os.path.join(weight_path4, obj_file.replace(".obj", "_skin.npy"))
        if not os.path.exists(wt_gt_file) or not os.path.exists(wt_ours_file): # or not os.path.exists(wt_rignet_file) or not os.path.exists(wt_gvb_file):
            continue

        # 导入OBJ文件
        bpy.ops.wm.obj_import(filepath=obj_file_path)
        imported_objects = bpy.context.selected_objects
        if not imported_objects:
            print("Failed to import OBJ file:", obj_file_path)
            continue
        obj = imported_objects[0]

        weight_gt = np.load(wt_gt_file)
        weight_ours = np.load(wt_ours_file)
        # weight_rignet = np.load(wt_rignet_file)
        # weight_gvb = np.load(wt_gvb_file)
        if weight_gt.shape != weight_ours.shape: # or weight_gt.shape != weight_rignet.shape or weight_gt.shape != weight_gvb.shape:
            print("Weight shape mismatch among files")
            bpy.ops.object.select_all(action='SELECT')
            bpy.ops.object.delete(use_global=False)
            continue

        if len(weight_gt.shape) != 2:
            print("Weight file should be 2D array")
            bpy.ops.object.select_all(action='SELECT')
            bpy.ops.object.delete(use_global=False)
            continue

        n_vertices = len(obj.data.vertices)
        if weight_gt.shape[0] != n_vertices:
            print(f"Vertex count mismatch: OBJ has {n_vertices}, but weight file has {weight_gt.shape[0]}")
            bpy.ops.object.select_all(action='SELECT')
            bpy.ops.object.delete(use_global=False)
            continue

        # generate colors
        n_color = weight_gt.shape[1]
        colors = generate_color_table(n_color, base_color_path)
        vertex_colors_gt = weight2color(weight_gt, colors)
        vertex_colors_ours = weight2color(weight_ours, colors)
        # vertex_colors_rignet = weight2color(weight_rignet, colors)
        # vertex_colors_gvb = weight2color(weight_gvb, colors)

        # Save obj with vertex colors
        output_path_gt = os.path.join(save_path, obj_file.replace(".obj", "_gt.obj"))
        output_path_ours = os.path.join(save_path, obj_file.replace(".obj", "_ours.obj"))
        # output_path_rignet = os.path.join(save_path, obj_file.replace(".obj", "_rignet.obj"))
        # output_path_gvb = os.path.join(save_path, obj_file.replace(".obj", "_gvb.obj"))
        export_obj_with_vertex_colors(obj, vertex_colors_gt, output_path_gt)
        export_obj_with_vertex_colors(obj, vertex_colors_ours, output_path_ours)
        # export_obj_with_vertex_colors(obj, vertex_colors_rignet, output_path_rignet)
        # export_obj_with_vertex_colors(obj, vertex_colors_gvb, output_path_gvb)

        # Save object with error color map
        # 1. Ours vs GT
        error_ours = compute_error_map(weight_ours, weight_gt)
        colors_error_ours = error_map_to_color(error_ours, colormap='Reds')
        output_path_error_ours = os.path.join(save_path, obj_file.replace(".obj", "_error_ours.obj"))
        export_obj_with_vertex_colors(obj, colors_error_ours, output_path_error_ours)

        # 2. Rignet vs GT
        # error_rignet = compute_error_map(weight_rignet, weight_gt)
        # colors_error_rignet = error_map_to_color(error_rignet, colormap='Reds')
        # output_path_error_rignet = os.path.join(save_path, obj_file.replace(".obj", "_error_rignet.obj"))
        # export_obj_with_vertex_colors(obj, colors_error_rignet, output_path_error_rignet)

        # 3. GVB vs GT
        # error_gvb = compute_error_map(weight_gvb, weight_gt)
        # colors_error_gvb = error_map_to_color(error_gvb, colormap='Reds')
        # output_path_error_gvb = os.path.join(save_path, obj_file.replace(".obj", "_error_gvb.obj"))
        # export_obj_with_vertex_colors(obj, colors_error_gvb, output_path_error_gvb)

        bpy.ops.object.select_all(action='SELECT')
        bpy.ops.object.delete(use_global=False)

    print("Done")

if __name__ == "__main__":
    main()