[Feature] Add two traditional filters to smooth poses (open-mmlab#1127)

* add savgol&gaus1d filters

* Update __init__.py

* Update test_temporal_filter.py

* Update and rename gaus1d_filter.py to gauss1d_filter.py

* Update __init__.py

* pass pre-commit hooks

* Add smoother

* add Smoother
* refactor filters with base filter
* add unittest for Smoother

* add temporal filter unittests

* remove misadded file

* fix bugs in smoother and filters
update get_track_id and demos

* fix unittest

* fix smoother bug with empty input

* fix smoothing in 3d video demo

Co-authored-by: ly015 <liyining0712@gmail.com>

configs/_base_/filters/gaussian.py

@@ -0,0 +1,5 @@
+filter_cfg = dict(
+    type='GaussianFilter',
+    window_size=11,
+    sigma=4.0,
+)

configs/_base_/filters/one_euro.py

@@ -0,0 +1,5 @@
+filter_cfg = dict(
+    type='OneEuroFilter',
+    min_cutoff=0.004,
+    beta=0.7,
+)

configs/_base_/filters/savizky_golay.py

@@ -0,0 +1,5 @@
+filter_cfg = dict(
+    type='SavizkyGolayFilter',
+    window_size=11,
+    polyorder=2,
+)

demo/body3d_two_stage_video_demo.py

@@ -12,6 +12,7 @@
                          inference_pose_lifter_model,
                          inference_top_down_pose_model, init_pose_model,
                          process_mmdet_results, vis_3d_pose_result)
+from mmpose.core import Smoother
 
 try:
     from mmdet.apis import inference_detector, init_detector
@@ -123,10 +124,6 @@ def main():
         '--use-oks-tracking', action='store_true', help='Using OKS tracking')
     parser.add_argument(
         '--tracking-thr', type=float, default=0.3, help='Tracking threshold')
-    parser.add_argument(
-        '--euro',
-        action='store_true',
-        help='Using One_Euro_Filter for smoothing')
     parser.add_argument(
         '--radius',
         type=int,
@@ -137,6 +134,17 @@ def main():
         type=int,
         default=2,
         help='Link thickness for visualization')
+    parser.add_argument(
+        '--smooth',
+        action='store_true',
+        help='Apply a temporal filter to smooth the pose estimation results. '
+        'See also --smooth-filter-cfg.')
+    parser.add_argument(
+        '--smooth-filter-cfg',
+        type=str,
+        default='configs/_base_/filters/one_euro.py',
+        help='Config file of the filter to smooth the pose estimation '
+        'results. See also --smooth.')
 
     assert has_mmdet, 'Please install mmdet to run the demo.'
 
@@ -167,7 +175,7 @@ def main():
     pose_det_results_list = []
     next_id = 0
     pose_det_results = []
-    for frame in video:
+    for frame in mmcv.track_iter_progress(video):
         pose_det_results_last = pose_det_results
 
         # test a single image, the resulting box is (x1, y1, x2, y2)
@@ -194,9 +202,7 @@ def main():
             pose_det_results_last,
             next_id,
             use_oks=args.use_oks_tracking,
-            tracking_thr=args.tracking_thr,
-            use_one_euro=args.euro,
-            fps=video.fps)
+            tracking_thr=args.tracking_thr)
 
         pose_det_results_list.append(copy.deepcopy(pose_det_results))
 
@@ -237,6 +243,12 @@ def main():
     else:
         data_cfg = pose_lift_model.cfg.data_cfg
 
+    # build pose smoother for temporal refinement
+    if args.smooth:
+        smoother = Smoother(filter_cfg=args.smooth_filter_cfg, keypoint_dim=3)
+    else:
+        smoother = None
+
     num_instances = args.num_instances
     for i, pose_det_results in enumerate(
             mmcv.track_iter_progress(pose_det_results_list)):
@@ -279,6 +291,10 @@ def main():
             res['track_id'] = instance_id
             pose_lift_results_vis.append(res)
 
+        # Smoothing
+        if smoother:
+            pose_lift_results = smoother.smooth(pose_lift_results)
+
         # Visualization
         if num_instances < 0:
             num_instances = len(pose_lift_results_vis)

demo/bottom_up_pose_tracking_demo.py

@@ -7,6 +7,7 @@
 
 from mmpose.apis import (get_track_id, inference_bottom_up_pose_model,
                          init_pose_model, vis_pose_tracking_result)
+from mmpose.core import Smoother
 from mmpose.datasets import DatasetInfo
 
 
@@ -42,7 +43,19 @@ def main():
     parser.add_argument(
         '--euro',
         action='store_true',
-        help='Using One_Euro_Filter for smoothing')
+        help='(Deprecated, please use --smooth and --smooth-filter-cfg) '
+        'Using One_Euro_Filter for smoothing.')
+    parser.add_argument(
+        '--smooth',
+        action='store_true',
+        help='Apply a temporal filter to smooth the pose estimation results. '
+        'See also --smooth-filter-cfg.')
+    parser.add_argument(
+        '--smooth-filter-cfg',
+        type=str,
+        default='configs/_base_/filters/one_euro.py',
+        help='Config file of the filter to smooth the pose estimation '
+        'results. See also --smooth.')
     parser.add_argument(
         '--radius',
         type=int,
@@ -97,6 +110,20 @@ def main():
     # optional
     return_heatmap = False
 
+    # build pose smoother for temporal refinement
+    if args.euro:
+        warnings.warn(
+            'Argument --euro will be deprecated in the future. '
+            'Please use --smooth to enable temporal smoothing, and '
+            '--smooth-filter-cfg to set the filter config.',
+            DeprecationWarning)
+        smoother = Smoother(
+            filter_cfg='configs/_base_/filters/one_euro.py', keypoint_dim=2)
+    elif args.smooth:
+        smoother = Smoother(filter_cfg=args.smooth_filter_cfg, keypoint_dim=2)
+    else:
+        smoother = None
+
     # e.g. use ('backbone', ) to return backbone feature
     output_layer_names = None
     next_id = 0
@@ -126,6 +153,10 @@ def main():
             use_one_euro=args.euro,
             fps=fps)
 
+        # post-process the pose results with smoother
+        if smoother:
+            pose_results = smoother.smooth(pose_results)
+
         # show the results
         vis_img = vis_pose_tracking_result(
             pose_model,

demo/top_down_pose_tracking_demo_with_mmdet.py

@@ -8,6 +8,7 @@
 from mmpose.apis import (get_track_id, inference_top_down_pose_model,
                          init_pose_model, process_mmdet_results,
                          vis_pose_tracking_result)
+from mmpose.core import Smoother
 from mmpose.datasets import DatasetInfo
 
 try:
@@ -59,7 +60,19 @@ def main():
     parser.add_argument(
         '--euro',
         action='store_true',
-        help='Using One_Euro_Filter for smoothing')
+        help='(Deprecated, please use --smooth and --smooth-filter-cfg) '
+        'Using One_Euro_Filter for smoothing.')
+    parser.add_argument(
+        '--smooth',
+        action='store_true',
+        help='Apply a temporal filter to smooth the pose estimation results. '
+        'See also --smooth-filter-cfg.')
+    parser.add_argument(
+        '--smooth-filter-cfg',
+        type=str,
+        default='configs/_base_/filters/one_euro.py',
+        help='Config file of the filter to smooth the pose estimation '
+        'results. See also --smooth.')
     parser.add_argument(
         '--radius',
         type=int,
@@ -122,6 +135,20 @@ def main():
     # e.g. use ('backbone', ) to return backbone feature
     output_layer_names = None
 
+    # build pose smoother for temporal refinement
+    if args.euro:
+        warnings.warn(
+            'Argument --euro will be deprecated in the future. '
+            'Please use --smooth to enable temporal smoothing, and '
+            '--smooth-filter-cfg to set the filter config.',
+            DeprecationWarning)
+        smoother = Smoother(
+            filter_cfg='configs/_base_/filters/one_euro.py', keypoint_dim=2)
+    elif args.smooth:
+        smoother = Smoother(filter_cfg=args.smooth_filter_cfg, keypoint_dim=2)
+    else:
+        smoother = None
+
     next_id = 0
     pose_results = []
     while (cap.isOpened()):
@@ -154,9 +181,11 @@ def main():
             pose_results_last,
             next_id,
             use_oks=args.use_oks_tracking,
-            tracking_thr=args.tracking_thr,
-            use_one_euro=args.euro,
-            fps=fps)
+            tracking_thr=args.tracking_thr)
+
+        # post-process the pose results with smoother
+        if smoother:
+            pose_results = smoother.smooth(pose_results)
 
         # show the results
         vis_img = vis_pose_tracking_result(

demo/top_down_pose_tracking_demo_with_mmtracking.py

@@ -7,6 +7,7 @@
 
 from mmpose.apis import (inference_top_down_pose_model, init_pose_model,
                          vis_pose_tracking_result)
+from mmpose.core import Smoother
 from mmpose.datasets import DatasetInfo
 
 try:
@@ -78,6 +79,17 @@ def main():
         type=int,
         default=1,
         help='Link thickness for visualization')
+    parser.add_argument(
+        '--smooth',
+        action='store_true',
+        help='Apply a temporal filter to smooth the pose estimation results. '
+        'See also --smooth-filter-cfg.')
+    parser.add_argument(
+        '--smooth-filter-cfg',
+        type=str,
+        default='configs/_base_/filters/one_euro.py',
+        help='Config file of the filter to smooth the pose estimation '
+        'results. See also --smooth.')
 
     assert has_mmtrack, 'Please install mmtrack to run the demo.'
 
@@ -127,6 +139,12 @@ def main():
     # e.g. use ('backbone', ) to return backbone feature
     output_layer_names = None
 
+    # build pose smoother for temporal refinement
+    if args.smooth:
+        smoother = Smoother(filter_cfg=args.smooth_filter_cfg, keypoint_dim=2)
+    else:
+        smoother = None
+
     frame_id = 0
     while (cap.isOpened()):
         flag, img = cap.read()
@@ -151,6 +169,9 @@ def main():
             return_heatmap=return_heatmap,
             outputs=output_layer_names)
 
+        if smoother:
+            pose_results = smoother.smooth(pose_results)
+
         # show the results
         vis_img = vis_pose_tracking_result(
             pose_model,

demo/webcam_demo.py

@@ -10,7 +10,7 @@
 
 from mmpose.apis import (get_track_id, inference_top_down_pose_model,
                          init_pose_model, vis_pose_result)
-from mmpose.core import apply_bugeye_effect, apply_sunglasses_effect
+from mmpose.core import Smoother, apply_bugeye_effect, apply_sunglasses_effect
 from mmpose.utils import StopWatch
 
 try:
@@ -152,6 +152,18 @@ def parse_args():
         help='Enable synchronous mode that video I/O and inference will be '
         'temporally aligned. Note that this will reduce the display FPS.')
 
+    parser.add_argument(
+        '--smooth',
+        action='store_true',
+        help='Apply a temporal filter to smooth the pose estimation results. '
+        'See also --smooth-filter-cfg.')
+    parser.add_argument(
+        '--smooth-filter-cfg',
+        type=str,
+        default='configs/_base_/filters/one_euro.py',
+        help='Config file of the filter to smooth the pose estimation '
+        'results. See also --smooth.')
+
     return parser.parse_args()
 
 
@@ -265,8 +277,9 @@ def inference_pose():
             ts_input, frame, t_info, mmdet_results = det_result_queue.popleft()
 
         pose_results_list = []
-        for model_info, pose_history in zip(pose_model_list,
-                                            pose_history_list):
+        for model_info, pose_history, smoother in zip(pose_model_list,
+                                                      pose_history_list,
+                                                      pose_smoother_list):
             model_name = model_info['name']
             pose_model = model_info['model']
             cat_ids = model_info['cat_ids']
@@ -295,9 +308,10 @@ def inference_pose():
                     pose_results_last,
                     next_id,
                     use_oks=False,
-                    tracking_thr=0.3,
-                    use_one_euro=True,
-                    fps=None)
+                    tracking_thr=0.3)
+
+                if smoother:
+                    pose_results = smoother.smooth(pose_results)
 
                 pose_results_list.append(pose_results)
 
@@ -497,6 +511,7 @@ def main():
     global pose_result_queue, pose_result_queue_mutex
     global det_model, pose_model_list, pose_history_list
     global event_exit, event_inference_done
+    global pose_smoother_list
 
     args = parse_args()
 
@@ -540,6 +555,16 @@ def main():
     for _ in range(len(pose_model_list)):
         pose_history_list.append({'pose_results_last': [], 'next_id': 0})
 
+    # build pose smoother for temporal refinement
+    pose_smoother_list = []
+    for _ in range(len(pose_model_list)):
+        if args.smooth:
+            smoother = Smoother(
+                filter_cfg=args.smooth_filter_cfg, keypoint_dim=2)
+        else:
+            smoother = None
+        pose_smoother_list.append(smoother)
+
     # frame buffer
     if args.buffer_size > 0:
         buffer_size = args.buffer_size

mmpose/apis/inference_tracking.py

@@ -177,8 +177,9 @@ def get_track_id(results,
     Args:
         results (list[dict]): The bbox & pose results of the current frame
             (bbox_result, pose_result).
-        results_last (list[dict]): The bbox & pose & track_id info of the
-            last frame (bbox_result, pose_result, track_id).
+        results_last (list[dict], optional): The bbox & pose & track_id info
+            of the last frame (bbox_result, pose_result, track_id). None is
+            equivalent to an empty result list. Default: None
         next_id (int): The track id for the new person instance.
         min_keypoints (int): Minimum number of keypoints recognized as person.
             default: 3.
@@ -194,6 +195,18 @@ def get_track_id(results,
             current frame (bbox_result, pose_result, track_id).
         - next_id (int): The track id for the new person instance.
     """
+    if use_one_euro:
+        warnings.warn(
+            'In the future, get_track_id() will no longer perform '
+            'temporal refinement and the arguments `use_one_euro` and '
+            '`fps` will be deprecated. This part of function has been '
+            'migrated to Smoother (mmpose.core.Smoother). See '
+            'demo/top_down_pose_trackign_demo_with_mmdet.py for an '
+            'example.', DeprecationWarning)
+
+    if results_last is None:
+        results_last = []
+
     results = _get_area(results)
 
     if use_oks:
@@ -216,6 +229,7 @@ def get_track_id(results,
                 result['track_id'] = -1
         else:
             result['track_id'] = track_id
+
         if use_one_euro:
             result['keypoints'] = _temporal_refine(
                 result, match_result, fps=fps)