Update app.py

app.py

@@ -39,37 +39,67 @@ def fn(vid, bg_type="Color", bg_image=None, color="#00FF00", fps=0):
     # Extract audio from the video
     audio = video.audio
 
-    # Extract frames at the specified FPS
-    frames = video.iter_frames(fps=fps)
+    # Process video in chunks of 1 second
+    chunk_duration = 1  # seconds
+    total_duration = video.duration
+    start_time = 0
 
-    # Process each frame for background removal
     processed_frames = []
-    yield gr.update(visible=True), gr.update(visible=False)
-    for frame in frames:
-        pil_image = Image.fromarray(frame)
-        if bg_type == "Color":
-            processed_image = process(pil_image, color)
-        else:
-            processed_image = process(pil_image, bg_image)
-        processed_frames.append(np.array(processed_image))
-        yield processed_image, None
+    yield gr.update(visible=True), gr.update(visible=False), gr.update(value=0)
+
+    while start_time < total_duration:
+        end_time = min(start_time + chunk_duration, total_duration)
+        chunk = video.subclip(start_time, end_time)
+        chunk_frames = chunk.iter_frames(fps=fps)
+
+        for frame in chunk_frames:
+            pil_image = Image.fromarray(frame)
+            if bg_type == "Color":
+                processed_image = process(pil_image, color)
+            else:
+                processed_image = process(pil_image, bg_image)
+            processed_frames.append(np.array(processed_image))
+            yield processed_image, None, None
+
+        # Save processed frames for the current chunk
+        temp_dir = "temp"
+        os.makedirs(temp_dir, exist_ok=True)
+        for i, frame in enumerate(processed_frames):
+            Image.fromarray(frame).save(os.path.join(temp_dir, f"frame_{start_time}_{i}.png"))
+
+        # Clear processed frames for the current chunk
+        processed_frames = []
+
+        # Update progress bar
+        progress = start_time / total_duration
+        yield None, None, gr.update(value=progress)
+
+        start_time += chunk_duration
+
+    # Load all saved frames
+    all_frames = []
+    for filename in sorted(os.listdir(temp_dir)):
+        if filename.startswith("frame_") and filename.endswith(".png"):
+            frame = np.array(Image.open(os.path.join(temp_dir, filename)))
+            all_frames.append(frame)
 
     # Create a new video from the processed frames
-    processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
+    processed_video = mp.ImageSequenceClip(all_frames, fps=fps)
 
     # Add the original audio back to the processed video
     processed_video = processed_video.set_audio(audio)
 
     # Save the processed video to a temporary file
-    temp_dir = "temp"
-    os.makedirs(temp_dir, exist_ok=True)
-    unique_filename = str(uuid.uuid4()) + ".mp4"
-    temp_filepath = os.path.join(temp_dir, unique_filename)
+    temp_filepath = os.path.join(temp_dir, "processed_video.mp4")
     processed_video.write_videofile(temp_filepath, codec="libx264")
 
-    yield gr.update(visible=False), gr.update(visible=True)
+    # Clean up temporary files
+    for filename in os.listdir(temp_dir):
+        os.remove(os.path.join(temp_dir, filename))
+
+    yield gr.update(visible=False), gr.update(visible=True), gr.update(value=1)
     # Return the path to the temporary file
-    yield processed_image, temp_filepath
+    yield processed_image, temp_filepath, None
 
 
 def process(image, bg):
@@ -120,11 +150,12 @@ with gr.Blocks() as demo:
 
     bg_type.change(update_visibility, inputs=bg_type, outputs=[color_picker, bg_image])
 
+    progress_bar = gr.ProgressBar(label="Processing Video")
 
     examples = gr.Examples(
         [["rickroll-2sec.mp4", "Image", "images.webp"], ["rickroll-2sec.mp4", "Color", None]],
         inputs=[in_video, bg_type, bg_image],
-        outputs=[stream_image, out_video],
+        outputs=[stream_image, out_video, progress_bar],
         fn=fn,
         cache_examples=True,
         cache_mode="eager",
@@ -134,7 +165,7 @@ with gr.Blocks() as demo:
     submit_button.click(
         fn,
         inputs=[in_video, bg_type, bg_image, color_picker, fps_slider],
-        outputs=[stream_image, out_video],
+        outputs=[stream_image, out_video, progress_bar],
     )
 
 if __name__ == "__main__":