Update app.py

app.py

@@ -29,6 +29,7 @@ from sam2.build_sam import build_sam2_video_predictor
 
 from moviepy.editor import ImageSequenceClip
 from tqdm import tqdm
+import re
 
 def get_video_fps(video_path):
     # Open the video file
@@ -335,13 +336,28 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         clip = ImageSequenceClip(jpeg_images, fps=fps)
         # Write the result to a file
         final_vid_output_path = "output_video.mp4"
-        # Use tqdm progress bar while writing the video file
-        with tqdm(total=total_frames, desc="Rendering video") as pbar:
-            def update_progress(current_frame):
-                pbar.update(current_frame - pbar.n)
-            
-            clip.write_videofile("output_video.mp4", codec='libx264', progress_bar=False, verbose=False, logger=None, progress_callback=update_progress)
-
+        
+        # Initialize the tqdm progress bar with the total number of frames
+        pbar = tqdm(total=total_frames, desc="Rendering video")
+        
+        # Define a function to update the progress bar
+        def update_progress(txt):
+            # Extract progress from the log message
+            progress_match = re.search(r"frame\s+(\d+)", txt)
+            if progress_match:
+                current_frame = int(progress_match.group(1))
+                pbar.n = current_frame
+                pbar.refresh()
+        
+        # Write the result to a file
+        clip.write_videofile(
+            final_vid_output_path,
+            codec='libx264',
+            logger=lambda txt: update_progress(txt)
+        )
+        
+        # Close the progress bar after the process is finished
+        pbar.close()