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
@@ -333,15 +334,34 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         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)
+        
+        # Create a custom logger class to use tqdm for progress display
+        class TqdmLogger:
+            def __init__(self, total):
+                self.pbar = tqdm(total=total)
+        
+            def __call__(self, message):
+                # Extract progress percentage from the message
+                progress_match = re.search(r"frame\s+(\d+)", message)
+                if progress_match:
+                    current_frame = int(progress_match.group(1))
+                    self.pbar.n = current_frame
+                    self.pbar.refresh()
+        
+            def close(self):
+                self.pbar.close()
+        
+        # Initialize the custom logger with total frames
+        logger = TqdmLogger(total_frames)
+
         # Write the result to a file
         final_vid_output_path = "output_video.mp4"
-        # 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)
+        
+        # Write the result to a file using the custom logger for progress
+        clip.write_videofile(final_vid_output_path, codec='libx264', logger=logger)
+        
+        # Close the progress bar after the process is finished
+        logger.close()
         return gr.update(value=None), gr.update(value=final_vid_output_path)
 
 def update_ui(vis_frame_type):