Update app.py

app.py

@@ -1,33 +1,124 @@
 #!/usr/bin/env python3
-
-import sys
-import os
-import gradio as gr
-import trimesh
-import numpy as np
+"""
+RigNet Gradio Demo for Hugging Face Spaces (CPU)
+"""
 import os
 import sys
 import tempfile
 import shutil
 import traceback
 from pathlib import Path
+
+# CRITICAL: Patch quick_start.py BEFORE importing
+# Fix binvox path issue - must happen before RigNet imports
+import_patch_code = '''
+import os
+from sys import platform
+
+# Monkey-patch quick_start.create_single_data to fix binvox paths
+original_create_single_data = None
+
+def patched_create_single_data(mesh_filename):
+    """Patched version that fixes binvox path issues"""
+    import open3d as o3d
+    import numpy as np
+    import torch
+    from torch_geometric.data import Data
+    from torch_geometric.utils import add_self_loops
+    sys.path.insert(0, '/app/RigNet')
+    from utils import binvox_rw
+    from gen_dataset import get_tpl_edges, get_geo_edges
+    from geometric_proc.common_ops import calc_surface_geodesic
+    from quick_start import normalize_obj
+    
+    # Load and normalize mesh
+    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)
+    
+    # Save normalized mesh
+    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)
+    
+    # Prepare data
+    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
+    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)
+    
+    # Voxelization - FIX THE PATH ISSUE HERE
+    binvox_file = normalized_obj.replace('.obj', '.binvox')
+    
+    if not os.path.exists(binvox_file):
+        print(f"  Creating voxel file: {binvox_file}")
+        # Use full path to binvox
+        cmd = f"/usr/local/bin/binvox -d 88 -pb {normalized_obj}"
+        print(f"  Running: {cmd}")
+        ret = os.system(cmd)
+        if ret != 0:
+            raise RuntimeError(f"binvox failed with return code {ret}")
+    
+    # Verify binvox file was created
+    if not os.path.exists(binvox_file):
+        raise FileNotFoundError(f"Binvox file not created: {binvox_file}")
+    
+    # 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
+'''
+
+# Execute the patch
+exec(import_patch_code)
+
+import gradio as gr
 import torch
+import numpy as np
 
 # Add RigNet to Python path
 sys.path.insert(0, '/app/RigNet')
 
 # Import RigNet modules
+from models.GCN import JOINTNET_MASKNET_MEANSHIFT as JOINTNET
+from models.ROOT_GCN import ROOTNET
+from models.PairCls_GCN import PairCls as BONENET
+from models.SKINNING import SKINNET
+
+# Import other functions from quick_start (we'll use our patched version)
 from quick_start import (
-    create_single_data,
     predict_joints,
     predict_skeleton,
     predict_skinning,
     normalize_obj
 )
-from models.GCN import JOINTNET_MASKNET_MEANSHIFT as JOINTNET
-from models.ROOT_GCN import ROOTNET
-from models.PairCls_GCN import PairCls as BONENET
-from models.SKINNING import SKINNET
 
 # Global variables for models
 device = torch.device("cpu")
@@ -112,11 +203,12 @@ def process_mesh(input_obj_path, bandwidth, threshold, downsample_skinning=True)
         shutil.copy(input_obj_path, mesh_filename)
         
         print(f"\nProcessing: {base_name}")
+        print(f"Working directory: {work_dir}")
         
-        # Step 1: Create data
+        # Step 1: Create data - USE PATCHED VERSION
         print("  [1/4] Creating input data...")
         data, vox, surface_geodesic, translation_normalize, scale_normalize = \
-            create_single_data(mesh_filename)
+            patched_create_single_data(mesh_filename)
         data.to(device)
         
         # Step 2: Predict joints
@@ -162,73 +254,32 @@ def process_mesh(input_obj_path, bandwidth, threshold, downsample_skinning=True)
 
 def rignet_inference(input_obj, bandwidth, threshold):
     """
-    Gradio inference function with extensive debugging
+    Gradio inference function
     """
     print("\n" + "="*60)
-    print("🔍 DEBUG: rignet_inference CALLED!")
+    print("🔍 rignet_inference CALLED!")
     print(f"   input_obj type: {type(input_obj)}")
-    print(f"   input_obj value: {input_obj}")
-    print(f"   bandwidth: {bandwidth}")
-    print(f"   threshold: {threshold}")
     
-    # Check if input is None or empty
     if input_obj is None:
-        msg = "⚠️ Please upload an OBJ file first"
-        print(f"   ERROR: {msg}")
-        print("="*60 + "\n")
-        return None, msg
+        return None, "⚠️ Please upload an OBJ file first"
     
     try:
-        # Ensure models are loaded
         load_models()
         
-        # Extract file path - handle multiple Gradio formats
+        # Extract file path
         input_path = None
-        
-        # Case 1: File object with .name attribute
         if hasattr(input_obj, 'name'):
             input_path = input_obj.name
-            print(f"   ✓ Got path from .name: {input_path}")
-        
-        # Case 2: Already a string path
         elif isinstance(input_obj, str):
             input_path = input_obj
-            print(f"   ✓ Already a string path: {input_path}")
-        
-        # Case 3: Dictionary with 'name' key
-        elif isinstance(input_obj, dict):
-            if 'name' in input_obj:
-                input_path = input_obj['name']
-                print(f"   ✓ Got path from dict['name']: {input_path}")
-            else:
-                print(f"   ERROR: Dict without 'name' key. Keys: {input_obj.keys()}")
-        
-        # Case 4: Unknown type - debug it
-        else:
-            print(f"   ERROR: Unknown input type!")
-            print(f"   Attributes: {dir(input_obj)}")
-            if hasattr(input_obj, '__dict__'):
-                print(f"   __dict__: {input_obj.__dict__}")
-            msg = f"❌ Unexpected file input type: {type(input_obj)}"
-            print("="*60 + "\n")
-            return None, msg
-        
-        # Validate file path
-        if not input_path:
-            msg = "❌ Could not extract file path from input"
-            print(f"   ERROR: {msg}")
-            print("="*60 + "\n")
-            return None, msg
+        elif isinstance(input_obj, dict) and 'name' in input_obj:
+            input_path = input_obj['name']
         
-        if not os.path.exists(input_path):
-            msg = f"❌ File does not exist: {input_path}"
-            print(f"   ERROR: {msg}")
-            print("="*60 + "\n")
-            return None, msg
+        if not input_path or not os.path.exists(input_path):
+            return None, f"❌ Invalid file path: {input_path}"
         
-        file_size = os.path.getsize(input_path)
-        print(f"   ✓ File validated: {input_path}")
-        print(f"   ✓ File size: {file_size:,} bytes")
+        print(f"   Processing: {input_path}")
+        print(f"   File size: {os.path.getsize(input_path):,} bytes")
         print("="*60 + "\n")
         
         # Process the mesh
@@ -239,147 +290,46 @@ def rignet_inference(input_obj, bandwidth, threshold):
             downsample_skinning=True
         )
         
-        # Validate output
         if not os.path.exists(output_rig_path):
-            msg = "❌ Output file was not created"
-            print(f"ERROR: {msg}")
-            return None, msg
+            return None, "❌ Output file was not created"
         
         output_size = os.path.getsize(output_rig_path)
         status_msg = f"✅ Rigging completed!\n\nFile: {os.path.basename(output_rig_path)}\nSize: {output_size:,} bytes"
         
-        print(f"✓ SUCCESS! Returning output file")
         return output_rig_path, status_msg
         
     except Exception as e:
-        error_msg = f"❌ Error during processing:\n\n{str(e)}\n\nDetails:\n{traceback.format_exc()}"
-        print("\n" + "="*60)
-        print("❌ EXCEPTION CAUGHT:")
+        error_msg = f"❌ Error:\n\n{str(e)}\n\n{traceback.format_exc()}"
         print(error_msg)
-        print("="*60 + "\n")
         return None, error_msg
-def process_obj_file(file_obj):
-    """
-    Process OBJ file and return first 10 lines of analysis results
-    """
-    sys.stdout.flush()
-    print(f"[DEBUG] Processing file: {file_obj.name if file_obj else 'None'}", flush=True)
-    
-    if not file_obj:
-        return "⚠️ No file provided"
-    
-    try:
-        results = []
-        results.append("="*60)
-        results.append("OBJ FILE ANALYSIS - First 10 Lines of Results")
-        results.append("="*60)
-        
-        # Read raw OBJ file first 10 lines
-        results.append("\n📄 RAW OBJ FILE (First 10 Lines):")
-        results.append("-"*60)
-        with open(file_obj.name, 'r') as f:
-            for i, line in enumerate(f):
-                if i >= 10:
-                    break
-                results.append(f"Line {i+1}: {line.rstrip()}")
-        
-        # Load mesh using trimesh
-        results.append("\n🔷 MESH ANALYSIS:")
-        results.append("-"*60)
-        
-        mesh = trimesh.load(file_obj.name, force='mesh')
-        
-        # Check if it's a Scene or Mesh
-        if isinstance(mesh, trimesh.Scene):
-            results.append(f"Type: Scene with {len(mesh.geometry)} geometries")
-            # Get the first geometry
-            if len(mesh.geometry) > 0:
-                first_geom_name = list(mesh.geometry.keys())[0]
-                mesh = mesh.geometry[first_geom_name]
-                results.append(f"Using first geometry: {first_geom_name}")
-        
-        # Mesh statistics (ensures we don't exceed 10 total result lines)
-        results.append(f"Vertices: {len(mesh.vertices)}")
-        results.append(f"Faces: {len(mesh.faces)}")
-        results.append(f"Is Watertight: {mesh.is_watertight}")
-        results.append(f"Is Winding Consistent: {mesh.is_winding_consistent}")
-        results.append(f"Bounds: {mesh.bounds.tolist()}")
-        results.append(f"Center Mass: {mesh.center_mass.tolist()}")
-        
-        # Join results
-        output = "\n".join(results[:25])  # Limit output
-        
-        print("[DEBUG] Processing completed successfully", flush=True)
-        return output
-        
-    except Exception as e:
-        error_msg = f"❌ Error processing file: {str(e)}\n\nStacktrace:\n{sys.exc_info()}"
-        print(error_msg, flush=True)
-        return error_msg
 
-# Gradio Interface
-# demo = gr.Interface(
-#     fn=process_obj_file,
-#     inputs=gr.File(
-#         label="Upload OBJ File",
-#         file_types=[".obj"],
-#         type="file"
-#     ),
-#     outputs=gr.Textbox(
-#         label="Analysis Results (First 10 Lines)",
-#         lines=20,
-#         max_lines=30
-#     ),
-#     title="🔷 OBJ File Analyzer",
-#     description="Upload a 3D OBJ file to see the first 10 lines of raw content and mesh analysis",
-#     examples=None,
-#     cache_examples=False
-# )
 
 if __name__ == "__main__":
-    print("="*60, flush=True)
-    print("🚀 Starting OBJ File Analyzer...", flush=True)
-    print("="*60, flush=True)
+    print("="*60)
+    print("RigNet Gradio Demo - Starting...")
+    print("="*60)
+    
     load_models()
-
-
+    
     demo = gr.Interface(
         fn=rignet_inference,
         inputs=[
             gr.File(label="Upload OBJ File", file_types=[".obj"], type="file"),
-            gr.Slider(0.02, 0.08, value=0.04, step=0.001, label="Bandwidth", info="Joint clustering density (default: 0.04)"),
-            gr.Slider(0.1, 3.0, value=1.0, step=0.1, label="Threshold (×10⁻⁵)", info="Joint filtering threshold (default: 1.0)")
+            gr.Slider(0.02, 0.08, value=0.04, step=0.001, label="Bandwidth"),
+            gr.Slider(0.1, 3.0, value=1.0, step=0.1, label="Threshold (×10⁻⁵)")
         ],
         outputs=[
             gr.File(label="Download Rig TXT"),
             gr.Textbox(label="Status", lines=5)
         ],
         title="🎭 RigNet: Neural Rigging for 3D Characters",
-        description="""
-Upload a 3D character mesh (OBJ format) to automatically generate skeletal rig and skinning weights.
-
-**Recommended:** OBJ files with 1K-5K vertices work best.
-**Processing time:** 1-3 minutes on CPU depending on mesh complexity.
-        """,
-        article="""
-### 📚 About the Output
-
-The generated `*_rig.txt` file contains:
-- **joints**: 3D positions of skeletal joints
-- **root**: Root joint of the hierarchy  
-- **hier**: Parent-child relationships (skeleton hierarchy)
-- **skin**: Skinning weights for each vertex
-
-This format can be imported into 3D animation software.
-
-**Reference:** [RigNet: Neural Rigging for Articulated Characters (SIGGRAPH 2020)](https://arxiv.org/abs/2005.00559)
-        """,
+        description="Upload OBJ (1K-5K vertices recommended). Processing takes 1-3 minutes on CPU.",
         allow_flagging="never"
     )
+    
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,
         show_error=True,
         debug=True
     )