add progress bar and draw human keypoints option

app.py

@@ -3,7 +3,6 @@ from tools.inferencer import PoseInferencerV2
 from tools.dtw import DTWForKeypoints
 from tools.visualizer import FastVisualizer
 from tools.utils import convert_video_to_playable_mp4
-from argparse import ArgumentParser
 from pathlib import Path
 from tqdm import tqdm
 import mmengine
@@ -12,13 +11,6 @@ import mmcv
 import cv2
 import gradio as gr
 
-def parse_args():
-    parser = ArgumentParser()
-    parser.add_argument('--config', type=str, default='configs/mark2.py')
-    parser.add_argument('--video1', type=str, default='assets/tennis1.mp4')
-    parser.add_argument('--video2', type=str, default='assets/tennis2.mp4')
-    return parser.parse_args()
-
 def concat(img1, img2, height=1080):
     w1, h1, _ = img1.shape
     w2, h2, _ = img2.shape
@@ -35,15 +27,19 @@ def concat(img1, img2, height=1080):
     image = cv2.hconcat([img1, img2])
     return image
 
-def draw(vis: FastVisualizer, img, keypoint, box, oks, oks_unnorm, draw_score_bar=True):
+def draw(vis: FastVisualizer, img, keypoint, box, oks, oks_unnorm, 
+         draw_human_keypoints=True,
+         draw_score_bar=True):
     vis.set_image(img)
     vis.draw_non_transparent_area(box)
     if draw_score_bar:
         vis.draw_score_bar(oks)
-    vis.draw_human_keypoints(keypoint, oks_unnorm)
+    if draw_human_keypoints:
+        vis.draw_human_keypoints(keypoint, oks_unnorm)
     return vis.get_image()
 
-def main(video1, video2):
+def main(video1, video2, draw_human_keypoints,
+         progress=gr.Progress(track_tqdm=True)):
     # build PoseInferencerV2
     config = 'configs/mark2.py'
     cfg = mmengine.Config.fromfile(config)
@@ -68,14 +64,14 @@ def main(video1, video2):
 
     vis = FastVisualizer()
     
-    for i, j in tqdm(dtw_path): 
+    for i, j in tqdm(dtw_path, desc='Visualizing'): 
         frame1 = v1[i]
         frame2 = v2[j]
 
         frame1_ = draw(vis, frame1.copy(), keypoints1[i], boxes1[i],
-                       oks[i, j], oks_unnorm[i, j])
+                       oks[i, j], oks_unnorm[i, j], draw_human_keypoints)
         frame2_ = draw(vis, frame2.copy(), keypoints2[j], boxes2[j],
-                       oks[i, j], oks_unnorm[i, j], draw_score_bar=False)
+                       oks[i, j], oks_unnorm[i, j], draw_human_keypoints, draw_score_bar=False)
         # concate two frames
         frame = concat(frame1_, frame2_)
         # draw logo
@@ -100,10 +96,12 @@ if __name__ == '__main__':
 
     inputs = [
         gr.Video(label="Input video 1"),
-        gr.Video(label="Input video 2")
+        gr.Video(label="Input video 2"),
+        "checkbox"
     ]
 
     output = gr.Video(label="Output video")
 
-    demo = gr.Interface(fn=main, inputs=inputs, outputs=output).queue()
+    demo = gr.Interface(fn=main, inputs=inputs, outputs=output,
+                        allow_flagging='never').queue()
     demo.launch()

tools/inferencer.py

@@ -29,6 +29,8 @@ class PoseInferencer:
         self.pose_model = init_model(self.pose_model_cfg,
                                      self.pose_model_ckpt,
                                      device=device)
+        # use this count to tell the progress
+        self.video_count = 0
 
     def process_one_image(self, img):
         init_default_scope('mmdet')
@@ -101,6 +103,8 @@ class PoseInferencerV2:
         self.pose_model = init_model(self.pose_model_cfg,
                                      self.pose_model_ckpt,
                                      device)
+        # use this count to tell the progress
+        self.video_count = 0
 
     def process_one_image(self, img):
         init_default_scope('mmdet')
@@ -145,10 +149,12 @@ class PoseInferencerV2:
         video_reader = mmcv.VideoReader(video_path)
         all_pose, all_det = [], []
 
-        for frame in tqdm(video_reader):
+        count = self.video_count + 1
+        for frame in tqdm(video_reader, desc=f'Inference video {count}'):
             # inference with detector
             det, pose = self.process_one_image(frame)
             all_pose.append(pose)
             all_det.append(det)
+        self.video_count += 1
 
         return all_det, all_pose

tools/visualizer.py

@@ -157,8 +157,8 @@ class FastVisualizer:
         else: lvl_names = self.score_level_names(scores)
 
         for idx, (point, lvl_name) in enumerate(zip(keypoints, lvl_names)):
-            if idx in set((1, 2, 3, 4)): 
-                continue # do not draw eyes and years
+            if idx in set((0, 1, 2, 3, 4)): 
+                continue # do not draw head
             rectangle_xyhw = np.array((point[0], point[1], cube_size, cube_size))
             rectangle_xyxy = self.xyhw_to_xyxy(rectangle_xyhw)
             self.draw_rectangle(rectangle_xyxy,