Update app.py

app.py

@@ -20,7 +20,7 @@ from wilor.models import WiLoR, load_wilor
 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('cuda') if torch.cuda.is_available() else torch.device('cpu')
+device = torch.device('cpu') #if torch.cuda.is_available() else torch.device('cpu')
 
 LIGHT_PURPLE=(0.25098039,  0.274117647,  0.65882353)
 
@@ -59,70 +59,70 @@ def run_wilow_model(image, conf, IoU_threshold=0.5):
         cv2.rectangle(img_vis, (int(Bbox[0]), int(Bbox[1]) - 20), (int(Bbox[0]) + w, int(Bbox[1])), color, -1)
         cv2.putText(img_vis, label, (int(Bbox[0]), int(Bbox[1]) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,0), 2)
           
-    # if len(bboxes) != 0: 
-    #     boxes = np.stack(bboxes)
-    #     right = np.stack(is_right)
-    #     dataset = ViTDetDataset(model_cfg, img_cv2, boxes, right, rescale_factor=2.0 )
-    #     dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False, num_workers=0)
+    if len(bboxes) != 0: 
+        boxes = np.stack(bboxes)
+        right = np.stack(is_right)
+        dataset = ViTDetDataset(model_cfg, img_cv2, boxes, right, rescale_factor=2.0 )
+        dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False, num_workers=0)
     
-    #     all_verts = []
-    #     all_cam_t = []
-    #     all_right = []
-    #     all_joints= []
+        all_verts = []
+        all_cam_t = []
+        all_right = []
+        all_joints= []
     
-    #     for batch in dataloader: 
-    #         batch = recursive_to(batch, device)
+        for batch in dataloader: 
+            batch = recursive_to(batch, device)
     
-    #         with torch.no_grad():
-    #             out = model(batch) 
+            with torch.no_grad():
+                out = model(batch) 
                 
-    #         multiplier    = (2*batch['right']-1)
-    #         pred_cam      = out['pred_cam']
-    #         pred_cam[:,1] = multiplier*pred_cam[:,1]
-    #         box_center    = batch["box_center"].float()
-    #         box_size      = batch["box_size"].float()
-    #         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()
+            multiplier    = (2*batch['right']-1)
+            pred_cam      = out['pred_cam']
+            pred_cam[:,1] = multiplier*pred_cam[:,1]
+            box_center    = batch["box_center"].float()
+            box_size      = batch["box_size"].float()
+            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 = []
+            # Render the result
+            all_verts  = []
+            all_cam_t  = []
+            all_right  = []
+            all_joints = []
             
-    #         batch_size = batch['img'].shape[0]
-    #         for n in range(batch_size):
+            batch_size = batch['img'].shape[0]
+            for n in range(batch_size):
                 
-    #             verts  = out['pred_vertices'][n].detach().cpu().numpy()
-    #             joints = out['pred_keypoints_3d'][n].detach().cpu().numpy()
+                verts  = out['pred_vertices'][n].detach().cpu().numpy()
+                joints = out['pred_keypoints_3d'][n].detach().cpu().numpy()
                 
-    #             is_right = batch['right'][n].cpu().numpy()
-    #             verts[:,0] = (2*is_right-1)*verts[:,0]
-    #             joints[:,0] = (2*is_right-1)*joints[:,0]
+                is_right = batch['right'][n].cpu().numpy()
+                verts[:,0] = (2*is_right-1)*verts[:,0]
+                joints[:,0] = (2*is_right-1)*joints[:,0]
                 
-    #             cam_t = pred_cam_t_full[n]
+                cam_t = pred_cam_t_full[n]
                 
-    #             all_verts.append(verts)
-    #             all_cam_t.append(cam_t)
-    #             all_right.append(is_right)
-    #             all_joints.append(joints)
-    #         # Render front view
+                all_verts.append(verts)
+                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,
-    #         )
-    #         cam_view = renderer.render_rgba_multiple(all_verts, cam_t=all_cam_t, render_res=img_size[n], is_right=all_right, **misc_args)
+            misc_args = dict(
+                mesh_base_color=LIGHT_PURPLE,
+                scene_bg_color=(1, 1, 1),
+                focal_length=scaled_focal_length,
+            )
+            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
+            # 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:]
+            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
+    image = input_img_overlay
     return image, f'{len(detections)} hands detected'