Update app.py

app.py

@@ -1,158 +1,514 @@
+"""
+RigNet Gradio Demo
+Automatic rigging for 3D character models
+Based on: https://github.com/zhan-xu/RigNet
+"""
+
 import os
 import sys
-import shutil
-import platform
+import gradio as gr
 import torch
 import numpy as np
 import open3d as o3d
-import gradio as gr
+import trimesh
+from pathlib import Path
+import tempfile
+import shutil
 from torch_geometric.data import Data
+from torch_geometric.utils import add_self_loops
 
-# --- 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.
+# Import RigNet modules (assuming they're in the same directory structure)
+from utils import binvox_rw
+from utils.rig_parser import Info
+from utils.io_utils import assemble_skel_skin
+from utils.cluster_utils import meanshift_cluster, nms_meanshift
+from utils.mst_utils import inside_check, flip, increase_cost_for_outside_bone, primMST_symmetry, loadSkel_recur
+from geometric_proc.common_ops import get_bones, calc_surface_geodesic
+from gen_dataset import get_tpl_edges, get_geo_edges
+from mst_generate import sample_on_bone, getInitId
+from run_skinning import post_filter
+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 itertools as it
+from geometric_proc.compute_volumetric_geodesic import pts2line, calc_pts2bone_visible_mat
 
-# --- 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)
+# Global variables
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+models_loaded = False
+jointNet, rootNet, boneNet, skinNet = None, None, None, None
 
-# Ensure binvox 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}")
-
-# --- 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 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
-    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")
-print(f"Loading RigNet models on {device}...")
-
-def load_checkpoint(model, filename):
-    filepath = os.path.join(RIGNET_PATH, "checkpoints", filename)
-    # 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()
-load_checkpoint(jointNet, 'gcn_meanshift/model_best.pth.tar')
+def load_models():
+    """Load all pre-trained RigNet models"""
+    global jointNet, rootNet, boneNet, skinNet, models_loaded
+    
+    if models_loaded:
+        return
+    
+    print("Loading RigNet models...")
+    
+    # Joint prediction network
+    jointNet = JOINTNET()
+    jointNet.to(device)
+    jointNet.eval()
+    jointNet_checkpoint = torch.load('checkpoints/gcn_meanshift/model_best.pth.tar', 
+                                    map_location=device)
+    jointNet.load_state_dict(jointNet_checkpoint['state_dict'])
+    print("✓ Joint prediction network loaded")
+    
+    # Root prediction network
+    rootNet = ROOTNET()
+    rootNet.to(device)
+    rootNet.eval()
+    rootNet_checkpoint = torch.load('checkpoints/rootnet/model_best.pth.tar',
+                                   map_location=device)
+    rootNet.load_state_dict(rootNet_checkpoint['state_dict'])
+    print("✓ Root prediction network loaded")
+    
+    # Bone connection network
+    boneNet = BONENET()
+    boneNet.to(device)
+    boneNet.eval()
+    boneNet_checkpoint = torch.load('checkpoints/bonenet/model_best.pth.tar',
+                                   map_location=device)
+    boneNet.load_state_dict(boneNet_checkpoint['state_dict'])
+    print("✓ Connection prediction network loaded")
+    
+    # Skinning prediction network
+    skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True)
+    skinNet_checkpoint = torch.load('checkpoints/skinnet/model_best.pth.tar',
+                                   map_location=device)
+    skinNet.load_state_dict(skinNet_checkpoint['state_dict'])
+    skinNet.to(device)
+    skinNet.eval()
+    print("✓ Skinning prediction network loaded")
+    
+    models_loaded = True
+    print("All models loaded successfully!")
 
-rootNet = ROOTNET().to(device); rootNet.eval()
-load_checkpoint(rootNet, 'rootnet/model_best.pth.tar')
+def normalize_obj(mesh_v):
+    """Normalize mesh vertices to unit scale"""
+    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
 
-boneNet = BONENET().to(device); boneNet.eval()
-load_checkpoint(boneNet, 'bonenet/model_best.pth.tar')
+def create_single_data(mesh_filename):
+    """Create input data for the network"""
+    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_path = mesh_filename.replace("_remesh.obj", "_normalized.obj")
+    o3d.io.write_triangle_mesh(normalized_path, mesh_normalized)
+    
+    # Vertices
+    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 distance matrix
+    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
+    vox_path = mesh_filename.replace('_remesh.obj', '_normalized.binvox')
+    if not os.path.exists(vox_path):
+        # Use binvox command
+        if sys.platform == "linux" or sys.platform == "linux2":
+            os.system(f"./binvox -d 88 -pb {normalized_path}")
+        elif sys.platform == "win32":
+            os.system(f"binvox.exe -d 88 {normalized_path}")
+    
+    with open(vox_path, '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
 
-skinNet = SKINNET(nearest_bone=5, use_Dg=True, use_Lf=True).to(device); skinNet.eval()
-load_checkpoint(skinNet, 'skinnet/model_best.pth.tar')
+def predict_joints(input_data, vox, threshold=1e-5, bandwidth=None):
+    """Predict skeleton joints"""
+    data_displacement, _, attn_pred, bandwidth_pred = jointNet(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 points by reflecting
+    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 pair-wise bone data
+    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()
+    
+    joints_batch = torch.zeros(len(pred_joints), dtype=torch.long)
+    pairs_batch = torch.zeros(len(pairs), dtype=torch.long)
+    
+    input_data.joints = pred_joints
+    input_data.pairs = pairs
+    input_data.pair_attr = pair_attr
+    input_data.joints_batch = joints_batch
+    input_data.pairs_batch = pairs_batch
+    
+    return input_data
 
-print("Models loaded.")
+def predict_skeleton(input_data, vox):
+    """Predict skeleton structure"""
+    root_id = getInitId(input_data, rootNet)
+    pred_joints = input_data.joints.data.cpu().numpy()
+    
+    with torch.no_grad():
+        connect_prob, _ = boneNet(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:
+            from utils.tree_utils import TreeNode
+            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
 
-# --- 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
+def calc_geodesic_matrix(bones, mesh_v, surface_geodesic, mesh_filename, subsampling=True):
+    """Calculate volumetric geodesic distance from vertices to bones"""
+    if subsampling:
+        mesh0 = o3d.io.read_triangle_mesh(mesh_filename)
+        mesh0 = mesh0.simplify_quadric_decimation(3000)
+        simplified_path = mesh_filename.replace(".obj", "_simplified.obj")
+        o3d.io.write_triangle_mesh(simplified_path, mesh0)
+        mesh_trimesh = trimesh.load(simplified_path)
+        
+        subsamples_ids = np.random.choice(len(mesh_v), np.min((len(mesh_v), 1500)), replace=False)
+        subsamples = mesh_v[subsamples_ids, :]
+        surface_geodesic = surface_geodesic[subsamples_ids, :][:, subsamples_ids]
+    else:
+        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()
+    
+    # Remove visible points which are too far
+    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]
+    
+    if subsampling:
+        nn_dist = np.sum((mesh_v[:, np.newaxis, :] - subsamples[np.newaxis, ...]) ** 2, axis=2)
+        nn_ind = np.argmin(nn_dist, axis=1)
+        visible_matrix = visible_matrix[nn_ind, :]
+        os.remove(simplified_path)
+    
+    return visible_matrix
 
-# --- 5. Inference Pipeline ---
-def rignet_inference(input_mesh_path):
-    if not input_mesh_path: return None
+def predict_skinning(input_data, pred_skel, surface_geodesic, mesh_filename, subsampling=True):
+    """Predict skinning weights"""
+    num_nearest_bone = 5
+    bones, bone_names, bone_isleaf = get_bones(pred_skel)
+    mesh_v = input_data.pos.data.cpu().numpy()
+    
+    print("  Calculating volumetric geodesic distance...")
+    geo_dist = calc_geodesic_matrix(bones, mesh_v, surface_geodesic, mesh_filename, subsampling=subsampling)
+    
+    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 = []
+        
+        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)
+    
+    skin_pred = skinNet(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)))
     
-    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")
+    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]
     
-    print(f"Processing: {input_mesh_path}")
+    print("  Filtering skinning prediction...")
+    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 process_model(input_file, bandwidth_val, threshold_val):
+    """Main processing function for Gradio interface"""
     try:
-        # 1. Preprocess
-        mesh_ori = o3d.io.read_triangle_mesh(input_mesh_path)
-        if len(np.asarray(mesh_ori.vertices)) == 0: raise ValueError("Empty mesh")
+        # Load models if not already loaded
+        load_models()
+        
+        # Create temporary directory for processing
+        temp_dir = tempfile.mkdtemp()
+        
+        # Copy input file
+        input_path = Path(input_file.name)
+        temp_input = os.path.join(temp_dir, "input_ori.obj")
+        shutil.copy(input_path, temp_input)
+        
+        # Remesh the input
+        print("Preprocessing: Remeshing input...")
+        mesh_ori = o3d.io.read_triangle_mesh(temp_input)
         mesh_remesh = mesh_ori.simplify_quadric_decimation(4000)
-        o3d.io.write_triangle_mesh(mesh_filename, mesh_remesh)
-
-        # 2. Create Data (Patched)
-        data, vox, surface_geodesic, t_norm, s_norm = create_single_data_patched(mesh_filename)
-        data = data.to(device)
-
-        # 3. Inference
-        mesh_norm = mesh_filename.replace("_remesh.obj", "_normalized.obj")
+        temp_remesh = os.path.join(temp_dir, "input_remesh.obj")
+        o3d.io.write_triangle_mesh(temp_remesh, mesh_remesh)
+        
+        # Create data
+        print("Creating data...")
+        data, vox, surface_geodesic, translation, scale = create_single_data(temp_remesh)
+        data.to(device)
         
+        # Predict joints
         print("Predicting joints...")
-        data = predict_joints(data, vox, jointNet, 1e-5, bandwidth=0.0429, mesh_filename=mesh_norm)
-        data = data.to(device)
-
-        print("Predicting connectivity...")
-        skel = predict_skeleton(data, vox, rootNet, boneNet, mesh_filename=mesh_norm)
-
-        print("Predicting skinning...")
-        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)
+        data = predict_joints(data, vox, threshold=threshold_val, bandwidth=bandwidth_val)
+        data.to(device)
+        
+        # Predict skeleton
+        print("Predicting skeleton connectivity...")
+        pred_skeleton = predict_skeleton(data, vox)
+        
+        # Predict skinning
+        print("Predicting skinning weights...")
+        normalized_mesh = temp_remesh.replace("_remesh.obj", "_normalized.obj")
+        pred_rig = predict_skinning(data, pred_skeleton, surface_geodesic, 
+                                    normalized_mesh, subsampling=True)
+        
+        # Denormalize
+        pred_rig.normalize(scale, -translation)
+        
+        # Save result
+        output_file = os.path.join(temp_dir, "output_rig.txt")
+        pred_rig.save(output_file)
+        
+        print("✓ Processing complete!")
+        
+        # Create info message
+        num_joints = len(pred_rig.joint_pos)
+        info_msg = f"Successfully generated rig with {num_joints} joints!"
+        
+        return output_file, info_msg
         
-        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: {e}")
-        raise gr.Error(f"Processing failed: {str(e)}")
+        import traceback
+        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        return None, error_msg
 
-# --- 6. 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. (Uses PyTorch 2.1 CPU)"
-)
+# Gradio Interface
+def create_demo():
+    """Create Gradio interface"""
+    
+    with gr.Blocks(title="RigNet: Automatic 3D Character Rigging") as demo:
+        gr.Markdown("""
+        # 🎮 RigNet: Automatic Character Rigging
+        
+        Upload a 3D character model (OBJ format) and automatically generate a skeletal rig with skinning weights.
+        
+        **Based on:** [RigNet: Neural Rigging for Articulated Characters (SIGGRAPH 2020)](https://github.com/zhan-xu/RigNet)
+        
+        ### Instructions:
+        1. Upload your 3D character mesh in OBJ format
+        2. Adjust parameters if needed (default values work for most cases)
+        3. Click "Generate Rig" and wait for processing
+        4. Download the generated rig file
+        
+        **Note:** For best results, simplify your mesh to 1K-5K vertices before uploading.
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                input_mesh = gr.File(label="Upload 3D Model (.obj)", file_types=[".obj"])
+                
+                with gr.Accordion("Advanced Parameters", open=False):
+                    bandwidth = gr.Slider(
+                        minimum=0.02, maximum=0.08, value=0.0429, step=0.001,
+                        label="Bandwidth (for joint clustering)",
+                        info="Default: 0.0429. Adjust if joint prediction is too dense/sparse"
+                    )
+                    threshold = gr.Slider(
+                        minimum=0.1e-5, maximum=5e-5, value=1e-5, step=0.1e-5,
+                        label="Density Threshold",
+                        info="Default: 1e-5. Higher values = fewer joints"
+                    )
+                
+                process_btn = gr.Button("🚀 Generate Rig", variant="primary", size="lg")
+            
+            with gr.Column():
+                output_file = gr.File(label="Download Rig Output (.txt)")
+                status_msg = gr.Textbox(label="Status", lines=3)
+        
+        gr.Markdown("""
+        ### Output Format:
+        The generated `.txt` file contains:
+        - **Joint definitions:** Position of each joint in 3D space
+        - **Hierarchy:** Parent-child relationships between joints
+        - **Skinning weights:** How each vertex is influenced by nearby joints
+        
+        ### Next Steps:
+        - Import the mesh and rig file into animation software (Maya, Blender, etc.)
+        - Use provided scripts (e.g., `maya_save_fbx.py`) to convert to FBX format
+        - Start animating your character!
+        
+        ---
+        **References:**
+        - [RigNet Paper](https://arxiv.org/abs/2005.00559)
+        - [GitHub Repository](https://github.com/zhan-xu/RigNet)
+        - [Project Page](https://zhan-xu.github.io/rig-net/)
+        """)
+        
+        # Event handler
+        process_btn.click(
+            fn=process_model,
+            inputs=[input_mesh, bandwidth, threshold],
+            outputs=[output_file, status_msg]
+        )
+    
+    return demo
 
 if __name__ == "__main__":
-    iface.launch()
+    demo = create_demo()
+    demo.launch(share=False)