support bottom-up demo

demo/bottom_up_img_demo.py

@@ -0,0 +1,78 @@
+import os
+from argparse import ArgumentParser
+
+from mmpose.apis import (inference_bottom_up_pose_model, init_pose_model,
+                         vis_pose_result)
+
+
+def main():
+    """Visualize the demo images."""
+    parser = ArgumentParser()
+    parser.add_argument('pose_config', help='Config file for detection')
+    parser.add_argument('pose_checkpoint', help='Checkpoint file')
+    parser.add_argument('--img-root', type=str, default='', help='Image root')
+    parser.add_argument(
+        '--json-file',
+        type=str,
+        default='',
+        help='Json file containing image info.')
+    parser.add_argument(
+        '--show',
+        action='store_true',
+        default=False,
+        help='whether to show img')
+    parser.add_argument(
+        '--out-img-root',
+        type=str,
+        default='',
+        help='Root of the output img file. '
+        'Default not saving the visualization images.')
+    parser.add_argument(
+        '--device', default='cuda:0', help='Device used for inference')
+    parser.add_argument(
+        '--kpt-thr', type=float, default=0.3, help='Keypoint score threshold')
+
+    args = parser.parse_args()
+
+    assert args.show or (args.out_img_root != '')
+
+    skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12],
+                [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3],
+                [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]
+
+    from pycocotools.coco import COCO
+    coco = COCO(args.json_file)
+    # build the pose model from a config file and a checkpoint file
+    pose_model = init_pose_model(
+        args.pose_config, args.pose_checkpoint, device=args.device)
+
+    img_keys = list(coco.imgs.keys())
+
+    # process each image
+    for i in range(len(img_keys)):
+        image_id = img_keys[i]
+        image = coco.loadImgs(image_id)[0]
+        image_name = os.path.join(args.img_root, image['file_name'])
+
+        # test a single image, with a list of bboxes.
+        pose_results = inference_bottom_up_pose_model(pose_model, image_name)
+
+        if args.out_img_root == '':
+            out_file = None
+        else:
+            os.makedirs(args.out_img_root, exist_ok=True)
+            out_file = os.path.join(args.out_img_root, f'vis_{i}.jpg')
+
+        # show the results
+        vis_pose_result(
+            pose_model,
+            image_name,
+            pose_results,
+            skeleton=skeleton,
+            kpt_score_thr=args.kpt_thr,
+            show=args.show,
+            out_file=out_file)
+
+
+if __name__ == '__main__':
+    main()

demo/bottom_up_video_demo.py

@@ -0,0 +1,93 @@
+import os
+from argparse import ArgumentParser
+
+import cv2
+
+from mmpose.apis import (inference_bottom_up_pose_model, init_pose_model,
+                         vis_pose_result)
+
+
+def main():
+    """Visualize the demo images."""
+    parser = ArgumentParser()
+    parser.add_argument('pose_config', help='Config file for pose')
+    parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
+    parser.add_argument('--video-path', type=str, help='Video path')
+    parser.add_argument(
+        '--show',
+        action='store_true',
+        default=False,
+        help='whether to show visualizations.')
+    parser.add_argument(
+        '--out-video-root',
+        default='',
+        help='Root of the output video file. '
+        'Default not saving the visualization video.')
+    parser.add_argument(
+        '--device', default='cuda:0', help='Device used for inference')
+    parser.add_argument(
+        '--kpt-thr', type=float, default=0.3, help='Keypoint score threshold')
+
+    args = parser.parse_args()
+
+    skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12],
+                [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3],
+                [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]
+
+    assert args.show or (args.out_video_root != '')
+
+    # build the pose model from a config file and a checkpoint file
+    pose_model = init_pose_model(
+        args.pose_config, args.pose_checkpoint, device=args.device)
+
+    cap = cv2.VideoCapture(args.video_path)
+
+    if args.out_video_root == '':
+        save_out_video = False
+    else:
+        os.makedirs(args.out_video_root, exist_ok=True)
+        save_out_video = True
+
+    if save_out_video:
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
+                int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        videoWriter = cv2.VideoWriter(
+            os.path.join(args.out_video_root,
+                         f'vis_{os.path.basename(args.video_path)}'), fourcc,
+            fps, size)
+
+    while (cap.isOpened()):
+        flag, img = cap.read()
+        if not flag:
+            break
+
+        pose_results = inference_bottom_up_pose_model(pose_model, img)
+
+        # show the results
+        vis_img = vis_pose_result(
+            pose_model,
+            img,
+            pose_results,
+            skeleton=skeleton,
+            kpt_score_thr=args.kpt_thr,
+            show=False)
+
+        if args.show:
+            cv2.imshow('Image', vis_img)
+
+        if save_out_video:
+            videoWriter.write(vis_img)
+
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            break
+
+    cap.release()
+    if save_out_video:
+        videoWriter.release()
+    cv2.destroyAllWindows()
+
+
+if __name__ == '__main__':
+    main()

demo/top_down_img_demo.py

@@ -1,7 +1,8 @@
 import os
 from argparse import ArgumentParser
 
-from mmpose.apis import inference_pose_model, init_pose_model, vis_pose_result
+from mmpose.apis import (inference_top_down_pose_model, init_pose_model,
+                         vis_pose_result)
 
 
 def main():
@@ -37,7 +38,6 @@ def main():
     args = parser.parse_args()
 
     assert args.show or (args.out_img_root != '')
-    assert 'cuda' in args.device
 
     skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12],
                 [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3],
@@ -68,7 +68,7 @@ def main():
             person_bboxes.append(bbox)
 
         # test a single image, with a list of bboxes.
-        pose_results = inference_pose_model(
+        pose_results = inference_top_down_pose_model(
             pose_model, image_name, person_bboxes, format='xywh')
 
         if args.out_img_root == '':

demo/top_down_img_demo_with_mmdet.py

@@ -3,7 +3,8 @@
 
 from mmdet.apis import inference_detector, init_detector
 
-from mmpose.apis import inference_pose_model, init_pose_model, vis_pose_result
+from mmpose.apis import (inference_top_down_pose_model, init_pose_model,
+                         vis_pose_result)
 
 
 def main():
@@ -47,7 +48,6 @@ def main():
 
     assert args.show or (args.out_img_root != '')
     assert args.img != ''
-    assert 'cuda' in args.device
     assert args.det_config is not None
     assert args.det_checkpoint is not None
 
@@ -66,7 +66,7 @@ def main():
     person_bboxes = det_results[0].copy()
 
     # test a single image, with a list of bboxes.
-    pose_results = inference_pose_model(
+    pose_results = inference_top_down_pose_model(
         pose_model,
         image_name,
         person_bboxes,

demo/top_down_video_demo_with_mmdet.py

@@ -4,7 +4,8 @@
 import cv2
 from mmdet.apis import inference_detector, init_detector
 
-from mmpose.apis import inference_pose_model, init_pose_model, vis_pose_result
+from mmpose.apis import (inference_top_down_pose_model, init_pose_model,
+                         vis_pose_result)
 
 
 def main():
@@ -45,7 +46,6 @@ def main():
                 [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]
 
     assert args.show or (args.out_video_root != '')
-    assert 'cuda' in args.device
     assert args.det_config is not None
     assert args.det_checkpoint is not None
 
@@ -83,7 +83,7 @@ def main():
         person_bboxes = det_results[0].copy()
 
         # test a single image, with a list of bboxes.
-        pose_results = inference_pose_model(
+        pose_results = inference_top_down_pose_model(
             pose_model,
             img,
             person_bboxes,

docs/getting_started.md

@@ -198,7 +198,7 @@ Assume that you have already downloaded the checkpoints to the directory `checkp
 
 
 
-### Top-down Image demo
+### Top-down image demo
 
 #### Using gt human bounding boxes as input
 
@@ -245,13 +245,13 @@ Examples:
 python demo/top_down_img_demo_with_mmdet.py mmdetection/configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py \
     mmdetection/checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth \
     configs/top_down/hrnet/coco/hrnet_w48_coco_256x192.py \
-    hrnet_w48_coco_256x192/epoch_210.pth \
+    hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth \
     --img-root tests/data/coco/ \
     --img 000000196141.jpg \
     --out-img-root vis_results
 ```
 
-### Top-down Video demo
+### Top-down video demo
 
 We also provide a video demo to illustrate the results.
 
@@ -273,7 +273,55 @@ Examples:
 python demo/top_down_video_demo_with_mmdet.py mmdetection/configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py \
     mmdetection/checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth \
     configs/top_down/hrnet/coco/hrnet_w48_coco_256x192.py \
-    hrnet_w48_coco_256x192/epoch_210.pth \
+    hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth \
+    --video_path demo/demo_video.mp4 \
+    --output-video-root vis_results
+```
+
+
+### Bottom-up image demo
+
+We provide a demo script to test a single image.
+
+```shell
+python demo/bottom_up_img_demo.py \
+    ${MMPOSE_CONFIG_FILE} ${MMPOSE_CHECKPOINT_FILE} \
+    --img-root ${IMG_ROOT} --json-file ${JSON_FILE} \
+    --out-img-root ${OUTPUT_DIR} \
+    [--show --device ${GPU_ID}] \
+    [--kpt-thr ${KPT_SCORE_THR}]
+```
+
+Examples:
+
+```shell
+python demo/bottom_up_img_demo.py \
+    configs/bottom_up/hrnet/coco/hrnet_w32_coco_512x512.py \
+    hrnet_w32_coco_512x512-bcb8c247_20200816.pth \
+    --img-root tests/data/coco/ --json-file tests/data/coco/test_coco.json \
+    --out-img-root vis_results
+```
+
+
+### Bottom-up video demo
+
+We also provide a video demo to illustrate the results.
+
+```shell
+python demo/bottom_up_video_demo.py \
+    ${MMPOSE_CONFIG_FILE} ${MMPOSE_CHECKPOINT_FILE} \
+    --video-path ${VIDEO_FILE} \
+    --output-video-root ${OUTPUT_VIDEO_ROOT} \
+    [--show --device ${GPU_ID}] \
+    [--bbox-thr ${BBOX_SCORE_THR} --kpt-thr ${KPT_SCORE_THR}]
+```
+
+Examples:
+
+```shell
+python demo/bottom_up_video_demo.py \
+    configs/bottom_up/hrnet/coco/hrnet_w32_coco_512x512.py \
+    hrnet_w32_coco_512x512-bcb8c247_20200816.pth \
     --video_path demo/demo_video.mp4 \
     --output-video-root vis_results
 ```

mmpose/apis/__init__.py

@@ -1,8 +1,11 @@
-from .inference import inference_pose_model, init_pose_model, vis_pose_result
+from .inference import (inference_bottom_up_pose_model,
+                        inference_top_down_pose_model, init_pose_model,
+                        vis_pose_result)
 from .test import multi_gpu_test, single_gpu_test
 from .train import train_model
 
 __all__ = [
-    'train_model', 'init_pose_model', 'inference_pose_model', 'multi_gpu_test',
-    'single_gpu_test', 'vis_pose_result'
+    'train_model', 'init_pose_model', 'inference_top_down_pose_model',
+    'inference_bottom_up_pose_model', 'multi_gpu_test', 'single_gpu_test',
+    'vis_pose_result'
 ]