From f93850598e5df2ca5f797c4566d54098a445d59c Mon Sep 17 00:00:00 2001
From: Yining Li <liyining0712@gmail.com>
Date: Thu, 30 Jun 2022 16:20:10 +0800
Subject: [PATCH] add progress bar in bottom-up tracking demo (#1454)

---
 demo/bottom_up_pose_tracking_demo.py          | 30 ++++++++-----------
 .../top_down_pose_tracking_demo_with_mmdet.py |  2 +-
 2 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/demo/bottom_up_pose_tracking_demo.py b/demo/bottom_up_pose_tracking_demo.py
index 6cda0b3a32..10bb3a2d80 100644
--- a/demo/bottom_up_pose_tracking_demo.py
+++ b/demo/bottom_up_pose_tracking_demo.py
@@ -4,6 +4,7 @@
 from argparse import ArgumentParser
 
 import cv2
+import mmcv
 
 from mmpose.apis import (get_track_id, inference_bottom_up_pose_model,
                          init_pose_model, vis_pose_tracking_result)
@@ -86,10 +87,8 @@ def main():
     else:
         dataset_info = DatasetInfo(dataset_info)
 
-    cap = cv2.VideoCapture(args.video_path)
-    fps = None
-
-    assert cap.isOpened(), f'Faild to load video file {args.video_path}'
+    video = mmcv.VideoReader(args.video_path)
+    assert video.opened, f'Faild to load video file {args.video_path}'
 
     if args.out_video_root == '':
         save_out_video = False
@@ -98,9 +97,8 @@ def main():
         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)))
+        fps = video.fps
+        size = (video.width, video.height)
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         videoWriter = cv2.VideoWriter(
             os.path.join(args.out_video_root,
@@ -128,15 +126,12 @@ def main():
     output_layer_names = None
     next_id = 0
     pose_results = []
-    while (cap.isOpened()):
-        flag, img = cap.read()
-        if not flag:
-            break
+    for cur_frame in mmcv.track_iter_progress(video):
         pose_results_last = pose_results
 
-        pose_results, returned_outputs = inference_bottom_up_pose_model(
+        pose_results, _ = inference_bottom_up_pose_model(
             pose_model,
-            img,
+            cur_frame,
             dataset=dataset,
             dataset_info=dataset_info,
             pose_nms_thr=args.pose_nms_thr,
@@ -158,9 +153,9 @@ def main():
             pose_results = smoother.smooth(pose_results)
 
         # show the results
-        vis_img = vis_pose_tracking_result(
+        vis_frame = vis_pose_tracking_result(
             pose_model,
-            img,
+            cur_frame,
             pose_results,
             radius=args.radius,
             thickness=args.thickness,
@@ -170,15 +165,14 @@ def main():
             show=False)
 
         if args.show:
-            cv2.imshow('Image', vis_img)
+            cv2.imshow('Image', vis_frame)
 
         if save_out_video:
-            videoWriter.write(vis_img)
+            videoWriter.write(vis_frame)
 
         if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
             break
 
-    cap.release()
     if save_out_video:
         videoWriter.release()
     if args.show:
diff --git a/demo/top_down_pose_tracking_demo_with_mmdet.py b/demo/top_down_pose_tracking_demo_with_mmdet.py
index b5744157a1..65a1a458eb 100644
--- a/demo/top_down_pose_tracking_demo_with_mmdet.py
+++ b/demo/top_down_pose_tracking_demo_with_mmdet.py
@@ -187,7 +187,7 @@ def main():
                                           args.online)
 
         # test a single image, with a list of bboxes.
-        pose_results, returned_outputs = inference_top_down_pose_model(
+        pose_results, _ = inference_top_down_pose_model(
             pose_model,
             frames if args.use_multi_frames else cur_frame,
             person_results,