rignet

app.py

@@ -14,7 +14,6 @@ if RIGNET_PATH not in sys.path:
     sys.path.append(RIGNET_PATH)
 
 # Ensure binvox executable is available in the root (required by RigNet's os.system call)
-# RigNet code expects './binvox' in the current directory.
 BINVOX_SRC = os.path.join(RIGNET_PATH, "binvox")
 if platform.system() == "Windows":
     BINVOX_SRC += ".exe"
@@ -29,10 +28,9 @@ if os.path.exists(BINVOX_SRC):
     if platform.system() != "Windows":
         os.system(f"chmod +x {BINVOX_DEST}")
 else:
-    print(f"Warning: binvox not found at {BINVOX_SRC}. Voxelization may fail.")
+    print(f"Warning: binvox not found at {BINVOX_SRC}. Inference may fail.")
 
 # --- 2. Import RigNet Modules ---
-# These imports must happen AFTER adding RigNet to sys.path
 try:
     from quick_start import (
         create_single_data, predict_joints, predict_skeleton, 
@@ -43,7 +41,7 @@ try:
     from models.PairCls_GCN import PairCls as BONENET
     from models.SKINNING import SKINNET
 except ImportError as e:
-    print("Error importing RigNet modules. Ensure the 'RigNet' folder is present and contains the code.")
+    print("Error importing RigNet modules. Ensure the 'RigNet' folder is correctly placed.")
     raise e
 
 # --- 3. Load Models Globally ---
@@ -53,6 +51,8 @@ print(f"Loading RigNet models on {device}...")
 def load_checkpoint(model, filename):
     # Checkpoints are located inside RigNet/checkpoints/
     filepath = os.path.join(RIGNET_PATH, "checkpoints", filename)
+    if not os.path.exists(filepath):
+        raise FileNotFoundError(f"Checkpoint not found: {filepath}")
     checkpoint = torch.load(filepath, map_location=device)
     model.load_state_dict(checkpoint['state_dict'])
     return model
@@ -84,89 +84,79 @@ def rignet_inference(input_mesh_path):
     if not input_mesh_path:
         return None
 
-    # Setup temporary paths
-    # Gradio stores inputs in a specific hash folder. 
-    # We work in that folder or create a clean temp one to avoid side effects.
+    # Work in a temp folder derived from the input path
     working_dir = os.path.dirname(input_mesh_path)
     base_name = os.path.basename(input_mesh_path).replace(".obj", "")
     
-    # Prepare filenames (mimicking quick_start naming convention)
+    # Prepare filenames
     mesh_filename = os.path.join(working_dir, f"{base_name}_remesh.obj")
-    mesh_original_path = input_mesh_path
     
     print(f"Processing: {input_mesh_path}")
 
-    # 1. Preprocessing / Decimation
-    mesh_ori = o3d.io.read_triangle_mesh(mesh_original_path)
-    if len(np.asarray(mesh_ori.vertices)) == 0:
-        raise ValueError("Empty mesh uploaded.")
-
-    # Simplify mesh as per RigNet requirement (approx 4k vertices for inference)
-    mesh_remesh = mesh_ori.simplify_quadric_decimation(4000)
-    o3d.io.write_triangle_mesh(mesh_filename, mesh_remesh)
-
-    # 2. Create Data (Voxelization + Geodesic calculation)
-    # create_single_data calls ./binvox internally. We ensured it exists in root.
-    data, vox, surface_geodesic, translation_normalize, scale_normalize = create_single_data(mesh_filename)
-    data.to(device)
-
-    # 3. Pipeline Predictions
-    # Config (using default learned parameters from quick_start)
-    bandwidth = 0.0429 
-    threshold = 1e-5
-    
-    print("Predicting joints...")
-    # Note: We pass the normalized obj path which create_single_data generated
-    mesh_norm_path = mesh_filename.replace("_remesh.obj", "_normalized.obj")
-    
-    data = predict_joints(
-        data, vox, jointNet, threshold, bandwidth=bandwidth,
-        mesh_filename=mesh_norm_path
-    )
-    data.to(device)
-
-    print("Predicting connectivity...")
-    pred_skeleton = predict_skeleton(
-        data, vox, rootNet, boneNet,
-        mesh_filename=mesh_norm_path
-    )
-
-    print("Predicting skinning...")
-    # downsample_skinning=True for speed (as per quick_start default)
-    pred_rig = predict_skinning(
-        data, pred_skeleton, skinNet, surface_geodesic,
-        mesh_norm_path, subsampling=True
-    )
-
-    # 4. Post-processing
-    # Reverse normalization
-    pred_rig.normalize(scale_normalize, -translation_normalize)
-
-    # Transfer rig to original high-res mesh
-    final_rig = tranfer_to_ori_mesh(mesh_original_path, mesh_filename, pred_rig)
-
-    # Save output
-    output_path = os.path.join(working_dir, f"{base_name}_rig.txt")
-    final_rig.save(output_path)
-    
-    # Optional: Return the simplified mesh for visualization if needed, 
-    # but here we just return the RigNet text file output.
-    return output_path
+    try:
+        # 1. Preprocessing / Decimation
+        mesh_ori = o3d.io.read_triangle_mesh(input_mesh_path)
+        if len(np.asarray(mesh_ori.vertices)) == 0:
+            raise ValueError("Empty mesh uploaded.")
+
+        # Simplify mesh (approx 4k vertices for inference)
+        mesh_remesh = mesh_ori.simplify_quadric_decimation(4000)
+        o3d.io.write_triangle_mesh(mesh_filename, mesh_remesh)
+
+        # 2. Create Data (Voxelization + Geodesic calculation)
+        data, vox, surface_geodesic, translation_normalize, scale_normalize = create_single_data(mesh_filename)
+        data = data.to(device)
+
+        # 3. Pipeline Predictions
+        # Default parameters from quick_start
+        bandwidth = 0.0429 
+        threshold = 1e-5
+        
+        mesh_norm_path = mesh_filename.replace("_remesh.obj", "_normalized.obj")
+        
+        print("Predicting joints...")
+        data = predict_joints(
+            data, vox, jointNet, threshold, bandwidth=bandwidth,
+            mesh_filename=mesh_norm_path
+        )
+        data = data.to(device)
+
+        print("Predicting connectivity...")
+        pred_skeleton = predict_skeleton(
+            data, vox, rootNet, boneNet,
+            mesh_filename=mesh_norm_path
+        )
+
+        print("Predicting skinning...")
+        pred_rig = predict_skinning(
+            data, pred_skeleton, skinNet, surface_geodesic,
+            mesh_norm_path, subsampling=True
+        )
+
+        # 4. Post-processing
+        pred_rig.normalize(scale_normalize, -translation_normalize)
+        final_rig = tranfer_to_ori_mesh(input_mesh_path, mesh_filename, pred_rig)
+
+        # Save output to a text file
+        output_path = os.path.join(working_dir, f"{base_name}_rig.txt")
+        final_rig.save(output_path)
+        
+        return output_path
+
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        raise gr.Error(f"Processing failed: {str(e)}")
 
 # --- 5. Gradio Interface ---
-title = "RigNet: Neural Rigging for Articulated Characters"
-description = (
-    "Upload a 3D model (.obj) to automatically generate a skeleton and skinning weights. "
-    "The output is a text file containing the rig information."
-)
+title = "RigNet: Neural Rigging"
+description = "Upload an .obj file. The model will generate a skeleton and skinning weights, returned as a text file."
 
 iface = gr.Interface(
     fn=rignet_inference,
     inputs=gr.Model3D(label="Input Mesh (.obj)"),
-    outputs=gr.File(label="Download Rig Output (.txt)"),
+    outputs=gr.File(label="Download Rig (.txt)"),
     title=title,
-    description=description,
-    examples=[], # Add examples if you have sample .obj files in a folder
+    description=description
 )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,30 +1,13 @@
-# Core dependencies
-gradio==4.16.0
+--find-links https://data.pyg.org/whl/torch-1.12.0+cpu.html
 torch==1.12.0
 torchvision==0.13.0
-
-# PyTorch Geometric and related packages
+gradio
+numpy
+scipy
+trimesh
+open3d
 torch-geometric==1.7.2
 torch-scatter
 torch-sparse
 torch-cluster
 torch-spline-conv
-
-# Scientific computing
-numpy==1.21.6
-scipy==1.7.3
-scikit-learn
-
-# 3D processing
-trimesh[easy]==3.21.7
-open3d==0.9.0
-rtree>=0.8,<0.9
-
-# Visualization and utilities
-matplotlib==3.5.3
-opencv-python==4.7.0.72
-tensorboard==2.11.2
-Pillow==9.5.0
-
-# File handling
-pathlib