Update app.py

app.py

@@ -6,34 +6,50 @@ import torch
 import numpy as np
 import open3d as o3d
 import gradio as gr
+from torch_geometric.data import Data
+
+# --- 0. Compatibility Patch for PyG 2.0+ ---
+# RigNet uses older PyG where you could assign arbitrary attributes (data.joints = ...).
+# In PyG 2.0+, 'Data' stores everything in a store. We subclass to allow dot-notation access.
+class RigNetData(Data):
+    def __setattr__(self, key, value):
+        if key in ['joints', 'pairs', 'pair_attr', 'joints_batch', 'pairs_batch', 'skin_input']:
+            self[key] = value
+        else:
+            super().__setattr__(key, value)
+
+    def __getattr__(self, key):
+        # Fallback to getting from the dictionary if standard attribute fails
+        if key in self.keys:
+            return self[key]
+        return super().__getattr__(key)
+
+# Monkey-patch RigNet's create_single_data later to use this class, 
+# or just convert the object after creation.
 
 # --- 1. Setup Environment & Paths ---
-# Add RigNet submodule to python path so internal imports (e.g., 'from models import...') work
 RIGNET_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "RigNet")
 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)
+# Ensure binvox executable
 BINVOX_SRC = os.path.join(RIGNET_PATH, "binvox")
-if platform.system() == "Windows":
-    BINVOX_SRC += ".exe"
-    BINVOX_DEST = "binvox.exe"
-else:
-    BINVOX_DEST = "binvox"
+BINVOX_DEST = "binvox.exe" if platform.system() == "Windows" else "binvox"
+if platform.system() == "Windows": BINVOX_SRC += ".exe"
 
 if os.path.exists(BINVOX_SRC):
-    # Copy to root so ./binvox works
     shutil.copy(BINVOX_SRC, BINVOX_DEST)
-    # Make executable
-    if platform.system() != "Windows":
-        os.system(f"chmod +x {BINVOX_DEST}")
+    if platform.system() != "Windows": os.system(f"chmod +x {BINVOX_DEST}")
 else:
-    print(f"Warning: binvox not found at {BINVOX_SRC}. Inference may fail.")
+    print(f"Warning: binvox not found at {BINVOX_SRC}")
 
 # --- 2. Import RigNet Modules ---
 try:
+    # We need to intercept imports to inject our patched Data class if possible,
+    # but easier to just wrap the result of create_single_data.
     from quick_start import (
-        create_single_data, predict_joints, predict_skeleton, 
+        create_single_data as original_create_single_data, # Rename to wrap
+        predict_joints, predict_skeleton, 
         predict_skinning, tranfer_to_ori_mesh
     )
     from models.GCN import JOINTNET_MASKNET_MEANSHIFT as JOINTNET
@@ -41,122 +57,101 @@ 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 correctly placed.")
-    raise e
+    print(f"Error importing RigNet: {e}")
 
-# --- 3. Load Models Globally ---
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# --- 3. Load Models ---
+device = torch.device("cpu")
 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}")
+    # map_location=device is crucial for CPU-only spaces
     checkpoint = torch.load(filepath, map_location=device)
     model.load_state_dict(checkpoint['state_dict'])
     return model
 
 # Initialize models
-jointNet = JOINTNET().to(device)
-jointNet.eval()
+jointNet = JOINTNET().to(device); jointNet.eval()
 load_checkpoint(jointNet, 'gcn_meanshift/model_best.pth.tar')
 
-rootNet = ROOTNET().to(device)
-rootNet.eval()
+rootNet = ROOTNET().to(device); rootNet.eval()
 load_checkpoint(rootNet, 'rootnet/model_best.pth.tar')
 
-boneNet = BONENET().to(device)
-boneNet.eval()
+boneNet = BONENET().to(device); boneNet.eval()
 load_checkpoint(boneNet, 'bonenet/model_best.pth.tar')
 
-skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True).to(device)
-skinNet.eval()
+skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True).to(device); skinNet.eval()
 load_checkpoint(skinNet, 'skinnet/model_best.pth.tar')
 
-print("Models loaded successfully.")
+print("Models loaded.")
 
-# --- 4. Inference Pipeline ---
-def rignet_inference(input_mesh_path):
-    """
-    Main pipeline adapted from quick_start.py logic.
-    """
-    if not input_mesh_path:
-        return None
+# --- 4. Wrapper to fix Data object ---
+def create_single_data_patched(mesh_filename):
+    # Call original function
+    data, vox, surf_geo, t_norm, s_norm = original_create_single_data(mesh_filename)
+    
+    # Convert to our flexible class
+    new_data = RigNetData(
+        x=data.x, 
+        pos=data.pos, 
+        tpl_edge_index=data.tpl_edge_index, 
+        geo_edge_index=data.geo_edge_index, 
+        batch=data.batch
+    )
+    return new_data, vox, surf_geo, t_norm, s_norm
 
-    # Work in a temp folder derived from the input path
+# --- 5. Inference Pipeline ---
+def rignet_inference(input_mesh_path):
+    if not input_mesh_path: return None
+    
     working_dir = os.path.dirname(input_mesh_path)
     base_name = os.path.basename(input_mesh_path).replace(".obj", "")
-    
-    # Prepare filenames
     mesh_filename = os.path.join(working_dir, f"{base_name}_remesh.obj")
     
     print(f"Processing: {input_mesh_path}")
-
     try:
-        # 1. Preprocessing / Decimation
+        # 1. Preprocess
         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)
+        if len(np.asarray(mesh_ori.vertices)) == 0: raise ValueError("Empty mesh")
         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)
+        # 2. Create Data (Patched)
+        data, vox, surface_geodesic, t_norm, s_norm = create_single_data_patched(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")
+        # 3. Inference
+        mesh_norm = 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 = predict_joints(data, vox, jointNet, 1e-5, bandwidth=0.0429, mesh_filename=mesh_norm)
         data = data.to(device)
 
         print("Predicting connectivity...")
-        pred_skeleton = predict_skeleton(
-            data, vox, rootNet, boneNet,
-            mesh_filename=mesh_norm_path
-        )
+        skel = predict_skeleton(data, vox, rootNet, boneNet, mesh_filename=mesh_norm)
 
         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)
+        rig = predict_skinning(data, skel, skinNet, surface_geodesic, mesh_norm, subsampling=True)
+
+        # 4. Export
+        rig.normalize(s_norm, -t_norm)
+        final_rig = tranfer_to_ori_mesh(input_mesh_path, mesh_filename, rig)
         
-        return output_path
+        out_path = os.path.join(working_dir, f"{base_name}_rig.txt")
+        final_rig.save(out_path)
+        return out_path
 
     except Exception as e:
-        print(f"Error during inference: {e}")
+        print(f"Error: {e}")
         raise gr.Error(f"Processing failed: {str(e)}")
 
-# --- 5. Gradio Interface ---
-title = "RigNet: Neural Rigging"
-description = "Upload an .obj file. The model will generate a skeleton and skinning weights, returned as a text file."
-
+# --- 6. Launch ---
 iface = gr.Interface(
     fn=rignet_inference,
-    inputs=gr.Model3D(label="Input Mesh (.obj)"),
-    outputs=gr.File(label="Download Rig (.txt)"),
-    title=title,
-    description=description
+    inputs=gr.Model3D(label="Input .obj"),
+    outputs=gr.File(label="Rig Output .txt"),
+    title="RigNet Demo",
+    description="Upload a mesh to generate a rig. (Uses PyTorch 2.1 CPU)"
 )
 
 if __name__ == "__main__":