Update app.py

app.py

@@ -28,6 +28,7 @@ from PIL import Image, ImageFilter
 from sam2.build_sam import build_sam2_video_predictor
 
 from moviepy.editor import ImageSequenceClip
+from tqdm import tqdm
 
 def get_video_fps(video_path):
     # Open the video file
@@ -330,10 +331,17 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         # Create a video clip from the image sequence
         original_fps = get_video_fps(video_in)
         fps = original_fps  # Frames per second
+        total_frames = len(jpeg_images)
         clip = ImageSequenceClip(jpeg_images, fps=fps)
+        # Define a callback function to update the tqdm progress bar
+        def tqdm_callback(progress):
+            pbar.update(progress - pbar.n)
         # Write the result to a file
         final_vid_output_path = "output_video.mp4"
-        clip.write_videofile(final_vid_output_path, codec='libx264')
+        # Create a tqdm progress bar with the total set to the number of frames
+        with tqdm(total=total_frames) as pbar:
+            # Write the result to a file with the progress callback
+            clip.write_videofile(final_vid_output_path, codec='libx264', progress_bar=False, verbose=False, logger=None, progress_callback=tqdm_callback)
         return gr.update(value=None), gr.update(value=final_vid_output_path)
 
 def update_ui(vis_frame_type):
@@ -342,6 +350,7 @@ def update_ui(vis_frame_type):
     elif vis_frame_type == "render":
         return gr.update(visible=False), gr.update(visible=True)
 
+
 with gr.Blocks() as demo:
     first_frame_path = gr.State()
     tracking_points = gr.State([])