Update RigNet/quick_start.py

RigNet/quick_start.py

@@ -51,6 +51,7 @@ def normalize_obj(mesh_v):
     mesh_v *= scale
     return mesh_v, pivot, scale
 
+
 def create_single_data(mesh_filename):
     """
     create input data for the network. The data is wrapped by Data structure in pytorch-geometric library
@@ -64,8 +65,7 @@ def create_single_data(mesh_filename):
     mesh_f = np.asarray(mesh.triangles)
     mesh_v, translation_normalize, scale_normalize = normalize_obj(mesh_v)
     mesh_normalized = o3d.geometry.TriangleMesh(vertices=o3d.utility.Vector3dVector(mesh_v), triangles=o3d.utility.Vector3iVector(mesh_f))
-    normalized_obj = mesh_filename.replace("_remesh.obj", "_normalized.obj")
-    o3d.io.write_triangle_mesh(normalized_obj, mesh_normalized)
+    o3d.io.write_triangle_mesh(mesh_filename.replace("_remesh.obj", "_normalized.obj"), mesh_normalized)
 
     # vertices
     v = np.concatenate((mesh_v, mesh_vn), axis=1)
@@ -90,153 +90,37 @@ def create_single_data(mesh_filename):
     # batch
     batch = torch.zeros(len(v), dtype=torch.long)
 
-    # voxel - Use trimesh instead of binvox (no X11 required)
+    # voxel - FIXED: Use absolute path and better error handling
     binvox_file = mesh_filename.replace('_remesh.obj', '_normalized.binvox')
+    normalized_obj = mesh_filename.replace("_remesh.obj", "_normalized.obj")
     
     if not os.path.exists(binvox_file):
-        print(f"  voxelizing mesh (trimesh-based, no binvox)...")
-        
-        try:
-            # Load mesh with trimesh
-            mesh_tri = trimesh.load(normalized_obj)
-            
-            # Voxelize: create a 88x88x88 grid
-            # Calculate pitch to fit mesh in 88^3 grid
-            bounds = mesh_tri.bounds
-            dims = bounds[1] - bounds[0]
-            max_dim = max(dims)
-            pitch = max_dim / 88.0
-            
-            # Create voxel grid
-            voxel_grid = mesh_tri.voxelized(pitch=pitch)
-            
-            # Get current voxel matrix
-            vox_matrix = voxel_grid.matrix
-            current_shape = vox_matrix.shape
-            
-            print(f"  Original voxel shape: {current_shape}")
-            
-            # Resize to exactly 88x88x88 by padding/cropping each dimension
-            target_shape = (88, 88, 88)
-            resized = np.zeros(target_shape, dtype=bool)
-            
-            # Calculate how much to copy in each dimension
-            x_size = min(current_shape[0], target_shape[0])
-            y_size = min(current_shape[1], target_shape[1])
-            z_size = min(current_shape[2], target_shape[2])
-            
-            # Copy the overlapping region
-            resized[:x_size, :y_size, :z_size] = vox_matrix[:x_size, :y_size, :z_size]
-            
-            vox_matrix = resized
-            print(f"  Resized voxel shape: {vox_matrix.shape}")
-            
-            # Create binvox-compatible object with ALL required attributes
-            class Voxels:
-                def __init__(self, data, dims, translate, scale, axis_order):
-                    self.data = data
-                    self.dims = dims
-                    self.translate = translate
-                    self.scale = scale
-                    self.axis_order = axis_order
-            
-            vox_obj = Voxels(
-                data=vox_matrix,
-                dims=[88, 88, 88],
-                translate=[0.0, 0.0, 0.0],
-                scale=1.0,
-                axis_order='xyz'
-            )
-            
-            # Save as binvox format for caching
-            with open(binvox_file, 'wb') as f:
-                binvox_rw.write(vox_obj, f)
-            
-            print(f"  ✓ Voxelization complete: {binvox_file}")
-            
-        except Exception as e:
-            print(f"  ERROR: Trimesh voxelization failed: {e}")
-            import traceback
-            traceback.print_exc()
-            raise Exception(f"Voxelization failed: {e}")
-    
-    # Load voxel data
-    with open(binvox_file, 'rb') as fvox:
-        vox = binvox_rw.read_as_3d_array(fvox)
-
-    data = Data(x=v[:, 3:6], pos=v[:, 0:3], tpl_edge_index=tpl_e, geo_edge_index=geo_e, batch=batch)
-    return data, vox, surface_geodesic, translation_normalize, scale_normalize
-
-
-# def create_single_data(mesh_filename):
-#     """
-#     create input data for the network. The data is wrapped by Data structure in pytorch-geometric library
-#     :param mesh_filename: name of the input mesh
-#     :return: wrapped data, voxelized mesh, and geodesic distance matrix of all vertices
-#     """
-#     mesh = o3d.io.read_triangle_mesh(mesh_filename)
-#     mesh.compute_vertex_normals()
-#     mesh_v = np.asarray(mesh.vertices)
-#     mesh_vn = np.asarray(mesh.vertex_normals)
-#     mesh_f = np.asarray(mesh.triangles)
-#     mesh_v, translation_normalize, scale_normalize = normalize_obj(mesh_v)
-#     mesh_normalized = o3d.geometry.TriangleMesh(vertices=o3d.utility.Vector3dVector(mesh_v), triangles=o3d.utility.Vector3iVector(mesh_f))
-#     o3d.io.write_triangle_mesh(mesh_filename.replace("_remesh.obj", "_normalized.obj"), mesh_normalized)
-
-#     # vertices
-#     v = np.concatenate((mesh_v, mesh_vn), axis=1)
-#     v = torch.from_numpy(v).float()
-
-#     # topology edges
-#     print("  gathering topological edges.")
-#     tpl_e = get_tpl_edges(mesh_v, mesh_f).T
-#     tpl_e = torch.from_numpy(tpl_e).long()
-#     tpl_e, _ = add_self_loops(tpl_e, num_nodes=v.size(0))
-
-#     # surface geodesic distance matrix
-#     print("  calculating surface geodesic matrix.")
-#     surface_geodesic = calc_surface_geodesic(mesh)
-
-#     # geodesic edges
-#     print("  gathering geodesic edges.")
-#     geo_e = get_geo_edges(surface_geodesic, mesh_v).T
-#     geo_e = torch.from_numpy(geo_e).long()
-#     geo_e, _ = add_self_loops(geo_e, num_nodes=v.size(0))
-
-#     # batch
-#     batch = torch.zeros(len(v), dtype=torch.long)
-
-#     # voxel - FIXED: Use absolute path and better error handling
-#     binvox_file = mesh_filename.replace('_remesh.obj', '_normalized.binvox')
-#     normalized_obj = mesh_filename.replace("_remesh.obj", "_normalized.obj")
-    
-#     if not os.path.exists(binvox_file):
-#         print(f"  voxelizing mesh with binvox...")
+        print(f"  voxelizing mesh with binvox...")
         
-#         # Use absolute path to binvox (installed in Dockerfile)
-#         if platform == "linux" or platform == "linux2":
-#             cmd = f"binvox -d 88 -pb {normalized_obj}"
-#         elif platform == "win32":
-#             cmd = f"binvox.exe -d 88 {normalized_obj}"
-#         else:
-#             raise Exception('Sorry, we currently only support windows and linux.')
+        # Use absolute path to binvox (installed in Dockerfile)
+        if platform == "linux" or platform == "linux2":
+            cmd = f"binvox -d 88 -pb {normalized_obj}"
+        elif platform == "win32":
+            cmd = f"binvox.exe -d 88 {normalized_obj}"
+        else:
+            raise Exception('Sorry, we currently only support windows and linux.')
         
-#         print(f"  Running: {cmd}")
-#         exit_code = os.system(cmd)
+        print(f"  Running: {cmd}")
+        exit_code = os.system(cmd)
         
-#         if exit_code != 0:
-#             raise Exception(f"binvox command failed with exit code {exit_code}. Command: {cmd}")
+        if exit_code != 0:
+            raise Exception(f"binvox command failed with exit code {exit_code}. Command: {cmd}")
         
-#         if not os.path.exists(binvox_file):
-#             raise Exception(f"binvox did not create output file: {binvox_file}")
+        if not os.path.exists(binvox_file):
+            raise Exception(f"binvox did not create output file: {binvox_file}")
         
-#         print(f"  ✓ Voxelization complete: {binvox_file}")
+        print(f"  ✓ Voxelization complete: {binvox_file}")
     
-#     with open(binvox_file, 'rb') as fvox:
-#         vox = binvox_rw.read_as_3d_array(fvox)
+    with open(binvox_file, 'rb') as fvox:
+        vox = binvox_rw.read_as_3d_array(fvox)
 
-#     data = Data(x=v[:, 3:6], pos=v[:, 0:3], tpl_edge_index=tpl_e, geo_edge_index=geo_e, batch=batch)
-#     return data, vox, surface_geodesic, translation_normalize, scale_normalize
+    data = Data(x=v[:, 3:6], pos=v[:, 0:3], tpl_edge_index=tpl_e, geo_edge_index=geo_e, batch=batch)
+    return data, vox, surface_geodesic, translation_normalize, scale_normalize
 
 
 def predict_joints(input_data, vox, joint_pred_net, threshold, bandwidth=None, mesh_filename=None):
@@ -512,97 +396,97 @@ if __name__ == '__main__':
     # Change to False to be more accurate but less efficient.
     downsample_skinning = True
 
-    # load all weights
-    print("loading all networks...")
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-    jointNet = JOINTNET()
-    jointNet.to(device)
-    jointNet.eval()
-    jointNet_checkpoint = torch.load('checkpoints/gcn_meanshift/model_best.pth.tar')
-    jointNet.load_state_dict(jointNet_checkpoint['state_dict'])
-    print("     joint prediction network loaded.")
-
-    rootNet = ROOTNET()
-    rootNet.to(device)
-    rootNet.eval()
-    rootNet_checkpoint = torch.load('checkpoints/rootnet/model_best.pth.tar')
-    rootNet.load_state_dict(rootNet_checkpoint['state_dict'])
-    print("     root prediction network loaded.")
-
-    boneNet = BONENET()
-    boneNet.to(device)
-    boneNet.eval()
-    boneNet_checkpoint = torch.load('checkpoints/bonenet/model_best.pth.tar')
-    boneNet.load_state_dict(boneNet_checkpoint['state_dict'])
-    print("     connection prediction network loaded.")
-
-    skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True)
-    skinNet_checkpoint = torch.load('checkpoints/skinnet/model_best.pth.tar')
-    skinNet.load_state_dict(skinNet_checkpoint['state_dict'])
-    skinNet.to(device)
-    skinNet.eval()
-    print("     skinning prediction network loaded.")
-
-    # Here we provide 16~17 examples. For best results, we will need to override the learned bandwidth and its associated threshold
-    # To process other input characters, please first try the learned bandwidth (0.0429 in the provided model), and the default threshold 1e-5.
-    # We also use these two default parameters for processing all test models in batch.
-
-    #model_id, bandwidth, threshold = "smith", None, 1e-5
-    model_id, bandwidth, threshold = "17872", 0.045, 0.75e-5
-    #model_id, bandwidth, threshold = "8210", 0.05, 1e-5
-    #model_id, bandwidth, threshold = "8330", 0.05, 0.8e-5
-    #model_id, bandwidth, threshold = "9477", 0.043, 2.5e-5
-    #model_id, bandwidth, threshold = "17364", 0.058, 0.3e-5
-    #model_id, bandwidth, threshold = "15930", 0.055, 0.4e-5
-    #model_id, bandwidth, threshold = "8333", 0.04, 2e-5
-    #model_id, bandwidth, threshold = "8338", 0.052, 0.9e-5
-    #model_id, bandwidth, threshold = "3318", 0.03, 0.92e-5
-    #model_id, bandwidth, threshold = "15446", 0.032, 0.58e-5
-    #model_id, bandwidth, threshold = "1347", 0.062, 3e-5
-    #model_id, bandwidth, threshold = "11814", 0.06, 0.6e-5
-    #model_id, bandwidth, threshold = "2982", 0.045, 0.3e-5
-    #model_id, bandwidth, threshold = "2586", 0.05, 0.6e-5
-    #model_id, bandwidth, threshold = "8184", 0.05, 0.4e-5
-    #model_id, bandwidth, threshold = "9000", 0.035, 0.16e-5
-
-    # create data used for inferece
-    print("creating data for model ID {:s}".format(model_id))
-    mesh_filename = os.path.join(input_folder, '{:s}_remesh.obj'.format(model_id))
-    if not os.path.exists(mesh_filename):
-        mesh_ori_filename = os.path.join(input_folder, '{:s}_ori.obj'.format(model_id))
-        mesh_ori = o3d.io.read_triangle_mesh(mesh_ori_filename)
-        if len(np.asarray(mesh_ori.vertices)) == 0:
-            print(f"Please name your input model as {model_id}_ori.obj")
-            exit()
-        mesh_remesh = mesh_ori.simplify_quadric_decimation(4000)  # adjust vertices between 1K - 5K
-        o3d.io.write_triangle_mesh(mesh_filename, mesh_remesh)
-
-    data, vox, surface_geodesic, translation_normalize, scale_normalize = create_single_data(mesh_filename)
-    data.to(device)
-
-    print("predicting joints")
-    data = predict_joints(data, vox, jointNet, threshold, bandwidth=bandwidth,
-                          mesh_filename=mesh_filename.replace("_remesh.obj", "_normalized.obj"))
-    data.to(device)
-    print("predicting connectivity")
-    pred_skeleton = predict_skeleton(data, vox, rootNet, boneNet,
-                                     mesh_filename=mesh_filename.replace("_remesh.obj", "_normalized.obj"))
-    print("predicting skinning")
-    pred_rig = predict_skinning(data, pred_skeleton, skinNet, surface_geodesic,
-                                mesh_filename.replace("_remesh.obj", "_normalized.obj"),
-                                subsampling=downsample_skinning)
-
-    # here we reverse the normalization to the original scale and position
-    pred_rig.normalize(scale_normalize, -translation_normalize)
-
-    print("Saving result")
-    if True:
-        # here we use original mesh tesselation (without remeshing)
-        mesh_filename_ori = os.path.join(input_folder, '{:s}_ori.obj'.format(model_id))
-        pred_rig = tranfer_to_ori_mesh(mesh_filename_ori, mesh_filename, pred_rig)
-        pred_rig.save(mesh_filename_ori.replace('.obj', '_rig.txt'))
-    else:
-        # here we use remeshed mesh
-        pred_rig.save(mesh_filename.replace('.obj', '_rig.txt'))
-    print("Done!")
+    # # load all weights
+    # print("loading all networks...")
+    # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # jointNet = JOINTNET()
+    # jointNet.to(device)
+    # jointNet.eval()
+    # jointNet_checkpoint = torch.load('checkpoints/gcn_meanshift/model_best.pth.tar')
+    # jointNet.load_state_dict(jointNet_checkpoint['state_dict'])
+    # print("     joint prediction network loaded.")
+
+    # rootNet = ROOTNET()
+    # rootNet.to(device)
+    # rootNet.eval()
+    # rootNet_checkpoint = torch.load('checkpoints/rootnet/model_best.pth.tar')
+    # rootNet.load_state_dict(rootNet_checkpoint['state_dict'])
+    # print("     root prediction network loaded.")
+
+    # boneNet = BONENET()
+    # boneNet.to(device)
+    # boneNet.eval()
+    # boneNet_checkpoint = torch.load('checkpoints/bonenet/model_best.pth.tar')
+    # boneNet.load_state_dict(boneNet_checkpoint['state_dict'])
+    # print("     connection prediction network loaded.")
+
+    # skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True)
+    # skinNet_checkpoint = torch.load('checkpoints/skinnet/model_best.pth.tar')
+    # skinNet.load_state_dict(skinNet_checkpoint['state_dict'])
+    # skinNet.to(device)
+    # skinNet.eval()
+    # print("     skinning prediction network loaded.")
+
+    # # Here we provide 16~17 examples. For best results, we will need to override the learned bandwidth and its associated threshold
+    # # To process other input characters, please first try the learned bandwidth (0.0429 in the provided model), and the default threshold 1e-5.
+    # # We also use these two default parameters for processing all test models in batch.
+
+    # #model_id, bandwidth, threshold = "smith", None, 1e-5
+    # model_id, bandwidth, threshold = "17872", 0.045, 0.75e-5
+    # #model_id, bandwidth, threshold = "8210", 0.05, 1e-5
+    # #model_id, bandwidth, threshold = "8330", 0.05, 0.8e-5
+    # #model_id, bandwidth, threshold = "9477", 0.043, 2.5e-5
+    # #model_id, bandwidth, threshold = "17364", 0.058, 0.3e-5
+    # #model_id, bandwidth, threshold = "15930", 0.055, 0.4e-5
+    # #model_id, bandwidth, threshold = "8333", 0.04, 2e-5
+    # #model_id, bandwidth, threshold = "8338", 0.052, 0.9e-5
+    # #model_id, bandwidth, threshold = "3318", 0.03, 0.92e-5
+    # #model_id, bandwidth, threshold = "15446", 0.032, 0.58e-5
+    # #model_id, bandwidth, threshold = "1347", 0.062, 3e-5
+    # #model_id, bandwidth, threshold = "11814", 0.06, 0.6e-5
+    # #model_id, bandwidth, threshold = "2982", 0.045, 0.3e-5
+    # #model_id, bandwidth, threshold = "2586", 0.05, 0.6e-5
+    # #model_id, bandwidth, threshold = "8184", 0.05, 0.4e-5
+    # #model_id, bandwidth, threshold = "9000", 0.035, 0.16e-5
+
+    # # create data used for inferece
+    # print("creating data for model ID {:s}".format(model_id))
+    # mesh_filename = os.path.join(input_folder, '{:s}_remesh.obj'.format(model_id))
+    # if not os.path.exists(mesh_filename):
+    #     mesh_ori_filename = os.path.join(input_folder, '{:s}_ori.obj'.format(model_id))
+    #     mesh_ori = o3d.io.read_triangle_mesh(mesh_ori_filename)
+    #     if len(np.asarray(mesh_ori.vertices)) == 0:
+    #         print(f"Please name your input model as {model_id}_ori.obj")
+    #         exit()
+    #     mesh_remesh = mesh_ori.simplify_quadric_decimation(4000)  # adjust vertices between 1K - 5K
+    #     o3d.io.write_triangle_mesh(mesh_filename, mesh_remesh)
+
+    # data, vox, surface_geodesic, translation_normalize, scale_normalize = create_single_data(mesh_filename)
+    # data.to(device)
+
+    # print("predicting joints")
+    # data = predict_joints(data, vox, jointNet, threshold, bandwidth=bandwidth,
+    #                       mesh_filename=mesh_filename.replace("_remesh.obj", "_normalized.obj"))
+    # data.to(device)
+    # print("predicting connectivity")
+    # pred_skeleton = predict_skeleton(data, vox, rootNet, boneNet,
+    #                                  mesh_filename=mesh_filename.replace("_remesh.obj", "_normalized.obj"))
+    # print("predicting skinning")
+    # pred_rig = predict_skinning(data, pred_skeleton, skinNet, surface_geodesic,
+    #                             mesh_filename.replace("_remesh.obj", "_normalized.obj"),
+    #                             subsampling=downsample_skinning)
+
+    # # here we reverse the normalization to the original scale and position
+    # pred_rig.normalize(scale_normalize, -translation_normalize)
+
+    # print("Saving result")
+    # if True:
+    #     # here we use original mesh tesselation (without remeshing)
+    #     mesh_filename_ori = os.path.join(input_folder, '{:s}_ori.obj'.format(model_id))
+    #     pred_rig = tranfer_to_ori_mesh(mesh_filename_ori, mesh_filename, pred_rig)
+    #     pred_rig.save(mesh_filename_ori.replace('.obj', '_rig.txt'))
+    # else:
+    #     # here we use remeshed mesh
+    #     pred_rig.save(mesh_filename.replace('.obj', '_rig.txt'))
+    # print("Done!")