Update app.py

app.py

@@ -1,124 +1,33 @@
 #!/usr/bin/env python3
-"""
-RigNet Gradio Demo for Hugging Face Spaces (CPU)
-"""
+
+import sys
+import os
+import gradio as gr
+import trimesh
+import numpy as np
 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")
@@ -148,7 +57,7 @@ def load_models():
         map_location=device
     )
     jointNet.load_state_dict(jointNet_checkpoint['state_dict'])
-    print("✓ Joint prediction network loaded")
+    print("✓ Joint prediction network loaded")
     
     # Root prediction network
     rootNet = ROOTNET()
@@ -159,7 +68,7 @@ def load_models():
         map_location=device
     )
     rootNet.load_state_dict(rootNet_checkpoint['state_dict'])
-    print("✓ Root prediction network loaded")
+    print("✓ Root prediction network loaded")
     
     # Bone connection network
     boneNet = BONENET()
@@ -170,7 +79,7 @@ def load_models():
         map_location=device
     )
     boneNet.load_state_dict(boneNet_checkpoint['state_dict'])
-    print("✓ Connectivity prediction network loaded")
+    print("✓ Connectivity prediction network loaded")
     
     # Skinning network
     skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True)
@@ -181,7 +90,7 @@ def load_models():
     skinNet.load_state_dict(skinNet_checkpoint['state_dict'])
     skinNet.to(device)
     skinNet.eval()
-    print("✓ Skinning prediction network loaded")
+    print("✓ Skinning prediction network loaded")
     
     models_loaded = True
     print("All models loaded successfully!\n")
@@ -203,12 +112,11 @@ 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 - USE PATCHED VERSION
+        # Step 1: Create data
         print("  [1/4] Creating input data...")
         data, vox, surface_geodesic, translation_normalize, scale_normalize = \
-            patched_create_single_data(mesh_filename)
+            create_single_data(mesh_filename)
         data.to(device)
         
         # Step 2: Predict joints
@@ -242,7 +150,7 @@ def process_mesh(input_obj_path, bandwidth, threshold, downsample_skinning=True)
         output_rig_path = os.path.join(work_dir, f'{base_name}_rig.txt')
         pred_rig.save(output_rig_path)
         
-        print(f"✓ Successfully generated rig: {base_name}_rig.txt\n")
+        print(f"✓ Successfully generated rig: {base_name}_rig.txt\n")
         
         return output_rig_path
         
@@ -254,32 +162,73 @@ def process_mesh(input_obj_path, bandwidth, threshold, downsample_skinning=True)
 
 def rignet_inference(input_obj, bandwidth, threshold):
     """
-    Gradio inference function
+    Gradio inference function with extensive debugging
     """
     print("\n" + "="*60)
-    print("🔍 rignet_inference CALLED!")
+    print("🔍 DEBUG: 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:
-        return None, "⚠️ Please upload an OBJ file first"
+        msg = "⚠️ Please upload an OBJ file first"
+        print(f"   ERROR: {msg}")
+        print("="*60 + "\n")
+        return None, msg
     
     try:
+        # Ensure models are loaded
         load_models()
         
-        # Extract file path
+        # Extract file path - handle multiple Gradio formats
         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
-        elif isinstance(input_obj, dict) and 'name' in input_obj:
-            input_path = input_obj['name']
+            print(f"   ✓ Already a string path: {input_path}")
         
-        if not input_path or not os.path.exists(input_path):
-            return None, f"❌ Invalid file 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()}")
         
-        print(f"   Processing: {input_path}")
-        print(f"   File size: {os.path.getsize(input_path):,} bytes")
+        # 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
+        
+        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
+        
+        file_size = os.path.getsize(input_path)
+        print(f"   ✓ File validated: {input_path}")
+        print(f"   ✓ File size: {file_size:,} bytes")
         print("="*60 + "\n")
         
         # Process the mesh
@@ -290,46 +239,147 @@ def rignet_inference(input_obj, bandwidth, threshold):
             downsample_skinning=True
         )
         
+        # Validate output
         if not os.path.exists(output_rig_path):
-            return None, "❌ Output file was not created"
+            msg = "❌ Output file was not created"
+            print(f"ERROR: {msg}")
+            return None, msg
         
         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"
+        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:\n\n{str(e)}\n\n{traceback.format_exc()}"
+        error_msg = f"❌ Error during processing:\n\n{str(e)}\n\nDetails:\n{traceback.format_exc()}"
+        print("\n" + "="*60)
+        print("❌ EXCEPTION CAUGHT:")
         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)
-    print("RigNet Gradio Demo - Starting...")
-    print("="*60)
-    
+    print("="*60, flush=True)
+    print("🚀 Starting OBJ File Analyzer...", flush=True)
+    print("="*60, flush=True)
     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"),
-            gr.Slider(0.1, 3.0, value=1.0, step=0.1, label="Threshold (×10⁻⁵)")
+            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)")
         ],
         outputs=[
             gr.File(label="Download Rig TXT"),
             gr.Textbox(label="Status", lines=5)
         ],
-        title="🎭 RigNet: Neural Rigging for 3D Characters",
-        description="Upload OBJ (1K-5K vertices recommended). Processing takes 1-3 minutes on CPU.",
+        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)
+        """,
         allow_flagging="never"
     )
-    
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
+        share=False,
         show_error=True,
         debug=True
-    )
+    )