add: progress bar

app.py

@@ -1,12 +1,21 @@
 import gradio as gr
 from videopose_PSTMO import gr_video2mc
+import os
 
 
-def Video2MC(video):
-    output_path = gr_video2mc(video)
+def Video2MC(video, progress = gr.Progress(track_tqdm=True)):
+
+    progress(1.0, desc="Step 0: Starting")
+    output_path = gr_video2mc(video, progress)
+
     return output_path, output_path
 
+    
 iface = gr.Interface(fn=Video2MC, 
                      inputs=gr.Video(), 
-                     outputs=["file", "text"])
-iface.launch()
+                     outputs=["file", "text"],
+                     examples=[os.path.join(os.path.dirname(__file__), 
+                                     "input_videos/kun_test_5sec.mp4")], 
+                     )
+
+iface.queue(concurrency_count=10).launch()

common/inference_3d.py

@@ -48,7 +48,7 @@ def input_augmentation(input_2D, input_2D_flip, model_trans, joints_left, joints
 
     return input_2D, output_3D, output_3D_VTE
 
-def step(opt, dataLoader, model, optimizer=None, epoch=None):
+def step(opt, dataLoader, model, progress, optimizer=None, epoch=None):
     model_trans = model['trans']
 
     model_trans.eval()
@@ -57,6 +57,7 @@ def step(opt, dataLoader, model, optimizer=None, epoch=None):
     joints_right = [1, 2, 3, 14, 15, 16]
     epoch_cnt=0
     out = []
+    # probar = progress.tqdm(total=1, desc="Step 2: Prediction (Please wait)")
     for _, batch, batch_2d, batch_2d_flip in dataLoader.next_epoch():
         #[gt_3D, input_2D] = get_varialbe([batch, batch_2d])
         #input_2D = Variable(batch_2d).contiguous().type(torch.cuda.FloatTensor)
@@ -102,6 +103,6 @@ def step(opt, dataLoader, model, optimizer=None, epoch=None):
         epoch_cnt +=1
     return out.numpy()
 
-def val(opt, val_loader, model):
+def val(opt, val_loader, model, progress):
     with torch.no_grad():
-        return step(opt, val_loader, model)
+        return step(opt, val_loader, model, progress)

common/visualization.py

@@ -50,7 +50,7 @@ def downsample_tensor(X, factor):
     return np.mean(X[:length].reshape(-1, factor, *X.shape[1:]), axis=1)
 
 
-def render_animation(keypoints, poses, skeleton, fps, bitrate, azim, output, viewport,
+def render_animation(keypoints, poses, skeleton, fps, bitrate, azim, output, progress, viewport, 
                      limit=-1, downsample=1, size=6, input_video_path=None, input_video_skip=0):
     """
     TODO
@@ -126,6 +126,7 @@ def render_animation(keypoints, poses, skeleton, fps, bitrate, azim, output, vie
 
     parents = skeleton.parents()
     pbar = tqdm(total=limit)
+    # probar = progress.tqdm(total=limit, desc="Step 3: 3D Rendering")
 
     def update_video(i):
         nonlocal initialized, image, lines, points
@@ -177,6 +178,7 @@ def render_animation(keypoints, poses, skeleton, fps, bitrate, azim, output, vie
             points.set_offsets(keypoints[i])
 
         pbar.update()
+        # probar.update()
 
     fig.tight_layout()
 

joints_detectors/Alphapose/gene_npz.py

@@ -32,8 +32,8 @@ def image_interface(model, image):
     pass
 
 
-def generate_kpts(video_file):
-    final_result, video_name = handle_video(video_file)
+def generate_kpts(video_file, progress):
+    final_result, video_name = handle_video(video_file, progress)
 
     # ============ Changing ++++++++++
 
@@ -67,7 +67,7 @@ def generate_kpts(video_file):
     return kpts
 
 
-def handle_video(video_file):
+def handle_video(video_file, progress):
     # =========== common ===============
     args.video = video_file
     base_name = os.path.basename(args.video)
@@ -136,7 +136,7 @@ def handle_video(video_file):
     print('Start pose estimation...')
     im_names_desc = tqdm(range(data_loader.length()))
     batchSize = args.posebatch
-    for i in im_names_desc:
+    for i in progress.tqdm(range(data_loader.length()), desc="Step 1: 2D Detecting"):
 
         start_time = getTime()
         with torch.no_grad():

videopose_PSTMO.py

@@ -60,7 +60,7 @@ class Skeleton:
         return [1, 2, 3, 14, 15, 16]
 
 
-def main(args):
+def main(args, progress):
     detector_2d = get_detector_2d(args.detector_2d)
 
     assert detector_2d, 'detector_2d should be in ({alpha, hr, open}_pose)'
@@ -69,7 +69,7 @@ def main(args):
     #args.input_npz = './outputs/alpha_pose_skiing_cut/skiing_cut.npz'
     if not args.input_npz:
         video_name = args.viz_video
-        keypoints = detector_2d(video_name)
+        keypoints = detector_2d(video_name, progress)
     else:
         npz = np.load(args.input_npz)
         keypoints = npz['kpts']  # (N, 17, 2)
@@ -133,7 +133,7 @@ def main(args):
                              pad=pad, causal_shift=causal_shift, augment=args.test_time_augmentation, shuffle=False,
                              kps_left=kps_left, kps_right=kps_right, joints_left=joints_left, joints_right=joints_right)
 
-    prediction = val(args, gen, model)
+    prediction = val(args, gen, model, progress)
 
     # save 3D joint points
     # np.save(f'outputs/test_3d_{args.video_name}_output.npy', prediction, allow_pickle=True)
@@ -160,7 +160,7 @@ def main(args):
 
     from common.visualization import render_animation
     render_animation(input_keypoints, anim_output,
-                     Skeleton(), 25, args.viz_bitrate, np.array(70., dtype=np.float32), args.viz_output,
+                     Skeleton(), 25, args.viz_bitrate, np.array(70., dtype=np.float32), args.viz_output, progress,
                      limit=args.viz_limit, downsample=args.viz_downsample, size=args.viz_size,
                      input_video_path=args.viz_video, viewport=(1000, 1002),
                      input_video_skip=args.viz_skip)
@@ -171,7 +171,7 @@ def main(args):
     return output_dir_dict
 
 
-def inference_video(video_path, detector_2d):
+def inference_video(video_path, detector_2d, progress):
     """
     Do image -> 2d points -> 3d points to video.
     :param detector_2d: used 2d joints detector. Can be {alpha_pose, hr_pose}
@@ -189,14 +189,14 @@ def inference_video(video_path, detector_2d):
     args.evaluate = 'pretrained_h36m_detectron_coco.bin'
 
     with Timer(video_path):
-        output_dir_dict = main(args)
+        output_dir_dict = main(args, progress)
 
     output_dir_dict["output_videos"] = args.viz_output
     output_dir_dict["video_name"] = args.video_name
     return output_dir_dict
 
 
-def gr_video2mc(video_path):
+def gr_video2mc(video_path, progress):
 
     if not os.path.exists('output_3Dpose_npy'):
         os.makedirs('output_3Dpose_npy')
@@ -208,7 +208,7 @@ def gr_video2mc(video_path):
         os.makedirs('output_videos')
 
     FPS_mine_imator = 30
-    output_dir_dict = inference_video(video_path, 'alpha_pose')
+    output_dir_dict = inference_video(video_path, 'alpha_pose', progress)
     Hk.hpe2keyframes(output_dir_dict['npy'], FPS_mine_imator, f"output_miframes/{output_dir_dict['video_name']}.miframes")
     return os.path.abspath(f"output_miframes/{output_dir_dict['video_name']}.miframes")