[Fix] Fix the OOM error for demo (#134)

* fix script name typo

* fix the OOM Error for image demo

* fix the OOM Error for video demo

* change the default argument value fof `img_format`

* release GPU memorry

demo/estimate_smpl_image.py

@@ -1,12 +1,16 @@
+import os
+import os.path as osp
 from argparse import ArgumentParser
+from pathlib import Path
 
 import mmcv
 import numpy as np
+import torch
 
 from mmhuman3d.apis import inference_image_based_model, init_model
-from mmhuman3d.core.visualization import visualize_smpl_calibration
+from mmhuman3d.core.visualization import visualize_smpl_hmr
+from mmhuman3d.utils import array_to_images
 from mmhuman3d.utils.demo_utils import (
-    conver_verts_to_cam_coord,
     prepare_frames,
     process_mmdet_results,
     process_mmtracking_results,
@@ -70,8 +74,7 @@ def single_person_with_mmdet(args, frames_iter):
             continue
         # vis bboxes
         if args.draw_bbox:
-            bboxes = [res['bbox'] for res in person_results]
-            bboxes = np.vstack(bboxes)
+            bboxes = person_results[0]['bbox'][None]
             mmcv.imshow_bboxes(
                 frame, bboxes, top_k=-1, thickness=2, show=False)
         img_index.append(i)
@@ -84,29 +87,33 @@ def single_person_with_mmdet(args, frames_iter):
     verts = np.array(verts)
     bboxes_xyxy = np.array(bboxes_xyxy)
 
-    # convert vertices from world to camera
-    verts, K0 = conver_verts_to_cam_coord(
-        verts, pred_cams, bboxes_xyxy, focal_length=args.focal_length)
+    del mesh_model
+    del person_det_model
+    torch.cuda.empty_cache()
 
     # smooth
     if args.smooth_type is not None:
         verts = smooth_process(verts, smooth_type=args.smooth_type)
 
     if args.show_path is not None:
+        frames_folder = osp.join(Path(args.show_path).parent, 'images')
+        os.makedirs(frames_folder, exist_ok=True)
+        array_to_images(
+            np.array(frames_iter)[img_index], output_folder=frames_folder)
         body_model_config = dict(model_path=args.body_model_dir, type='smpl')
         # Visualization
-        visualize_smpl_calibration(
+        visualize_smpl_hmr(
+            cam_transl=pred_cams,
+            bbox=bboxes_xyxy,
             verts=verts,
             output_path=args.show_path,
             render_choice=args.render_choice,
             resolution=frames_iter[0].shape[:2],
-            image_array=np.array(frames_iter)[img_index],
+            origin_frames=frames_folder,
             body_model_config=body_model_config,
-            K=K0,
-            R=None,
-            T=None,
             overwrite=True,
-            palette=args.palette)
+            palette=args.palette,
+            read_frames_batch=True)
 
 
 def multi_person_with_mmtracking(args, frames_iter):
@@ -183,35 +190,46 @@ def multi_person_with_mmtracking(args, frames_iter):
 
         track_ids_lists.append(track_ids)
 
-    verts, K0 = conver_verts_to_cam_coord(
-        verts, pred_cams, bboxes_xyxy, focal_length=args.focal_length)
+    del mesh_model
+    del tracking_model
+    torch.cuda.empty_cache()
 
     # smooth
     if args.smooth_type is not None:
         verts = smooth_process(verts, smooth_type=args.smooth_type)
 
     # To compress vertices array
-    V = np.zeros([frame_num, max_instance, 6890, 3])
+    compressed_verts = np.zeros([frame_num, max_instance, 6890, 3])
+    compressed_cams = np.zeros([frame_num, max_instance, 3])
+    compressed_bboxs = np.zeros([frame_num, max_instance, 5])
     for i, track_ids_list in enumerate(
             mmcv.track_iter_progress(track_ids_lists)):
         instance_num = len(track_ids_list)
-        V[i, :instance_num] = verts[i, track_ids_list]
+        compressed_verts[i, :instance_num] = verts[i, track_ids_list]
+        compressed_cams[i, :instance_num] = pred_cams[i, track_ids_list]
+        compressed_bboxs[i, :instance_num] = bboxes_xyxy[i, track_ids_list]
 
     if args.show_path is not None:
+        frames_folder = osp.join(Path(args.show_path).parent, 'images')
+        os.makedirs(frames_folder, exist_ok=True)
+        array_to_images(
+            np.array(frames_iter)[img_index], output_folder=frames_folder)
+
         body_model_config = dict(model_path=args.body_model_dir, type='smpl')
+
         # Visualization
-        visualize_smpl_calibration(
-            verts=V,
+        visualize_smpl_hmr(
+            cam_transl=compressed_cams,
+            bbox=compressed_bboxs,
+            verts=compressed_verts,
             output_path=args.show_path,
             render_choice=args.render_choice,
             resolution=frames_iter[0].shape[:2],
-            image_array=np.array(frames_iter)[img_index],
+            origin_frames=frames_folder,
             body_model_config=body_model_config,
-            K=K0,
-            R=None,
-            T=None,
             overwrite=True,
-            palette=args.palette)
+            palette=args.palette,
+            read_frames_batch=True)
 
 
 def main(args):

demo/etimate_smpl_video.py → demo/estimate_smpl_video.py

@@ -1,14 +1,19 @@
+import os
+import os.path as osp
 from argparse import ArgumentParser
+from pathlib import Path
 
 import mmcv
 import numpy as np
+import torch
 
 from mmhuman3d.apis import (
     feature_extract,
     inference_video_based_model,
     init_model,
 )
 from mmhuman3d.core.visualization import visualize_smpl_vibe
+from mmhuman3d.utils import array_to_images
 from mmhuman3d.utils.demo_utils import (
     extract_feature_sequence,
     prepare_frames,
@@ -102,11 +107,21 @@ def single_person_with_mmdet(args, frames_iter):
     verts = np.array(verts)
     bboxes_xyxy = np.array(bboxes_xyxy)
 
+    del mesh_model
+    del extractor
+    del person_det_model
+    torch.cuda.empty_cache()
+
     # smooth
     if args.smooth_type is not None:
         verts = smooth_process(verts, smooth_type=args.smooth_type)
 
     if args.show_path is not None:
+        frames_folder = osp.join(Path(args.show_path).parent, 'images')
+        os.makedirs(frames_folder, exist_ok=True)
+        array_to_images(
+            np.array(frames_iter)[img_index], output_folder=frames_folder)
+
         body_model_config = dict(model_path=args.body_model_dir, type='smpl')
         visualize_smpl_vibe(
             verts=verts,
@@ -115,10 +130,11 @@ def single_person_with_mmdet(args, frames_iter):
             output_path=args.show_path,
             render_choice=args.render_choice,
             resolution=frames_iter[0].shape[:2],
-            image_array=np.array(frames_iter)[img_index],
+            origin_frames=frames_folder,
             body_model_config=body_model_config,
             overwrite=True,
-            palette=args.palette)
+            palette=args.palette,
+            read_frames_batch=True)
 
 
 def multi_person_with_mmtracking(args, frames_iter):
@@ -200,35 +216,44 @@ def multi_person_with_mmtracking(args, frames_iter):
 
         track_ids_lists.append(track_ids)
 
+    del mesh_model
+    del extractor
+    del tracking_model
+    torch.cuda.empty_cache()
+
     # smooth
     if args.smooth_type is not None:
         verts = smooth_process(verts, smooth_type=args.smooth_type)
 
     # To compress vertices array
-    V = np.zeros([frame_num, max_instance, 6890, 3])
-    C = np.zeros([frame_num, max_instance, 3])
-    B = np.zeros([frame_num, max_instance, 5])
+    compressed_verts = np.zeros([frame_num, max_instance, 6890, 3])
+    compressed_cams = np.zeros([frame_num, max_instance, 3])
+    compressed_bboxs = np.zeros([frame_num, max_instance, 5])
     for i, track_ids_list in enumerate(track_ids_lists):
         instance_num = len(track_ids_list)
-        V[i, :instance_num] = verts[i, track_ids_list]
-        C[i, :instance_num] = pred_cams[i, track_ids_list]
-        B[i, :instance_num] = bboxes_xyxy[i, track_ids_list]
+        compressed_verts[i, :instance_num] = verts[i, track_ids_list]
+        compressed_cams[i, :instance_num] = pred_cams[i, track_ids_list]
+        compressed_bboxs[i, :instance_num] = bboxes_xyxy[i, track_ids_list]
     assert len(img_index) > 0
 
     if args.show_path is not None:
-
+        frames_folder = osp.join(Path(args.show_path).parent, 'images')
+        os.makedirs(frames_folder, exist_ok=True)
+        array_to_images(
+            np.array(frames_iter)[img_index], output_folder=frames_folder)
         body_model_config = dict(model_path=args.body_model_dir, type='smpl')
         visualize_smpl_vibe(
-            verts=verts,
-            pred_cam=pred_cams,
-            bbox=bboxes_xyxy,
+            verts=compressed_verts,
+            pred_cam=compressed_cams,
+            bbox=compressed_bboxs,
             output_path=args.show_path,
             render_choice=args.render_choice,
             resolution=frames_iter[0].shape[:2],
-            image_array=np.array(frames_iter)[img_index],
+            origin_frames=frames_folder,
             body_model_config=body_model_config,
             overwrite=True,
-            palette=args.palette)
+            palette=args.palette,
+            read_frames_batch=True)
 
 
 def main(args):

mmhuman3d/core/visualization/visualize_smpl.py

@@ -464,7 +464,7 @@ def render_smpl(
     origin_frames: Optional[str] = None,
     frame_list: Optional[List[str]] = None,
     image_array: Optional[Union[np.ndarray, torch.Tensor]] = None,
-    img_format: str = 'frame_%06d.jpg',
+    img_format: str = '%06d.png',
     overwrite: bool = False,
     mesh_file_path: Optional[str] = None,
     read_frames_batch: bool = False,