refine inferencer doc & support displaying alias (#2137)

open-mmlab · Apr 6, 2023 · 5fbd0e1 · 5fbd0e1
1 parent df773cd
commit 5fbd0e1
Show file tree

Hide file tree

Showing 6 changed files with 186 additions and 102 deletions.
diff --git a/configs/body_2d_keypoint/topdown_heatmap/coco/hrnet_coco.yml b/configs/body_2d_keypoint/topdown_heatmap/coco/hrnet_coco.yml
@@ -22,7 +22,7 @@ Models:
       AR: 0.799
       [email protected]: 0.942
     Task: Body 2D Keypoint
-  Weights: https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w32_coco_256x192-c78dce93_20200708.pth
+  Weights: https://download.openmmlab.com/mmpose/v1/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-256x192-81c58e40_20220909.pth
 - Config: configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-384x288.py
   In Collection: HRNet
   Metadata:
@@ -38,7 +38,7 @@ Models:
       AR: 0.81
       [email protected]: 0.943
     Task: Body 2D Keypoint
-  Weights: https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w32_coco_384x288-d9f0d786_20200708.pth
+  Weights: https://download.openmmlab.com/mmpose/v1/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-384x288-ca5956af_20220909.pth
 - Config: configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w48_8xb32-210e_coco-256x192.py
   In Collection: HRNet
   Metadata:
@@ -54,7 +54,7 @@ Models:
       AR: 0.806
       [email protected]: 0.942
     Task: Body 2D Keypoint
-  Weights: https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth
+  Weights: https://download.openmmlab.com/mmpose/v1/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w48_8xb32-210e_coco-256x192-0e67c616_20220913.pth
 - Config: configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w48_8xb32-210e_coco-384x288.py
   In Collection: HRNet
   Metadata:
@@ -70,7 +70,7 @@ Models:
       AR: 0.816
       [email protected]: 0.946
     Task: Body 2D Keypoint
-  Weights: https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_384x288-314c8528_20200708.pth
+  Weights: https://download.openmmlab.com/mmpose/v1/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w48_8xb32-210e_coco-384x288-c161b7de_20220915.pth
 - Config: configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-aic-256x192-merge.py
   In Collection: HRNet
   Metadata:

diff --git a/demo/inferencer_demo.py b/demo/inferencer_demo.py
@@ -1,13 +1,17 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from argparse import ArgumentParser
+from typing import Dict
 
-from mmpose.apis.inferencers import MMPoseInferencer
+from mmpose.apis.inferencers import MMPoseInferencer, get_model_aliases
 
 
 def parse_args():
     parser = ArgumentParser()
     parser.add_argument(
-        'inputs', type=str, help='Input image/video path or folder path.')
+        'inputs',
+        type=str,
+        nargs='?',
+        help='Input image/video path or folder path.')
     parser.add_argument(
         '--pose2d',
         type=str,
@@ -84,6 +88,10 @@ def parse_args():
         type=str,
         default='',
         help='Directory for saving inference results.')
+    parser.add_argument(
+        '--show-alias',
+        action='store_true',
+        help='Display all the available model aliases.')
 
     call_args = vars(parser.parse_args())
 
@@ -95,14 +103,30 @@ def parse_args():
     for init_kw in init_kws:
         init_args[init_kw] = call_args.pop(init_kw)
 
-    return init_args, call_args
+    diaplay_alias = call_args.pop('show_alias')
+
+    return init_args, call_args, diaplay_alias
+
+
+def display_model_aliases(model_aliases: Dict[str, str]) -> None:
+    """Display the available model aliases and their corresponding model
+    names."""
+    aliases = list(model_aliases.keys())
+    max_alias_length = max(map(len, aliases))
+    print(f'{"ALIAS".ljust(max_alias_length+2)}MODEL_NAME')
+    for alias in sorted(aliases):
+        print(f'{alias.ljust(max_alias_length+2)}{model_aliases[alias]}')
 
 
 def main():
-    init_args, call_args = parse_args()
-    inferencer = MMPoseInferencer(**init_args)
-    for _ in inferencer(**call_args):
-        pass
+    init_args, call_args, diaplay_alias = parse_args()
+    if diaplay_alias:
+        model_alises = get_model_aliases(init_args['scope'])
+        display_model_aliases(model_alises)
+    else:
+        inferencer = MMPoseInferencer(**init_args)
+        for _ in inferencer(**call_args):
+            pass
 
 
 if __name__ == '__main__':