Update app.py

app.py

@@ -1,251 +1,407 @@
-import os
-import sys
-import tempfile
-import shutil
 import gradio as gr
+import os
 import torch
 import numpy as np
 import open3d as o3d
-from pathlib import Path
-
-# Add RigNet modules to path
-sys.path.insert(0, 'Rignet')
-
-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
-from utils.rig_parser import Info
-from gen_dataset import get_tpl_edges, get_geo_edges
+import trimesh
 from torch_geometric.data import Data
 from torch_geometric.utils import add_self_loops
-from utils import binvox_rw
-from geometric_proc.common_ops import calc_surface_geodesic, get_bones
-from utils.io_utils import assemble_skel_skin
-from run_skinning import post_filter
-
-# Import helper functions from quick_start
-from quick_start import (
-    normalize_obj, predict_joints, predict_skeleton, 
-    predict_skinning, calc_geodesic_matrix
-)
-
-# Initialize device
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-# Load models globally
-print("Loading neural networks...")
-jointNet = JOINTNET()
-jointNet.to(device)
-jointNet.eval()
-jointNet.load_state_dict(
-    torch.load('Rignet/checkpoints/gcn_meanshift/model_best.pth.tar', 
-               map_location=device)['state_dict']
-)
+import sys
 
-rootNet = ROOTNET()
-rootNet.to(device)
-rootNet.eval()
-rootNet.load_state_dict(
-    torch.load('Rignet/checkpoints/rootnet/model_best.pth.tar',
-               map_location=device)['state_dict']
-)
+# Add rignet to path
+sys.path.append('./rignet')
 
-boneNet = BONENET()
-boneNet.to(device)
-boneNet.eval()
-boneNet.load_state_dict(
-    torch.load('Rignet/checkpoints/bonenet/model_best.pth.tar',
-               map_location=device)['state_dict']
-)
+from rignet.utils import binvox_rw
+from rignet.utils.rig_parser import Info
+from rignet.utils.io_utils import assemble_skel_skin
+from rignet.utils.cluster_utils import meanshift_cluster, nms_meanshift
+from rignet.utils.mst_utils import inside_check, flip, increase_cost_for_outside_bone, primMST_symmetry, loadSkel_recur
+from rignet.geometric_proc.common_ops import get_bones, calc_surface_geodesic
+from rignet.geometric_proc.compute_volumetric_geodesic import calc_pts2bone_visible_mat, pts2line
+from rignet.gen_dataset import get_tpl_edges, get_geo_edges
+from rignet.mst_generate import sample_on_bone, getInitId
+from rignet.run_skinning import post_filter
+from rignet.models.GCN import JOINTNET_MASKNET_MEANSHIFT as JOINTNET
+from rignet.models.ROOT_GCN import ROOTNET
+from rignet.models.PairCls_GCN import PairCls as BONENET
+from rignet.models.SKINNING import SKINNET
 
-skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True)
-skinNet.to(device)
-skinNet.eval()
-skinNet.load_state_dict(
-    torch.load('Rignet/checkpoints/skinnet/model_best.pth.tar',
-               map_location=device)['state_dict']
-)
+# Global models
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+joint_net = None
+root_net = None
+bone_net = None
+skin_net = None
 
-print("All networks loaded successfully!")
+def load_models():
+    """Load all RigNet models"""
+    global joint_net, root_net, bone_net, skin_net
+    
+    print("Loading models...")
+    
+    # Joint prediction network
+    joint_net = JOINTNET()
+    joint_net.to(device)
+    joint_net.eval()
+    joint_checkpoint = torch.load('rignet/checkpoints/gcn_meanshift/model_best.pth.tar', map_location=device)
+    joint_net.load_state_dict(joint_checkpoint['state_dict'])
+    
+    # Root prediction network
+    root_net = ROOTNET()
+    root_net.to(device)
+    root_net.eval()
+    root_checkpoint = torch.load('rignet/checkpoints/rootnet/model_best.pth.tar', map_location=device)
+    root_net.load_state_dict(root_checkpoint['state_dict'])
+    
+    # Bone prediction network
+    bone_net = BONENET()
+    bone_net.to(device)
+    bone_net.eval()
+    bone_checkpoint = torch.load('rignet/checkpoints/bonenet/model_best.pth.tar', map_location=device)
+    bone_net.load_state_dict(bone_checkpoint['state_dict'])
+    
+    # Skinning prediction network
+    skin_net = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True)
+    skin_checkpoint = torch.load('rignet/checkpoints/skinnet/model_best.pth.tar', map_location=device)
+    skin_net.load_state_dict(skin_checkpoint['state_dict'])
+    skin_net.to(device)
+    skin_net.eval()
+    
+    print("All models loaded successfully!")
 
+def normalize_obj(mesh_v):
+    """Normalize mesh vertices"""
+    dims = [max(mesh_v[:, 0]) - min(mesh_v[:, 0]),
+            max(mesh_v[:, 1]) - min(mesh_v[:, 1]),
+            max(mesh_v[:, 2]) - min(mesh_v[:, 2])]
+    scale = 1.0 / max(dims)
+    pivot = np.array([(min(mesh_v[:, 0]) + max(mesh_v[:, 0])) / 2, min(mesh_v[:, 1]),
+                      (min(mesh_v[:, 2]) + max(mesh_v[:, 2])) / 2])
+    mesh_v[:, 0] -= pivot[0]
+    mesh_v[:, 1] -= pivot[1]
+    mesh_v[:, 2] -= pivot[2]
+    mesh_v *= scale
+    return mesh_v, pivot, scale
 
 def create_single_data(mesh_filename):
-    """Create input data for the network"""
+    """Create input data from mesh file"""
     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)
     
-    # Normalize mesh
-    mesh_v, translation, scale = normalize_obj(mesh_v)
+    mesh_v, translation_normalize, scale_normalize = normalize_obj(mesh_v)
     mesh_normalized = o3d.geometry.TriangleMesh(
         vertices=o3d.utility.Vector3dVector(mesh_v),
         triangles=o3d.utility.Vector3iVector(mesh_f)
     )
     
-    normalized_path = mesh_filename.replace('.obj', '_normalized.obj')
+    normalized_path = mesh_filename.replace(".obj", "_normalized.obj")
     o3d.io.write_triangle_mesh(normalized_path, mesh_normalized)
     
-    # Create vertex features
+    # vertices
     v = np.concatenate((mesh_v, mesh_vn), axis=1)
     v = torch.from_numpy(v).float()
     
-    # Topological edges
+    # topology 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
+    # surface geodesic
     surface_geodesic = calc_surface_geodesic(mesh)
     
-    # Geodesic edges
+    # 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))
     
-    # Voxelization
-    vox_path = normalized_path.replace('.obj', '.binvox')
-    if not os.path.exists(vox_path):
-        os.system(f"binvox -d 88 -pb {normalized_path}")
+    batch = torch.zeros(len(v), dtype=torch.long)
+    
+    # voxelization
+    binvox_path = mesh_filename.replace('.obj', '_normalized.binvox')
+    if not os.path.exists(binvox_path):
+        os.system(f"./binvox -d 88 -pb {normalized_path}")
     
-    with open(vox_path, 'rb') as fvox:
+    with open(binvox_path, 'rb') as fvox:
         vox = binvox_rw.read_as_3d_array(fvox)
     
-    batch = torch.zeros(len(v), dtype=torch.long)
-    data = Data(x=v[:, 3:6], pos=v[:, 0:3], 
-                tpl_edge_index=tpl_e, geo_edge_index=geo_e, batch=batch)
+    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, scale
-
+    return data, vox, surface_geodesic, translation_normalize, scale_normalize
 
-def process_mesh(input_file, bandwidth, threshold):
-    """
-    Main processing function for RigNet inference
-    
-    Args:
-        input_file: Uploaded OBJ file
-        bandwidth: Bandwidth parameter for joint clustering
-        threshold: Threshold for joint density filtering
-    
-    Returns:
-        Path to generated rig txt file and status message
-    """
+def predict_rig(input_obj_file):
+    """Main inference function"""
     try:
-        # Create temporary directory
-        temp_dir = tempfile.mkdtemp()
+        # Save uploaded file
+        temp_dir = "/tmp/rignet_temp"
+        os.makedirs(temp_dir, exist_ok=True)
+        
+        mesh_filename = os.path.join(temp_dir, "input_mesh.obj")
         
-        # Copy input file
-        input_path = os.path.join(temp_dir, 'input_mesh.obj')
-        shutil.copy(input_file.name, input_path)
+        # Copy uploaded file
+        if isinstance(input_obj_file, str):
+            os.system(f"cp {input_obj_file} {mesh_filename}")
+        else:
+            with open(mesh_filename, 'wb') as f:
+                f.write(input_obj_file.read())
         
-        # Simplify mesh if too large
-        mesh_ori = o3d.io.read_triangle_mesh(input_path)
-        if len(np.asarray(mesh_ori.vertices)) > 5000:
+        # Remesh if necessary
+        mesh_ori = o3d.io.read_triangle_mesh(mesh_filename)
+        if len(np.asarray(mesh_ori.vertices)) > 5000 or len(np.asarray(mesh_ori.vertices)) < 1000:
             mesh_remesh = mesh_ori.simplify_quadric_decimation(4000)
-            o3d.io.write_triangle_mesh(input_path, mesh_remesh)
+            mesh_filename = os.path.join(temp_dir, "input_mesh_remesh.obj")
+            o3d.io.write_triangle_mesh(mesh_filename, mesh_remesh)
         
         # Create data
-        print("Creating network input data...")
-        data, vox, surface_geodesic, translation, scale = create_single_data(input_path)
+        print("Creating input data...")
+        data, vox, surface_geodesic, translation_normalize, scale_normalize = create_single_data(mesh_filename)
         data.to(device)
         
         # Predict joints
         print("Predicting joints...")
-        bandwidth = bandwidth if bandwidth > 0 else None
-        data = predict_joints(
-            data, vox, jointNet, threshold, 
-            bandwidth=bandwidth,
-            mesh_filename=input_path.replace('.obj', '_normalized.obj')
-        )
+        data = predict_joints(data, vox, joint_net, threshold=1e-5, bandwidth=0.0429)
         data.to(device)
         
         # Predict skeleton
-        print("Predicting skeleton connectivity...")
-        pred_skeleton = predict_skeleton(
-            data, vox, rootNet, boneNet,
-            mesh_filename=input_path.replace('.obj', '_normalized.obj')
-        )
+        print("Predicting skeleton...")
+        pred_skeleton = predict_skeleton(data, vox, root_net, bone_net)
         
         # Predict skinning
-        print("Predicting skinning weights...")
-        pred_rig = predict_skinning(
-            data, pred_skeleton, skinNet, surface_geodesic,
-            input_path.replace('.obj', '_normalized.obj'),
-            subsampling=True
-        )
+        print("Predicting skinning...")
+        pred_rig = predict_skinning(data, pred_skeleton, skin_net, surface_geodesic,
+                                    mesh_filename.replace(".obj", "_normalized.obj"))
         
-        # Reverse normalization
-        pred_rig.normalize(scale, -translation)
+        # Denormalize
+        pred_rig.normalize(scale_normalize, -translation_normalize)
         
         # Save rig
-        output_path = os.path.join(temp_dir, 'output_rig.txt')
-        pred_rig.save(output_path)
+        output_rig_path = os.path.join(temp_dir, "output_rig.txt")
+        pred_rig.save(output_rig_path)
         
-        # Read rig content for display
-        with open(output_path, 'r') as f:
-            rig_content = f.read()
-        
-        return output_path, f"✅ Rigging completed successfully!\n\nGenerated {len(pred_skeleton.joint_pos)} joints", rig_content
+        print("Rig generation completed!")
+        return output_rig_path
         
     except Exception as e:
-        return None, f"❌ Error: {str(e)}", ""
+        return f"Error: {str(e)}"
 
+def predict_joints(input_data, vox, joint_pred_net, threshold, bandwidth=None):
+    """Predict joint locations"""
+    import itertools as it
+    
+    with torch.no_grad():
+        data_displacement, _, attn_pred, bandwidth_pred = joint_pred_net(input_data)
+    
+    y_pred = data_displacement + input_data.pos
+    y_pred_np = y_pred.data.cpu().numpy()
+    attn_pred_np = attn_pred.data.cpu().numpy()
+    
+    y_pred_np, index_inside = inside_check(y_pred_np, vox)
+    attn_pred_np = attn_pred_np[index_inside, :]
+    y_pred_np = y_pred_np[attn_pred_np.squeeze() > 1e-3]
+    attn_pred_np = attn_pred_np[attn_pred_np.squeeze() > 1e-3]
+    
+    # symmetrize
+    y_pred_np_reflect = y_pred_np * np.array([[-1, 1, 1]])
+    y_pred_np = np.concatenate((y_pred_np, y_pred_np_reflect), axis=0)
+    attn_pred_np = np.tile(attn_pred_np, (2, 1))
+    
+    if bandwidth is None:
+        bandwidth = bandwidth_pred.item()
+    
+    y_pred_np = meanshift_cluster(y_pred_np, bandwidth, attn_pred_np, max_iter=40)
+    
+    Y_dist = np.sum(((y_pred_np[np.newaxis, ...] - y_pred_np[:, np.newaxis, :]) ** 2), axis=2)
+    density = np.maximum(bandwidth ** 2 - Y_dist, np.zeros(Y_dist.shape))
+    density = np.sum(density, axis=0)
+    density_sum = np.sum(density)
+    
+    y_pred_np = y_pred_np[density / density_sum > threshold]
+    attn_pred_np = attn_pred_np[density / density_sum > threshold][:, 0]
+    density = density[density / density_sum > threshold]
+    
+    pred_joints = nms_meanshift(y_pred_np, density, bandwidth)
+    pred_joints, _ = flip(pred_joints)
+    
+    # prepare bone pairs
+    pairs = list(it.combinations(range(pred_joints.shape[0]), 2))
+    pair_attr = []
+    for pr in pairs:
+        dist = np.linalg.norm(pred_joints[pr[0]] - pred_joints[pr[1]])
+        bone_samples = sample_on_bone(pred_joints[pr[0]], pred_joints[pr[1]])
+        bone_samples_inside, _ = inside_check(bone_samples, vox)
+        outside_proportion = len(bone_samples_inside) / (len(bone_samples) + 1e-10)
+        attr = np.array([dist, outside_proportion, 1])
+        pair_attr.append(attr)
+    
+    pairs = np.array(pairs)
+    pair_attr = np.array(pair_attr)
+    pairs = torch.from_numpy(pairs).float()
+    pair_attr = torch.from_numpy(pair_attr).float()
+    pred_joints = torch.from_numpy(pred_joints).float()
+    
+    input_data.joints = pred_joints
+    input_data.pairs = pairs
+    input_data.pair_attr = pair_attr
+    input_data.joints_batch = torch.zeros(len(pred_joints), dtype=torch.long)
+    input_data.pairs_batch = torch.zeros(len(pairs), dtype=torch.long)
+    
+    return input_data
 
-# Create Gradio interface
-with gr.Blocks(title="RigNet: Neural Rigging for 3D Characters") as demo:
-    gr.Markdown("""
-    # 🎮 RigNet: Neural Rigging for Articulated Characters
-    
-    Upload a 3D character mesh (OBJ format) and get an automatic rig with skeleton and skinning weights.
-    
-    **Paper**: [RigNet SIGGRAPH 2020](https://zhan-xu.github.io/rig-net/)  
-    **Code**: [GitHub Repository](https://github.com/zhan-xu/RigNet)
-    """)
-    
-    with gr.Row():
-        with gr.Column():
-            input_file = gr.File(
-                label="Upload OBJ File",
-                file_types=[".obj"],
-                type="filepath"
-            )
-            
-            with gr.Accordion("Advanced Parameters", open=False):
-                bandwidth = gr.Slider(
-                    minimum=0.0, maximum=0.1, value=0.0429, step=0.001,
-                    label="Bandwidth (0 for auto)",
-                    info="Controls joint clustering. Try 0.04-0.06 if results are poor."
-                )
-                threshold = gr.Slider(
-                    minimum=0.0, maximum=5.0, value=1.0, step=0.1,
-                    label="Threshold (×1e-5)",
-                    info="Filters joint density. Higher = fewer joints."
-                )
-            
-            process_btn = gr.Button("🚀 Generate Rig", variant="primary")
+def predict_skeleton(input_data, vox, root_pred_net, bone_pred_net):
+    """Predict skeleton connectivity"""
+    from rignet.utils.tree_utils import TreeNode
+    
+    root_id = getInitId(input_data, root_pred_net)
+    pred_joints = input_data.joints.data.cpu().numpy()
+    
+    with torch.no_grad():
+        connect_prob, _ = bone_pred_net(input_data, permute_joints=False)
+        connect_prob = torch.sigmoid(connect_prob)
+    
+    pair_idx = input_data.pairs.long().data.cpu().numpy()
+    prob_matrix = np.zeros((len(input_data.joints), len(input_data.joints)))
+    prob_matrix[pair_idx[:, 0], pair_idx[:, 1]] = connect_prob.data.cpu().numpy().squeeze()
+    prob_matrix = prob_matrix + prob_matrix.transpose()
+    
+    cost_matrix = -np.log(prob_matrix + 1e-10)
+    cost_matrix = increase_cost_for_outside_bone(cost_matrix, pred_joints, vox)
+    
+    pred_skel = Info()
+    parent, key, root_id = primMST_symmetry(cost_matrix, root_id, pred_joints)
+    
+    for i in range(len(parent)):
+        if parent[i] == -1:
+            pred_skel.root = TreeNode('root', tuple(pred_joints[i]))
+            break
+    
+    loadSkel_recur(pred_skel.root, i, None, pred_joints, parent)
+    pred_skel.joint_pos = pred_skel.get_joint_dict()
+    
+    return pred_skel
+
+def predict_skinning(input_data, pred_skel, skin_pred_net, surface_geodesic, mesh_filename):
+    """Predict skinning weights"""
+    num_nearest_bone = 5
+    bones, bone_names, bone_isleaf = get_bones(pred_skel)
+    mesh_v = input_data.pos.data.cpu().numpy()
+    
+    # Calculate geodesic distance
+    geo_dist = calc_geodesic_matrix(bones, mesh_v, surface_geodesic, mesh_filename)
+    
+    input_samples = []
+    loss_mask = []
+    skin_nn = []
+    
+    for v_id in range(len(mesh_v)):
+        geo_dist_v = geo_dist[v_id]
+        bone_id_near_to_far = np.argsort(geo_dist_v)
+        this_sample = []
+        this_nn = []
+        this_mask = []
         
-        with gr.Column():
-            status_text = gr.Textbox(label="Status", lines=3)
-            rig_content = gr.Textbox(
-                label="Generated Rig Content (Preview)",
-                lines=10,
-                max_lines=20
-            )
-            output_file = gr.File(label="Download Rig File (.txt)")
-    
-    # Examples
-    gr.Markdown("### 📎 Example Models")
-    gr.Markdown("Upload your own OBJ file or use default parameters (bandwidth=0.0429, threshold=1e-5)")
-    
-    # Process button click
-    process_btn.click(
-        fn=process_mesh,
-        inputs=[input_file, bandwidth, threshold],
-        outputs=[output_file, status_text, rig_content]
-    )
+        for i in range(num_nearest_bone):
+            if i >= len(bones):
+                this_sample += bones[bone_id_near_to_far[0]].tolist()
+                this_sample.append(1.0 / (geo_dist_v[bone_id_near_to_far[0]] + 1e-10))
+                this_sample.append(bone_isleaf[bone_id_near_to_far[0]])
+                this_nn.append(0)
+                this_mask.append(0)
+            else:
+                skel_bone_id = bone_id_near_to_far[i]
+                this_sample += bones[skel_bone_id].tolist()
+                this_sample.append(1.0 / (geo_dist_v[skel_bone_id] + 1e-10))
+                this_sample.append(bone_isleaf[skel_bone_id])
+                this_nn.append(skel_bone_id)
+                this_mask.append(1)
+        
+        input_samples.append(np.array(this_sample)[np.newaxis, :])
+        skin_nn.append(np.array(this_nn)[np.newaxis, :])
+        loss_mask.append(np.array(this_mask)[np.newaxis, :])
+    
+    skin_input = np.concatenate(input_samples, axis=0)
+    loss_mask = np.concatenate(loss_mask, axis=0)
+    skin_nn = np.concatenate(skin_nn, axis=0)
+    skin_input = torch.from_numpy(skin_input).float()
+    
+    input_data.skin_input = skin_input
+    input_data.to(device)
+    
+    with torch.no_grad():
+        skin_pred = skin_pred_net(input_data)
+        skin_pred = torch.softmax(skin_pred, dim=1)
+    
+    skin_pred = skin_pred.data.cpu().numpy()
+    skin_pred = skin_pred * loss_mask
+    skin_nn = skin_nn[:, 0:num_nearest_bone]
+    
+    skin_pred_full = np.zeros((len(skin_pred), len(bone_names)))
+    for v in range(len(skin_pred)):
+        for nn_id in range(len(skin_nn[v, :])):
+            skin_pred_full[v, skin_nn[v, nn_id]] = skin_pred[v, nn_id]
+    
+    # Post-processing
+    tpl_e = input_data.tpl_edge_index.data.cpu().numpy()
+    skin_pred_full = post_filter(skin_pred_full, tpl_e, num_ring=1)
+    skin_pred_full[skin_pred_full < np.max(skin_pred_full, axis=1, keepdims=True) * 0.35] = 0.0
+    skin_pred_full = skin_pred_full / (skin_pred_full.sum(axis=1, keepdims=True) + 1e-10)
+    
+    skel_res = assemble_skel_skin(pred_skel, skin_pred_full)
+    return skel_res
+
+def calc_geodesic_matrix(bones, mesh_v, surface_geodesic, mesh_filename):
+    """Calculate volumetric geodesic distance"""
+    mesh_trimesh = trimesh.load(mesh_filename)
+    subsamples = mesh_v
+    
+    origins, ends, pts_bone_dist = pts2line(subsamples, bones)
+    pts_bone_visibility = calc_pts2bone_visible_mat(mesh_trimesh, origins, ends)
+    pts_bone_visibility = pts_bone_visibility.reshape(len(bones), len(subsamples)).transpose()
+    pts_bone_dist = pts_bone_dist.reshape(len(bones), len(subsamples)).transpose()
+    
+    for b in range(pts_bone_visibility.shape[1]):
+        visible_pts = np.argwhere(pts_bone_visibility[:, b] == 1).squeeze(1)
+        if len(visible_pts) == 0:
+            continue
+        threshold_b = np.percentile(pts_bone_dist[visible_pts, b], 15)
+        pts_bone_visibility[pts_bone_dist[:, b] > 1.3 * threshold_b, b] = False
+    
+    visible_matrix = np.zeros(pts_bone_visibility.shape)
+    visible_matrix[np.where(pts_bone_visibility == 1)] = pts_bone_dist[np.where(pts_bone_visibility == 1)]
+    
+    for c in range(visible_matrix.shape[1]):
+        unvisible_pts = np.argwhere(pts_bone_visibility[:, c] == 0).squeeze(1)
+        visible_pts = np.argwhere(pts_bone_visibility[:, c] == 1).squeeze(1)
+        
+        if len(visible_pts) == 0:
+            visible_matrix[:, c] = pts_bone_dist[:, c]
+            continue
+        
+        for r in unvisible_pts:
+            dist1 = np.min(surface_geodesic[r, visible_pts])
+            nn_visible = visible_pts[np.argmin(surface_geodesic[r, visible_pts])]
+            if np.isinf(dist1):
+                visible_matrix[r, c] = 8.0 + pts_bone_dist[r, c]
+            else:
+                visible_matrix[r, c] = dist1 + visible_matrix[nn_visible, c]
+    
+    return visible_matrix
+
+# Load models at startup
+load_models()
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict_rig,
+    inputs=gr.File(label="Upload OBJ File", file_types=[".obj"]),
+    outputs=gr.File(label="Download Rig TXT"),
+    title="RigNet: Automatic Character Rigging",
+    description="Upload a 3D character mesh in OBJ format to automatically generate a rig. The model will predict joints, skeleton hierarchy, and skinning weights. Best results with meshes containing 1K-5K vertices.",
+    examples=[],
+    cache_examples=False
+)
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    iface.launch(server_name="0.0.0.0", server_port=7860)