Merge e60b836 into 71d74de

Author: Zheng-LinXiao

configs/_base_/datasets/coco_openpose.py

@@ -0,0 +1,157 @@
+dataset_info = dict(
+    dataset_name='coco_openpose',
+    paper_info=dict(
+        author='Zhe, Cao and Tomas, Simon and '
+        'Shih-En, Wei and Yaser, Sheikh',
+        title='OpenPose: Realtime Multi-Person 2D Pose '
+        'Estimation using Part Affinity Fields',
+        container='IEEE Transactions on Pattern Analysis '
+        'and Machine Intelligence',
+        year='2019',
+        homepage='https://github.com/CMU-Perceptual-Computing-Lab/openpose/',
+    ),
+    keypoint_info={
+        0:
+        dict(name='nose', id=0, color=[255, 0, 85], type='upper', swap=''),
+        1:
+        dict(name='neck', id=1, color=[255, 0, 0], type='upper', swap=''),
+        2:
+        dict(
+            name='right_shoulder',
+            id=2,
+            color=[255, 85, 0],
+            type='upper',
+            swap='left_shoulder'),
+        3:
+        dict(
+            name='right_elbow',
+            id=3,
+            color=[255, 170, 0],
+            type='upper',
+            swap='left_elbow'),
+        4:
+        dict(
+            name='right_wrist',
+            id=4,
+            color=[255, 255, 0],
+            type='upper',
+            swap='left_wrist'),
+        5:
+        dict(
+            name='left_shoulder',
+            id=5,
+            color=[170, 255, 0],
+            type='upper',
+            swap='right_shoulder'),
+        6:
+        dict(
+            name='left_elbow',
+            id=6,
+            color=[85, 255, 0],
+            type='upper',
+            swap='right_elbow'),
+        7:
+        dict(
+            name='left_wrist',
+            id=7,
+            color=[0, 255, 0],
+            type='upper',
+            swap='right_wrist'),
+        8:
+        dict(
+            name='right_hip',
+            id=8,
+            color=[255, 0, 170],
+            type='lower',
+            swap='left_hip'),
+        9:
+        dict(
+            name='right_knee',
+            id=9,
+            color=[255, 0, 255],
+            type='lower',
+            swap='left_knee'),
+        10:
+        dict(
+            name='right_ankle',
+            id=10,
+            color=[170, 0, 255],
+            type='lower',
+            swap='left_ankle'),
+        11:
+        dict(
+            name='left_hip',
+            id=11,
+            color=[85, 255, 0],
+            type='lower',
+            swap='right_hip'),
+        12:
+        dict(
+            name='left_knee',
+            id=12,
+            color=[0, 0, 255],
+            type='lower',
+            swap='right_knee'),
+        13:
+        dict(
+            name='left_ankle',
+            id=13,
+            color=[0, 85, 255],
+            type='lower',
+            swap='right_ankle'),
+        14:
+        dict(
+            name='right_eye',
+            id=14,
+            color=[0, 255, 170],
+            type='upper',
+            swap='left_eye'),
+        15:
+        dict(
+            name='left_eye',
+            id=15,
+            color=[0, 255, 255],
+            type='upper',
+            swap='right_eye'),
+        16:
+        dict(
+            name='right_ear',
+            id=16,
+            color=[0, 170, 255],
+            type='upper',
+            swap='left_ear'),
+        17:
+        dict(
+            name='left_ear',
+            id=17,
+            color=[0, 170, 255],
+            type='upper',
+            swap='right_ear'),
+    },
+    skeleton_info={
+        0: dict(link=('neck', 'right_shoulder'), id=0, color=[255, 0, 85]),
+        1: dict(link=('neck', 'left_shoulder'), id=1, color=[255, 0, 0]),
+        2:
+        dict(link=('right_shoulder', 'right_elbow'), id=2, color=[255, 85, 0]),
+        3:
+        dict(link=('right_elbow', 'right_wrist'), id=3, color=[255, 170, 0]),
+        4:
+        dict(link=('left_shoulder', 'left_elbow'), id=4, color=[255, 255, 0]),
+        5: dict(link=('left_elbow', 'left_wrist'), id=5, color=[170, 255, 0]),
+        6: dict(link=('neck', 'right_hip'), id=6, color=[85, 255, 0]),
+        7: dict(link=('right_hip', 'right_knee'), id=7, color=[0, 255, 0]),
+        8: dict(link=('right_knee', 'right_ankle'), id=8, color=[0, 255, 85]),
+        9: dict(link=('neck', 'left_hip'), id=9, color=[0, 255, 170]),
+        10: dict(link=('left_hip', 'left_knee'), id=10, color=[0, 255, 225]),
+        11: dict(link=('left_knee', 'left_ankle'), id=11, color=[0, 170, 255]),
+        12: dict(link=('neck', 'nose'), id=12, color=[0, 85, 255]),
+        13: dict(link=('nose', 'right_eye'), id=13, color=[0, 0, 255]),
+        14: dict(link=('right_eye', 'right_ear'), id=14, color=[255, 0, 170]),
+        15: dict(link=('nose', 'left_eye'), id=15, color=[170, 0, 255]),
+        16: dict(link=('left_eye', 'left_ear'), id=16, color=[255, 0, 255]),
+    },
+    joint_weights=[1.] * 18,
+    sigmas=[
+        0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062,
+        0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089, 0.082
+    ])

demo/image_demo.py

@@ -25,6 +25,29 @@ def parse_args():
         action='store_true',
         default=False,
         help='Whether to show the index of keypoints')
+    parser.add_argument(
+        '--skeleton-style',
+        default='mmpose',
+        type=str,
+        choices=['mmpose', 'openpose'],
+        help='Skeleton style selection')
+    parser.add_argument(
+        '--kpt-thr',
+        type=float,
+        default=0.3,
+        help='Visualizing keypoint thresholds')
+    parser.add_argument(
+        '--radius',
+        type=int,
+        default=3,
+        help='Keypoint radius for visualization')
+    parser.add_argument(
+        '--thickness',
+        type=int,
+        default=1,
+        help='Link thickness for visualization')
+    parser.add_argument(
+        '--alpha', type=float, default=0.8, help='The transparency of bboxes')
     parser.add_argument(
         '--show',
         action='store_true',
@@ -50,8 +73,13 @@ def main():
         cfg_options=cfg_options)
 
     # init visualizer
+    model.cfg.visualizer.radius = args.radius
+    model.cfg.visualizer.alpha = args.alpha
+    model.cfg.visualizer.line_width = args.thickness
+
     visualizer = VISUALIZERS.build(model.cfg.visualizer)
-    visualizer.set_dataset_meta(model.dataset_meta)
+    visualizer.set_dataset_meta(
+        model.dataset_meta, skeleton_style=args.skeleton_style)
 
     # inference a single image
     batch_results = inference_topdown(model, args.img)
@@ -65,8 +93,10 @@ def main():
         data_sample=results,
         draw_gt=False,
         draw_bbox=True,
+        kpt_thr=args.kpt_thr,
         draw_heatmap=args.draw_heatmap,
         show_kpt_idx=args.show_kpt_idx,
+        skeleton_style=args.skeleton_style,
         show=args.show,
         out_file=args.out_file)
 

demo/topdown_demo_with_mmdet.py

@@ -55,10 +55,11 @@ def process_one_image(args, img_path, detector, pose_estimator, visualizer,
         draw_heatmap=args.draw_heatmap,
         draw_bbox=args.draw_bbox,
         show_kpt_idx=args.show_kpt_idx,
+        skeleton_style=args.skeleton_style,
         show=args.show,
         wait_time=show_interval,
         out_file=out_file,
-        kpt_score_thr=args.kpt_thr)
+        kpt_thr=args.kpt_thr)
 
     # if there is no instance detected, return None
     return data_samples.get('pred_instances', None)
@@ -110,7 +111,10 @@ def main():
         default=0.3,
         help='IoU threshold for bounding box NMS')
     parser.add_argument(
-        '--kpt-thr', type=float, default=0.3, help='Keypoint score threshold')
+        '--kpt-thr',
+        type=float,
+        default=0.3,
+        help='Visualizing keypoint thresholds')
     parser.add_argument(
         '--draw-heatmap',
         action='store_true',
@@ -121,6 +125,12 @@ def main():
         action='store_true',
         default=False,
         help='Whether to show the index of keypoints')
+    parser.add_argument(
+        '--skeleton-style',
+        default='mmpose',
+        type=str,
+        choices=['mmpose', 'openpose'],
+        help='Skeleton style selection')
     parser.add_argument(
         '--radius',
         type=int,
@@ -131,6 +141,8 @@ def main():
         type=int,
         default=1,
         help='Link thickness for visualization')
+    parser.add_argument(
+        '--alpha', type=float, default=0.8, help='The transparency of bboxes')
     parser.add_argument(
         '--draw-bbox', action='store_true', help='Draw bboxes of instances')
 
@@ -164,11 +176,14 @@ def main():
 
     # init visualizer
     pose_estimator.cfg.visualizer.radius = args.radius
+    pose_estimator.cfg.visualizer.alpha = args.alpha
     pose_estimator.cfg.visualizer.line_width = args.thickness
+
     visualizer = VISUALIZERS.build(pose_estimator.cfg.visualizer)
     # the dataset_meta is loaded from the checkpoint and
     # then pass to the model in init_pose_estimator
-    visualizer.set_dataset_meta(pose_estimator.dataset_meta)
+    visualizer.set_dataset_meta(
+        pose_estimator.dataset_meta, skeleton_style=args.skeleton_style)
 
     input_type = mimetypes.guess_type(args.input)[0].split('/')[0]
     if input_type == 'image':

mmpose/apis/inferencers/base_mmpose_inferencer.py

@@ -314,7 +314,7 @@ def visualize(self,
                 show=show,
                 wait_time=wait_time,
                 out_file=out_file,
-                kpt_score_thr=kpt_thr)
+                kpt_thr=kpt_thr)
             results.append(visualization)
 
             if show and not hasattr(self, '_window_close_cid'):

mmpose/engine/hooks/visualization_hook.py

@@ -50,15 +50,15 @@ def __init__(
         self,
         enable: bool = False,
         interval: int = 50,
-        score_thr: float = 0.3,
+        kpt_thr: float = 0.3,
         show: bool = False,
         wait_time: float = 0.,
         out_dir: Optional[str] = None,
         backend_args: Optional[dict] = None,
     ):
         self._visualizer: Visualizer = Visualizer.get_current_instance()
         self.interval = interval
-        self.score_thr = score_thr
+        self.kpt_thr = kpt_thr
         self.show = show
         if self.show:
             # No need to think about vis backends.
@@ -112,7 +112,7 @@ def after_val_iter(self, runner: Runner, batch_idx: int, data_batch: dict,
                 draw_heatmap=True,
                 show=self.show,
                 wait_time=self.wait_time,
-                kpt_score_thr=self.score_thr,
+                kpt_thr=self.kpt_thr,
                 step=total_curr_iter)
 
     def after_test_iter(self, runner: Runner, batch_idx: int, data_batch: dict,
@@ -163,6 +163,6 @@ def after_test_iter(self, runner: Runner, batch_idx: int, data_batch: dict,
                 draw_bbox=True,
                 draw_heatmap=True,
                 wait_time=self.wait_time,
-                kpt_score_thr=self.score_thr,
+                kpt_thr=self.kpt_thr,
                 out_file=out_file,
                 step=self._test_index)