Update app.py

app.py

@@ -21,7 +21,6 @@ from wilor.utils import recursive_to
 from wilor.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD
 from wilor.utils.renderer import Renderer, cam_crop_to_full
 device = torch.device('cpu') if torch.cuda.is_available() else torch.device('cuda')
-print('CUDA AVAILABLE', torch.cuda.is_available())
 
 LIGHT_PURPLE=(0.25098039,  0.274117647,  0.65882353)
 
@@ -33,6 +32,31 @@ model.eval()
 
 detector = YOLO('./pretrained_models/detector.pt').to(device)
 
+def render_reconstruction(image, conf, IoU_threshold=0.5): 
+    input_img, num_dets, reconstructions = run_wilow_model(image, conf, IoU_threshold=0.5)
+    if num_dets> 0: 
+    # Render front view
+    
+        misc_args = dict(
+            mesh_base_color=LIGHT_PURPLE,
+            scene_bg_color=(1, 1, 1),
+            focal_length=reconstructions['focal'],
+        )
+
+        cam_view = renderer.render_rgba_multiple(reconstructions['verts'], 
+                                                 cam_t=reconstructions['cam_t'], 
+                                                 render_res=reconstructions['img_size'], 
+                                                 is_right=reconstructions['right'], **misc_args)
+
+        # Overlay image
+        
+        input_img = np.concatenate([input_img, np.ones_like(input_img[:,:,:1])], axis=2) # Add alpha channel
+        input_img_overlay = input_img[:,:,:3] * (1-cam_view[:,:,3:]) + cam_view[:,:,:3] * cam_view[:,:,3:]
+
+        return input_img_overlay, f'{num_dets} hands detected'  
+    else: 
+        return input_img, f'{num_dets} hands detected' 
+
 @spaces.GPU()
 def run_wilow_model(image, conf, IoU_threshold=0.5):
     img_cv2 = image[...,::-1]
@@ -75,6 +99,8 @@ def run_wilow_model(image, conf, IoU_threshold=0.5):
     
             with torch.no_grad():
                 out = model(batch) 
+                
+            print('CUDA AVAILABLE', torch.cuda.is_available())
             print(out['pred_vertices'])
             multiplier    = (2*batch['right']-1)
             pred_cam      = out['pred_cam']
@@ -84,12 +110,7 @@ def run_wilow_model(image, conf, IoU_threshold=0.5):
             img_size      = batch["img_size"].float()
             scaled_focal_length = model_cfg.EXTRA.FOCAL_LENGTH / model_cfg.MODEL.IMAGE_SIZE * img_size.max()
             pred_cam_t_full     = cam_crop_to_full(pred_cam, box_center, box_size, img_size, scaled_focal_length).detach().cpu().numpy()
-            
-            # Render the result
-            all_verts  = []
-            all_cam_t  = []
-            all_right  = []
-            all_joints = []
+
             
             batch_size = batch['img'].shape[0]
             for n in range(batch_size):
@@ -107,25 +128,11 @@ def run_wilow_model(image, conf, IoU_threshold=0.5):
                 all_cam_t.append(cam_t)
                 all_right.append(is_right)
                 all_joints.append(joints)
-            # Render front view
-    
-            misc_args = dict(
-                mesh_base_color=LIGHT_PURPLE,
-                scene_bg_color=(1, 1, 1),
-                focal_length=scaled_focal_length,
-            )
-            print(all_verts[0])
-            print(all_cam_t[0])
-            cam_view = renderer.render_rgba_multiple(all_verts, cam_t=all_cam_t, render_res=img_size[n], is_right=all_right, **misc_args)
-    
-            # Overlay image
-            
-            input_img = img_vis.astype(np.float32)/255.0
-            input_img = np.concatenate([input_img, np.ones_like(input_img[:,:,:1])], axis=2) # Add alpha channel
-            input_img_overlay = input_img[:,:,:3] * (1-cam_view[:,:,3:]) + cam_view[:,:,:3] * cam_view[:,:,3:]
-            
-    image = img_vis #input_img_overlay
-    return image, f'{len(detections)} hands detected'        
+
+        reconstructions = {'verts': all_verts, 'cam_t': all_cam_t, 'right': all_right, 'img_size': img_size[n], 'focal': scaled_focal_length}
+        return img_vis.astype(np.float32)/255.0, len(detections), reconstructions
+    else: 
+        return img_vis.astype(np.float32)/255.0, len(detections), None       
 
 
 
@@ -168,7 +175,7 @@ with gr.Blocks(title="WiLoR: End-to-end 3D hand localization and reconstruction
             reconstruction = gr.Image(label="Reconstructions", type="numpy")
             hands_detected = gr.Textbox(label="Hands Detected")
     
-        submit.click(fn=run_wilow_model, inputs=[input_image, threshold], outputs=[reconstruction, hands_detected])
+        submit.click(fn=render_reconstruction, inputs=[input_image, threshold], outputs=[reconstruction, hands_detected])
 
     with gr.Row():