Create app.py

app.py

@@ -0,0 +1,116 @@
+import os
+import sys
+import shutil
+import platform
+import torch
+import numpy as np
+import open3d as o3d
+import gradio as gr
+
+# --- 1. Setup Environment & Paths ---
+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 is executable
+BINVOX_SRC = os.path.join(RIGNET_PATH, "binvox")
+BINVOX_DEST = "binvox.exe" if platform.system() == "Windows" else "binvox"
+if platform.system() == "Windows":
+    BINVOX_SRC += ".exe"
+
+if os.path.exists(BINVOX_SRC):
+    shutil.copy(BINVOX_SRC, 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.")
+
+# --- 2. Import RigNet Modules ---
+try:
+    from quick_start import (
+        create_single_data, predict_joints, predict_skeleton, 
+        predict_skinning, tranfer_to_ori_mesh
+    )
+    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
+except ImportError as e:
+    print(f"Error importing RigNet: {e}")
+
+# --- 3. Load Models ---
+device = torch.device("cpu") # Force CPU for this demo environment
+print(f"Loading RigNet models on {device}...")
+
+def load_checkpoint(model, filename):
+    filepath = os.path.join(RIGNET_PATH, "checkpoints", filename)
+    checkpoint = torch.load(filepath, map_location=device)
+    model.load_state_dict(checkpoint['state_dict'])
+    return model
+
+# Initialize models (Global)
+jointNet = JOINTNET().to(device); jointNet.eval()
+load_checkpoint(jointNet, 'gcn_meanshift/model_best.pth.tar')
+
+rootNet = ROOTNET().to(device); rootNet.eval()
+load_checkpoint(rootNet, 'rootnet/model_best.pth.tar')
+
+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()
+load_checkpoint(skinNet, 'skinnet/model_best.pth.tar')
+
+print("Models loaded.")
+
+# --- 4. 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", "")
+    mesh_filename = os.path.join(working_dir, f"{base_name}_remesh.obj")
+    
+    print(f"Processing: {input_mesh_path}")
+    try:
+        # 1. Simplify
+        mesh_ori = o3d.io.read_triangle_mesh(input_mesh_path)
+        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. Data Prep
+        data, vox, surface_geodesic, t_norm, s_norm = create_single_data(mesh_filename)
+        data = data.to(device)
+
+        # 3. Predictions
+        mesh_norm = mesh_filename.replace("_remesh.obj", "_normalized.obj")
+        data = predict_joints(data, vox, jointNet, 1e-5, bandwidth=0.0429, mesh_filename=mesh_norm)
+        data = data.to(device)
+        
+        skel = predict_skeleton(data, vox, rootNet, boneNet, mesh_filename=mesh_norm)
+        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)
+        
+        out_path = os.path.join(working_dir, f"{base_name}_rig.txt")
+        final_rig.save(out_path)
+        return out_path
+
+    except Exception as e:
+        raise gr.Error(f"Error: {str(e)}")
+
+# --- 5. Launch ---
+iface = gr.Interface(
+    fn=rignet_inference,
+    inputs=gr.Model3D(label="Input .obj"),
+    outputs=gr.File(label="Rig Output .txt"),
+    title="RigNet Demo",
+    description="Upload a mesh to generate a rig."
+)
+
+if __name__ == "__main__":
+    # server_name="0.0.0.0" is CRITICAL for Docker Spaces
+    iface.launch(server_name="0.0.0.0", server_port=7860)