update

app.py

@@ -1,5 +1,6 @@
 import os
 import cv2
+import uuid
 import torch
 import numpy as np
 import gradio as gr
@@ -24,7 +25,7 @@ from mediapipe.tasks.python import BaseOptions
 # Configuration
 BACKBONE = 'hamer'
 EXPERIMENT_DIR = 'experiments_demo_image'
-EXAMPLE_DIR = '/home/user/app/asset/example_images'  # HF Spaces absolute path
+EXAMPLE_DIR = '/home/user/app/asset/example_images'
 
 # Setup
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -44,16 +45,17 @@ checkpoint = torch.load(checkpoint_path, map_location=device)
 model.load_state_dict(checkpoint['state_dict'])
 model.eval()
 
-# Load HandLandmarker
+# Load Mediapipe HandLandmarker
 base_options = BaseOptions(model_asset_path=cfg.MODEL.hand_landmarker_path)
 hand_options = vision.HandLandmarkerOptions(base_options=base_options, num_hands=2)
 detector = vision.HandLandmarker.create_from_options(hand_options)
 
+# Inference function
 def process_image(pil_img: Image.Image):
     orig_img = np.array(pil_img.convert("RGB"))
     mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=orig_img.copy())
     detection_result = detector.detect(mp_image)
-    annotated_image, right_hand_bbox = draw_landmarks_on_image(orig_img.copy(), detection_result)
+    _, right_hand_bbox = draw_landmarks_on_image(orig_img.copy(), detection_result)
 
     if right_hand_bbox is None:
         return None, "No hand detected."
@@ -61,6 +63,7 @@ def process_image(pil_img: Image.Image):
     crop_img, _, _, _, _, _ = augmentation_contact(
         orig_img.copy(), right_hand_bbox, 'test', enforce_flip=False)
 
+    # Prepare input tensor
     if BACKBONE in ['handoccnet'] or 'resnet' in cfg.MODEL.backbone_type or 'hrnet' in cfg.MODEL.backbone_type:
         img_tensor = transforms.ToTensor()(crop_img.astype(np.float32) / 255.0)
     elif BACKBONE in ['hamer'] or 'vit' in cfg.MODEL.backbone_type:
@@ -70,29 +73,45 @@ def process_image(pil_img: Image.Image):
     else:
         raise NotImplementedError(f"Unsupported backbone: {BACKBONE}")
 
+    # Inference
     with torch.no_grad():
         outputs = model({'input': {'image': img_tensor[None].to(device)}}, mode="test")
 
+    # Postprocess contact mask
     eval_thres = get_contact_thres(BACKBONE)
     contact_mask = (outputs['contact_out'][0] > eval_thres).detach().cpu().numpy()
     contact_mask = remove_small_contact_components(contact_mask, faces=mano.watertight_face['right'], min_size=20)
-    mesh_path = contact_renderer.export_contact_mesh(contact_mask)
-    return mesh_path
 
+    # Export .glb mesh
+    output_path = f"/tmp/contact_mesh_{uuid.uuid4().hex}.glb"
+    mesh_path = contact_renderer.export_contact_mesh(contact_mask, output_path=output_path)
 
-# Description markdown
+    return mesh_path, "Success"
+
+# Gradio UI
 description = '''
 ### HACO: Learning Dense Hand Contact Estimation from Imbalanced Data  
-Official demo for our paper on hand contact estimation.  
-Upload an image of a hand-object interaction scene, and HACO will predict and visualize contact regions on the hand.
+Upload an image of a hand-object interaction. HACO will predict contact regions and return a colored 3D hand mesh.
 '''
 
-# Gradio Interface
-with gr.Blocks(title="HACO Image Demo", css=".gradio-container") as demo:
-    input_image = gr.Image(label="Input Image", type="pil")
-    output_model = gr.Model3D(label="Contact Mesh")
-    run_button = gr.Button("Run HACO")
+with gr.Blocks(title="HACO Image Demo") as demo:
+    gr.Markdown(description)
 
-    run_button.click(fn=process_image, inputs=input_image, outputs=output_model)
+    with gr.Row():
+        input_image = gr.Image(label="Input Image", type="pil")
+        output_model = gr.Model3D(label="Predicted Contact Mesh")
+        output_status = gr.Textbox(label="Status")
+
+    run_button = gr.Button("Run HACO")
+    run_button.click(fn=process_image, inputs=input_image, outputs=[output_model, output_status])
+
+    # Example image list
+    example_list = [
+        [os.path.join(EXAMPLE_DIR, f)] for f in [
+            "holding_cup1.jpg", "holding_cup2.jpg", "holding_hammer1.jpg",
+            "holding_scissors2.jpg", "squidgame_demo2.png", "touching_wall1.jpg", "using_pen1.jpg"
+        ]
+    ]
+    gr.Examples(examples=example_list, inputs=[input_image], label="Example Images")
 
 demo.launch()

lib/models/decoder/decoder_hamer_style.py

@@ -17,10 +17,6 @@ from lib.core.config import cfg
 from lib.utils.human_models import mano
 
 
-V_regressor_336 = np.load(cfg.MODEL.V_regressor_336_path)
-V_regressor_84 = np.load(cfg.MODEL.V_regressor_84_path)
-
-
 # This function is from HaMeR (https://github.com/geopavlakos/hamer).
 def exists(val):
     return val is not None
@@ -629,9 +625,5 @@ class ContactTransformerDecoderHead(nn.Module):
         pred_contact = self.deccontact(token_out) + pred_contact
         # pred_contact = pred_contact.sigmoid()
 
-        # Joint contact
-        pred_joint_contact = (torch.tensor(mano.joint_regressor, dtype=torch.float32, device=device) @ pred_contact.T).T
-        pred_mesh_contact_336 = (torch.tensor(V_regressor_336, dtype=torch.float32, device=device) @ pred_contact.T).T
-        pred_mesh_contact_84 = (torch.tensor(V_regressor_84, dtype=torch.float32, device=device) @ pred_contact.T).T
 
-        return pred_contact, pred_mesh_contact_336, pred_mesh_contact_84, pred_joint_contact
+        return pred_contact

lib/models/model.py

@@ -24,9 +24,9 @@ class HACO(nn.Module):
             image = F.interpolate(image, size=(224, 224), mode='bilinear', align_corners=False)
         
         img_feat = self.backbone(image)
-        contact_out, contact_336_out, contact_84_out, contact_joint_out = self.decoder(img_feat)
+        contact_out = self.decoder(img_feat)
         
-        return dict(contact_out=contact_out, contact_336_out=contact_336_out, contact_84_out=contact_84_out, contact_joint_out=contact_joint_out)
+        return dict(contact_out=contact_out)
 
 
 

lib/utils/vis_utils.py

@@ -1,6 +1,7 @@
 import os
 import cv2
 import torch
+import uuid
 import trimesh
 import numpy as np
 import matplotlib.cm as cm
@@ -24,13 +25,12 @@ class ContactRenderer:
             mano_rest_out = mano.layer['right'](betas=torch.zeros((1, 10)), hand_pose=torch.from_numpy(hand_pose[None, 3:]).float(), global_orient=torch.zeros((1, 3)), transl=torch.zeros((1, 3)))
             self.hand_model_mano = trimesh.Trimesh(mano_rest_out.vertices[0], mano.watertight_face['right'])
 
-    def export_contact_mesh(self, contact_mask, output_path='contact_mesh.glb'):
+    def export_contact_mesh(self, contact_mask, output_path=None):
         vis_contact = contact_mask == 1.0
-        # Apply default color
         self.hand_model_mano.visual.vertex_colors = np.tile(self.default_mesh_color, (self.hand_model_mano.vertices.shape[0], 1))
-        # Apply contact color
         self.hand_model_mano.visual.vertex_colors[vis_contact] = self.contact_mesh_color
-        # Export to .glb
+        if output_path is None:
+            output_path = f"/tmp/contact_mesh_{uuid.uuid4().hex}.glb"
         self.hand_model_mano.export(output_path)
         return output_path