[Enhance] Add Opencv backend support for Visualizer (open-mmlab#2283)

Tau-J · Apr 25, 2023 · 16d12a9 · 16d12a9
1 parent 9eb6c24
commit 16d12a9
Show file tree

Hide file tree

Showing 8 changed files with 515 additions and 282 deletions.
diff --git a/demo/bottomup_demo.py b/demo/bottomup_demo.py
@@ -11,6 +11,7 @@
 import numpy as np
 
 from mmpose.apis import inference_bottomup, init_model
+from mmpose.registry import VISUALIZERS
 from mmpose.structures import split_instances
 
 
@@ -128,20 +129,18 @@ def main():
         device=args.device,
         cfg_options=cfg_options)
 
+    # build visualizer
+    model.cfg.visualizer.radius = args.radius
+    model.cfg.visualizer.line_width = args.thickness
+    visualizer = VISUALIZERS.build(model.cfg.visualizer)
+    visualizer.set_dataset_meta(model.dataset_meta)
+
     if args.input == 'webcam':
         input_type = 'webcam'
     else:
         input_type = mimetypes.guess_type(args.input)[0].split('/')[0]
 
     if input_type == 'image':
-        # init visualizer
-        from mmpose.registry import VISUALIZERS
-
-        model.cfg.visualizer.radius = args.radius
-        model.cfg.visualizer.line_width = args.thickness
-        visualizer = VISUALIZERS.build(model.cfg.visualizer)
-        visualizer.set_dataset_meta(model.dataset_meta)
-
         # inference
         pred_instances = process_one_image(
             args, args.input, model, visualizer, show_interval=0)
@@ -154,22 +153,6 @@ def main():
             mmcv.imwrite(mmcv.rgb2bgr(img_vis), output_file)
 
     elif input_type in ['webcam', 'video']:
-        from mmpose.visualization import FastVisualizer
-
-        visualizer = FastVisualizer(
-            model.dataset_meta,
-            radius=args.radius,
-            line_width=args.thickness,
-            kpt_thr=args.kpt_thr)
-
-        if args.draw_heatmap:
-            # init Localvisualizer
-            from mmpose.registry import VISUALIZERS
-
-            model.cfg.visualizer.radius = args.radius
-            model.cfg.visualizer.line_width = args.thickness
-            local_visualizer = VISUALIZERS.build(model.cfg.visualizer)
-            local_visualizer.set_dataset_meta(model.dataset_meta)
 
         if args.input == 'webcam':
             cap = cv2.VideoCapture(0)
@@ -187,15 +170,8 @@ def main():
             if not success:
                 break
 
-            # bottom-up pose estimation
-            if args.draw_heatmap:
-                pred_instances = process_one_image(args, frame, model,
-                                                   local_visualizer, 0.001)
-            else:
-                pred_instances = process_one_image(args, frame, model)
-                # visualization
-                visualizer.draw_pose(frame, pred_instances)
-                cv2.imshow('MMPose Demo [Press ESC to Exit]', frame)
+            pred_instances = process_one_image(args, frame, model, visualizer,
+                                               0.001)
 
             if args.save_predictions:
                 # save prediction results
@@ -206,10 +182,7 @@ def main():
 
             # output videos
             if output_file:
-                if args.draw_heatmap:
-                    frame_vis = local_visualizer.get_image()
-                else:
-                    frame_vis = frame.copy()[:, :, ::-1]
+                frame_vis = visualizer.get_image()
 
                 if video_writer is None:
                     fourcc = cv2.VideoWriter_fourcc(*'mp4v')

diff --git a/demo/topdown_demo_with_mmdet.py b/demo/topdown_demo_with_mmdet.py
@@ -13,6 +13,7 @@
 from mmpose.apis import inference_topdown
 from mmpose.apis import init_model as init_pose_estimator
 from mmpose.evaluation.functional import nms
+from mmpose.registry import VISUALIZERS
 from mmpose.structures import merge_data_samples, split_instances
 from mmpose.utils import adapt_mmdet_pipeline
 
@@ -186,24 +187,22 @@ def main():
         cfg_options=dict(
             model=dict(test_cfg=dict(output_heatmaps=args.draw_heatmap))))
 
+    # build 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, skeleton_style=args.skeleton_style)
+
     if args.input == 'webcam':
         input_type = 'webcam'
     else:
         input_type = mimetypes.guess_type(args.input)[0].split('/')[0]
 
     if input_type == 'image':
-        # init visualizer
-        from mmpose.registry import VISUALIZERS
-
-        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, skeleton_style=args.skeleton_style)
 
         # inference
         pred_instances = process_one_image(args, args.input, detector,
@@ -218,28 +217,6 @@ def main():
                 mmcv.imwrite(mmcv.rgb2bgr(img_vis), output_file)
 
     elif input_type in ['webcam', 'video']:
-        from mmpose.visualization import FastVisualizer
-
-        visualizer = FastVisualizer(
-            pose_estimator.dataset_meta,
-            radius=args.radius,
-            line_width=args.thickness,
-            kpt_thr=args.kpt_thr)
-
-        if args.draw_heatmap:
-            # init Localvisualizer
-            from mmpose.registry import VISUALIZERS
-
-            pose_estimator.cfg.visualizer.radius = args.radius
-            pose_estimator.cfg.visualizer.alpha = args.alpha
-            pose_estimator.cfg.visualizer.line_width = args.thickness
-            local_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
-            local_visualizer.set_dataset_meta(
-                pose_estimator.dataset_meta,
-                skeleton_style=args.skeleton_style)
 
         if args.input == 'webcam':
             cap = cv2.VideoCapture(0)
@@ -258,16 +235,9 @@ def main():
                 break
 
             # topdown pose estimation
-            if args.draw_heatmap:
-                pred_instances = process_one_image(args, frame, detector,
-                                                   pose_estimator,
-                                                   local_visualizer, 0.001)
-            else:
-                pred_instances = process_one_image(args, frame, detector,
-                                                   pose_estimator)
-                # visualization
-                visualizer.draw_pose(frame, pred_instances)
-                cv2.imshow('MMPose Demo [Press ESC to Exit]', frame)
+            pred_instances = process_one_image(args, frame, detector,
+                                               pose_estimator, visualizer,
+                                               0.001)
 
             if args.save_predictions:
                 # save prediction results
@@ -278,10 +248,7 @@ def main():
 
             # output videos
             if output_file:
-                if args.draw_heatmap:
-                    frame_vis = local_visualizer.get_image()
-                else:
-                    frame_vis = frame.copy()[:, :, ::-1]
+                frame_vis = visualizer.get_image()
 
                 if video_writer is None:
                     fourcc = cv2.VideoWriter_fourcc(*'mp4v')

diff --git a/mmpose/apis/inferencers/base_mmpose_inferencer.py b/mmpose/apis/inferencers/base_mmpose_inferencer.py
@@ -159,6 +159,9 @@ def _get_webcam_inputs(self, inputs: str) -> Generator:
         Raises:
             ValueError: If the inputs string is not in the expected format.
         """
+        assert getattr(self.visualizer, 'backend', None) == 'opencv', \
+            'Visualizer must utilize the OpenCV backend in order to ' \
+            'support webcam inputs.'
 
         # Ensure the inputs string is in the expected format.
         inputs = inputs.lower()
@@ -187,12 +190,9 @@ def _get_webcam_inputs(self, inputs: str) -> Generator:
         self.video_info = dict(
             fps=10, name='webcam.mp4', writer=None, predictions=[])
 
-        # Set up webcam reader generator function.
-        self._window_closing = False
-
         def _webcam_reader() -> Generator:
             while True:
-                if self._window_closing:
+                if cv2.waitKey(5) & 0xFF == 27:
                     vcap.release()
                     break
 
@@ -322,16 +322,6 @@ def visualize(self,
                 kpt_thr=kpt_thr)
             results.append(visualization)
 
-            if show and not hasattr(self, '_window_close_cid'):
-                if window_close_event_handler is None:
-                    window_close_event_handler = \
-                        self._visualization_window_on_close
-                self._window_close_cid = \
-                    self.visualizer.manager.canvas.mpl_connect(
-                        'close_event',
-                        window_close_event_handler
-                    )
-
             if vis_out_dir:
                 out_img = mmcv.rgb2bgr(visualization)
 

diff --git a/mmpose/visualization/__init__.py b/mmpose/visualization/__init__.py
@@ -1,5 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from .fast_visualizer import FastVisualizer
 from .local_visualizer import PoseLocalVisualizer
 
-__all__ = ['PoseLocalVisualizer', 'FastVisualizer']
+__all__ = ['PoseLocalVisualizer']
diff --git a/mmpose/visualization/fast_visualizer.py b/mmpose/visualization/fast_visualizer.py