Update app.py

app.py

@@ -2,6 +2,7 @@ import spaces
 import gradio as gr
 import cv2
 import numpy as np
+import os
 
 from PIL import Image
 from transparent_background import Remover
@@ -14,6 +15,12 @@ def doo(video):
     fps = cap.get(cv2.CAP_PROP_FPS)
     
     writer = None
+    frame_count = 0
+    max_frames = 400  # Maximum number of frames to process at a time
+    segment_count = 0  # Count of video segments processed
+    
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # Total number of frames in the video
+    processed_frames = 0  # Number of frames processed so far
     
     while cap.isOpened():
         ret, frame = cap.read()
@@ -25,14 +32,45 @@ def doo(video):
         img = Image.fromarray(frame).convert('RGB')
     
         if writer is None:
-            writer = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)
+            writer = cv2.VideoWriter(f'output_{segment_count}.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)
     
         out = remover.process(img, type='green')
         writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
+        
+        frame_count += 1
+        processed_frames += 1
+        gr.Interface.progress(processed_frames / total_frames)  # Update the progress bar
+        
+        if frame_count >= max_frames:
+            cap.release()
+            writer.release()
+            
+            # Process the next segment
+            cap = cv2.VideoCapture(video)
+            for _ in range(frame_count):
+                cap.read()  # Skip frames that have already been processed
+            writer = None
+            frame_count = 0
+            segment_count += 1
     
     cap.release()
-    writer.release()
-    return 'output.mp4'
+    if writer is not None:
+        writer.release()
+
+    # Concatenate all the video segments
+    final_output = cv2.VideoWriter('final_output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)
+    for i in range(segment_count):
+        cap = cv2.VideoCapture(f'output_{i}.mp4')
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if ret is False:
+                break
+            final_output.write(frame)
+        cap.release()
+        os.remove(f'output_{i}.mp4')  # Delete the segment file
+    final_output.release()
+
+    return 'final_output.mp4'
 
 iface = gr.Interface(fn=doo, inputs="video", outputs="video")
-iface.launch()
+iface.launch()