[Fix] 3d pose demo with multiple instances (#2483)

demo/body3d_pose_lifter_demo.py

@@ -70,6 +70,13 @@ def parse_args():
         'scale of the dataset, and move the bbox (along with the 2D pose) to '
         'the average bbox center of the dataset. This is useful when bbox '
         'is small, especially in multi-person scenarios.')
+    parser.add_argument(
+        '--num-instances',
+        type=int,
+        default=-1,
+        help='The number of 3D poses to be visualized in every frame. If '
+        'less than 0, it will be set to the number of pose results in the '
+        'first frame.')
     parser.add_argument(
         '--output-root',
         type=str,
@@ -227,7 +234,6 @@ def get_pose_lift_results(args, visualizer, pose_lifter, pose_est_results_list,
 
         pred_instances = pose_lift_res.pred_instances
         keypoints = pred_instances.keypoints
-        # print(keypoints)
         keypoint_scores = pred_instances.keypoint_scores
         if keypoint_scores.ndim == 3:
             keypoint_scores = np.squeeze(keypoint_scores, axis=1)
@@ -253,6 +259,9 @@ def get_pose_lift_results(args, visualizer, pose_lifter, pose_est_results_list,
     pred_3d_data_samples = merge_data_samples(pose_lift_results)
     det_data_sample = merge_data_samples(pose_est_results)
 
+    if args.num_instances < 0:
+        args.num_instances = len(pose_lift_results)
+
     # Visualization
     if visualizer is not None:
         visualizer.add_datasample(
@@ -264,6 +273,7 @@ def get_pose_lift_results(args, visualizer, pose_lifter, pose_est_results_list,
             show=args.show,
             draw_bbox=True,
             kpt_thr=args.kpt_thr,
+            num_instances=args.num_instances,
             wait_time=args.show_interval)
 
     return pred_3d_data_samples.get('pred_instances', None)

demo/docs/en/3d_human_pose_demo.md

@@ -20,6 +20,7 @@ ${MMPOSE_CHECKPOINT_FILE_3D} \
 [--show] \
 [--rebase-keypoint-height] \
 [--norm-pose-2d] \
+[--num-instances] \
 [--output-root ${OUT_VIDEO_ROOT}] \
 [--save-predictions]
 [--save-predictions] \

mmpose/visualization/local_visualizer_3d.py

@@ -126,9 +126,11 @@ def _draw_3d_data_samples(
                 num_instances = 0
         else:
             if len(pred_instances) > num_instances:
+                pred_instances_ = InstanceData()
                 for k in pred_instances.keys():
-                    new_val = pred_instances.k[:num_instances]
-                    pose_samples.pred_instances.k = new_val
+                    new_val = pred_instances[k][:num_instances]
+                    pred_instances_.set_field(new_val, k)
+                pred_instances = pred_instances_
             elif num_instances < len(pred_instances):
                 num_instances = len(pred_instances)
 
@@ -464,6 +466,7 @@ def add_datasample(self,
                        draw_bbox: bool = False,
                        show_kpt_idx: bool = False,
                        skeleton_style: str = 'mmpose',
+                       num_instances: int = -1,
                        show: bool = False,
                        wait_time: float = 0,
                        out_file: Optional[str] = None,
@@ -499,6 +502,10 @@ def add_datasample(self,
                 Defaults to ``False``
             skeleton_style (str): Skeleton style selection. Defaults to
                 ``'mmpose'``
+            num_instances (int): Number of instances to be shown in 3D. If
+                smaller than 0, all the instances in the pose_result will be
+                shown. Otherwise, pad or truncate the pose_result to a length
+                of num_instances. Defaults to -1
             show (bool): Whether to display the drawn image. Default to
                 ``False``
             wait_time (float): The interval of show (s). Defaults to 0
@@ -524,7 +531,10 @@ def add_datasample(self,
                         det_img_data, det_data_sample.pred_instances)
 
         pred_img_data = self._draw_3d_data_samples(
-            image.copy(), data_sample, draw_gt=draw_gt)
+            image.copy(),
+            data_sample,
+            draw_gt=draw_gt,
+            num_instances=num_instances)
 
         # merge visualization results
         if det_img_data is not None and gt_img_data is not None: