webcam demo

README.md

@@ -1,8 +1,12 @@
-# CenterHMR: a bottom-up single-shot method for multi-person 3D mesh recovery from a single image
+# CenterHMR: a bottom-up single-shot method for real-time multi-person 3D mesh recovery from a single image
 [![Google Colab demo](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1oz9E6uIbj4udOPZvA1Zi9pFx0SWH_UXg#scrollTo=s8gFtokdcEQo)
 
 The method achieves ECCV 2020 3DPW Challenge Runner Up. Please refer to [arxiv paper](https://arxiv.org/abs/2008.12272) for the details!
 
+### Update
+**2020/9/11: Real-time webcam demo using local/remote server.** Please refer to [config_guide.md](src/config_guide.md) for details.
+**2020/9/4:** Google Colab demo. Predicted results would be saved to a npy file per imag, please refer to [config_guide.md](src/config_guide.md) for details.
+
 <p float="center">
   <img src="../assets/demo/animation/c1_results_compressed.gif" width="32%" />
   <img src="../assets/demo/animation/c5_results_compressed.gif" width="32%" />
@@ -15,9 +19,6 @@ The method achieves ECCV 2020 3DPW Challenge Runner Up. Please refer to [arxiv p
   <img src="../assets/demo/animation/c3_results_compressed.gif" width="32%" />
 </p>
 
-### Update
-**2020/9/4:** Google Colab demo. Predicted results would be saved to a npy file per imag, please refer to [config_guide.md](src/config_guide.md) for details.
-
 ### Try on Google Colab
 Before installation, you can take a few minutes to try the prepared [Google Colab demo](https://colab.research.google.com/drive/1oz9E6uIbj4udOPZvA1Zi9pFx0SWH_UXg#scrollTo=s8gFtokdcEQo) a try.  
 It allows you to run the project in the cloud, free of charge. 
@@ -85,21 +86,38 @@ CUDA_VISIBLE_DEVICES=0 python core/test.py --gpu=0 --configs_yml=configs/basic_t
 ```
 Results will be saved in CenterHMR/demo/images_results.
 
-#### Internet images
+#### Internet images/videos
+
 You can also run the code on random internet images via putting the images under CenterHMR/demo/images before running sh run.sh.
 
 Or please refer to [config_guide.md](src/config_guide.md) for detail configurations.
 
 Please refer to [config_guide.md](src/config_guide.md) for **saving the estimated mesh/Center maps/parameters dict**.
 
+#### Webcam
+
+We also provide the webcam demo code, which can run at real-time on a 1070Ti GPU / remote server.  
+Currently, limited by the visualization pipeline, the code only support the single-person mesh recovery.
+
+To do this you just need to run
+```bash
+cd CenterHMR/src
+CUDA_VISIBLE_DEVICES=0 python core/test.py --gpu=0 --configs_yml=configs/basic_webcam.yml
+# or please set the TEST_MODE=0 WEBCAM_MODE=1 in run.sh, then run
+sh run.sh
+```
+Press Up/Down to end the demo. Pelease refer to [config_guide.md](src/config_guide.md) for setting mesh color or camera id.
+
+If you wish to run webcam demo using remote server, pelease refer to [config_guide.md](src/config_guide.md).
 
 ## TODO LIST
 
 The code will be gradually open sourced according to:
 - [ ] the schedule
-  - [x] demo code for internet images or videos
-  - [ ] evaluation code for re-implementation the results on 3DPW Challenge (really close)
-  - [ ] runtime optimization
+  - [x] demo code for internet images / videos / webcam
+  - [x] runtime optimization
+  - [ ] benchmark evaluation
+  - [ ] training
 
 ## Citation
 Please considering citing 
@@ -119,7 +137,8 @@ We thank [Peng Cheng](https://github.com/CPFLAME) for his constructive comments
 
 Here are some great resources we benefit:
 
-- SMPL models and layer is from [SMPL-X model](https://github.com/vchoutas/smplx).
+- SMPL models and layer is borrowed from MPII [SMPL-X model](https://github.com/vchoutas/smplx).
+- Webcam pipeline is borrowed from [minimal-hand](https://github.com/CalciferZh/minimal-hand).
 - Some functions are borrowed from [HMR-pytorch](https://github.com/MandyMo/pytorch_HMR).
 - Some functions for data augmentation are borrowed from [SPIN](https://github.com/nkolot/SPIN).
 - Synthetic occlusion is borrowed from [synthetic-occlusion](https://github.com/isarandi/synthetic-occlusion)

src/config_guide.md

@@ -53,3 +53,24 @@ If you installed Pytorch 1.6 or upper verion, then you can use the automatic mix
 ```bash
 model_precision: fp16
 ```
+#### Webcam setting 
+
+The webcam configure file is CenterHMR/src/configs/basic_webcam.yml
+
+###### webcam: whether run using webcam video
+
+###### cam_id: web camera id, default 0. 
+
+###### webcam_mesh_color: mesh color, default ghostwhite.
+
+Currently, we have LightCyan, ghostwhite, Azure, Cornislk, Honeydew, LavenderBlush. Feel free to paint the estimated mesh results in your favorite color.
+
+If your favorite color is not included, please add it to the mesh_color_dict (src/constants.py) and set the webcam_mesh_color.
+
+###### run_on_remote_server: whether run webcam (captured locally) demo on remote server 
+
+###### server_ip: IP address of remote server.
+
+###### server_port: Port of remote server, default 10086.
+
+Please change to the other port if 10086 has been used.

src/configs/basic_test.yml

@@ -22,3 +22,4 @@ ARGS:
  save_mesh: False
  save_centermap: True
  save_dict_results: True
+ webcam: False

src/configs/basic_test_video.yml

@@ -22,3 +22,4 @@ ARGS:
  save_mesh: False
  save_centermap: False
  save_dict_results: False
+ webcam: False

src/configs/basic_webcam.yml

@@ -0,0 +1,30 @@
+
+ARGS:
+ tab: 'webcam'
+ GPUS: 0 
+ multi_person: True
+ kernel_sizes: [5]
+ use_coordmaps: True
+ fine_tune: True
+ eval: False
+
+ baseline: 'hrnetv5'
+ input_size: 512
+ centermap_size: 64
+ model_precision: 'fp16'
+ val_batch_size: 2
+ gmodel_path: '../trained_models/pw3d_81.8_58.6.pkl'
+
+ # default: run on demo/images save at demo/images_results
+ # for video /path/to/project/CenterHMR/demo/videos/Messi_1
+ demo_image_folder: None
+ save_mesh: False
+ save_centermap: False
+ save_dict_results: False
+
+ webcam: True
+ cam_id: 0
+ webcam_mesh_color: 'ghostwhite'  #'LightCyan'
+ run_on_remote_server: False
+ server_ip: 'localhost'
+ server_port: 10086

src/constants.py

@@ -202,3 +202,6 @@ def joint_mapping(source_format, target_format):
         openpose25_2_smplx_map.append(i)
     else:
         openpose25_2_smplx_map.append(0)
+
+mesh_color_dict = {'LightCyan': [225,255,255], 'ghostwhite':[248, 248, 255], \
+'Azure':[240,255,255],'Cornislk':[255,248,220],'Honeydew':[240,255,240],'LavenderBlush':[255,240,245]}

src/core/base.py

@@ -58,7 +58,7 @@ def set_up_smplx(self):
         self.params_num = np.array(self.part_idx).sum()
         self.global_orient_nocam = torch.from_numpy(constants.global_orient_nocam).unsqueeze(0)
 
-    def _calc_smplx_params(self, param, data_3d):
+    def _calc_smplx_params(self, param):
         idx_list = [0]
         params_dict = {}
         # cam:4; poses: 87=3+63+6+6+3+3+3; expres: 10; shape: 10 = 111
@@ -95,10 +95,13 @@ def net_forward(self, data_3d, model, imgs=None,match_to_gt=False,mode='test'):
         params, center_maps, heatmap_AEs = model(imgs.contiguous())
 
         params, kps, data_3d, reorganize_idx = self.parse_maps(params, center_maps, heatmap_AEs, data_3d)
-        outputs = self._calc_smplx_params(params.contiguous(), data_3d)
+        if params is not None:
+            outputs = self._calc_smplx_params(params.contiguous())
+        else:
+            outputs = None
         return outputs, center_maps, kps, data_3d, reorganize_idx
 
-    def parse_maps(self,param_maps, center_maps, heatmap_AEs, data_3d):
+    def parse_maps(self,param_maps, center_maps, heatmap_AEs, data_3d=None):
         kps = heatmap_AEs
         centers_pred= []
         for batch_id in range(len(param_maps)):
@@ -111,12 +114,12 @@ def parse_maps(self,param_maps, center_maps, heatmap_AEs, data_3d):
             else:
                 centers_pred.append([])
 
-        info_vis = ['imgpath', 'image_org', 'offsets']
-        batch_size = param_maps.shape[0]
-        matched_data = {}
-        for key in info_vis:
-            matched_data[key] = []
         params_pred, reorganize_idx = [[] for i in range(2)]
+        if data_3d is not None:
+            info_vis = ['imgpath', 'image_org', 'offsets']
+            matched_data = {}
+            for key in info_vis:
+                matched_data[key] = []
 
         # while training, use gt center to extract the parameters from the estimated map
         # while evaluation, match the estimated center with the clostest gt center for parameter sampling.
@@ -125,22 +128,27 @@ def parse_maps(self,param_maps, center_maps, heatmap_AEs, data_3d):
             for person_id, center in enumerate(centers):
                 center_w, center_h = center.long()
                 params_pred.append(param_map[:,center_w,center_h])
-
-                for key in matched_data:
-                    data_gt = data_3d[key]
-                    if isinstance(data_gt, torch.Tensor):
-                        matched_data[key].append(data_gt[batch_id])
-                    elif isinstance(data_gt, list):
-                        matched_data[key].append(data_gt[batch_id])
                 reorganize_idx.append(batch_id)
 
-        params = torch.stack(params_pred)
-        for key in matched_data:
-            data_gt = data_3d[key]
-            if isinstance(data_gt, torch.Tensor):
-                data_3d[key] = torch.stack(matched_data[key])
-            elif isinstance(data_gt, list):
-                data_3d[key] = np.array(matched_data[key])
+                if data_3d is not None:
+                    for key in matched_data:
+                        data_gt = data_3d[key]
+                        if isinstance(data_gt, torch.Tensor):
+                            matched_data[key].append(data_gt[batch_id])
+                        elif isinstance(data_gt, list):
+                            matched_data[key].append(data_gt[batch_id])
+
+        if len(params_pred)>0:
+            params = torch.stack(params_pred)
+        else:
+            params = None
+        if data_3d is not None:
+            for key in matched_data:
+                data_gt = data_3d[key]
+                if isinstance(data_gt, torch.Tensor):
+                    data_3d[key] = torch.stack(matched_data[key])
+                elif isinstance(data_gt, list):
+                    data_3d[key] = np.array(matched_data[key])
 
         return params, kps, data_3d, np.array(reorganize_idx)
 
@@ -181,53 +189,4 @@ def load_model(self, path, model,prefix = 'module.',optimizer=None):
         else:
             print('model {} not exist!'.format(path))
         print('*'*20)
-        return model
-
-
-
-class Data_prefetcher():
-    def __init__(self, loader,fetch_item='image'):
-        self.loader = iter(loader)
-        self.stream = torch.cuda.Stream()
-        self.preload(fetch_item='image')
-    def preload(self,fetch_item='image'):
-        try:
-            self.next_input = next(self.loader)
-        except StopIteration:
-            self.next_input = None
-            self.next_image = None
-            return
-        with torch.cuda.stream(self.stream):
-            self.next_image = self.next_input[fetch_item].cuda(non_blocking=True)
-    def next(self,fetch_item='image'):
-        torch.cuda.current_stream().wait_stream(self.stream)
-        input, image = self.next_input, self.next_image
-        self.preload(fetch_item='image')
-        return input, image
-
-class MultiStreamBatchSampler(Sampler):
-    """Iterate two sets of indices
-    An 'epoch' is one iteration through the primary indices.
-    During the epoch, the secondary indices are iterated through
-    as many times as needed.
-    """
-    def __init__(self, datasets_list, batch_size_list):
-        self.primary_indices = primary_indices
-        self.secondary_indices = secondary_indices
-        self.secondary_batch_size = secondary_batch_size
-        self.primary_batch_size = batch_size - secondary_batch_size
-
-        assert len(self.primary_indices) >= self.primary_batch_size > 0
-        assert len(self.secondary_indices) >= self.secondary_batch_size > 0
-
-    def __iter__(self):
-        primary_iter = iterate_once(self.primary_indices)
-        secondary_iter = iterate_eternally(self.secondary_indices)
-        return (
-            primary_batch + secondary_batch
-            for (primary_batch, secondary_batch)
-            in  zip(grouper(primary_iter, self.primary_batch_size),
-                    grouper(secondary_iter, self.secondary_batch_size)))
-
-    def __len__(self):
-        return len(self.primary_indices) // self.primary_batch_size
+        return model

src/core/test.py

@@ -1,4 +1,7 @@
 from base import *
+from PIL import Image
+import torchvision
+
 
 class Demo(Base):
     def __init__(self):
@@ -9,11 +12,11 @@ def __init__(self):
         self.save_centermap = args.save_centermap
         self.save_dict_results = args.save_dict_results
         self.demo_dir = os.path.join(config.project_dir, 'demo')
+        self.generator.eval()
         print('Initialization finished!')
 
     def run(self, image_folder):
         vis_size = [1024,1024,3]#[1920,1080]
-        self.generator.eval()
         loader_val = self._create_single_data_loader(dataset='internet',train_flag=False, image_folder=image_folder)
         test_save_dir = image_folder+'_results'
         os.makedirs(test_save_dir,exist_ok=True)
@@ -66,13 +69,68 @@ def reorganize_results(self, outputs, img_paths, reorganize_idx,test_save_dir):
             # get the results: np.load('/path/to/person_overlap.npz',allow_pickle=True)['results'][()]
             np.savez(name, results=result_dict)
 
+    def single_image_forward(self,image):
+        image_size = image.shape[:2][::-1]
+        image_org = Image.fromarray(image)
+
+        resized_image_size = (float(self.input_size)/max(image_size) * np.array(image_size) // 2 * 2).astype(np.int)[::-1]
+        padding = tuple((self.input_size-resized_image_size)[::-1]//2)
+        transform = torchvision.transforms.Compose([
+            torchvision.transforms.Resize(resized_image_size, interpolation=3),
+            torchvision.transforms.Pad(padding, fill=0, padding_mode='constant'),
+            ])
+        image = torch.from_numpy(np.array(transform(image_org))).unsqueeze(0).cuda().contiguous().float()
+        outputs, centermaps, heatmap_AEs, _, reorganize_idx = self.net_forward(None,self.generator,image,mode='test')
+        return outputs
+
+    def webcam_run_local(self):
+        import keyboard
+        from utils.demo_utils import OpenCVCapture, Open3d_visualizer
+        capture = OpenCVCapture()
+        visualizer = Open3d_visualizer()
+
+        while True:
+            frame = capture.read()
+            if frame is None:
+                continue
+            with torch.no_grad():
+                outputs = self.single_image_forward(frame)
+            if outputs is not None:
+                verts = outputs['verts'][0].cpu().numpy()
+                verts = verts * 50 + np.array([0, 0, 100])
+                break_flag = visualizer.run(verts,frame)
+                if break_flag:
+                    break
+    def webcam_run_remote(self):
+        from utils.remote_server_utils import Server_port_receiver
+        capture = Server_port_receiver()
+
+        while True:
+            frame = capture.receive()
+            if isinstance(frame,list):
+                continue
+            with torch.no_grad():
+                outputs = self.single_image_forward(frame)
+            if outputs is not None:
+                verts = outputs['verts'][0].cpu().numpy()
+                verts = verts * 50 + np.array([0, 0, 100])
+                capture.send(verts)
+            else:
+                capture.send(['failed'])
+
 def main():
     demo = Demo()
-    # run the code on demo images
-    demo_image_folder = args.demo_image_folder
-    if not os.path.exists(demo_image_folder):
-        demo_image_folder = os.path.join(demo.demo_dir,'images')
-    demo.run(demo_image_folder)
+    if args.webcam:
+        if args.run_on_remote_server:
+            demo.webcam_run_remote()
+        else:
+            demo.webcam_run_local()
+    else:
+        # run the code on demo images
+        demo_image_folder = args.demo_image_folder
+        if not os.path.exists(demo_image_folder):
+            demo_image_folder = os.path.join(demo.demo_dir,'images')
+        demo.run(demo_image_folder)
 
 
 if __name__ == '__main__':

src/requirements.txt

@@ -12,3 +12,7 @@ pyrender
 shyaml
 PyYAML>=5.1.2
 numpy-quaternion
+open3d
+pygame
+keyboard
+transforms3d