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
@@ -15,7 +16,8 @@ def doo(video):
     
     writer = None
     frame_count = 0
-    max_frames = 50  # Maximum number of frames to process at a time
+    max_frames = 400  # Maximum number of frames to process at a time
+    segment_count = 0  # Count of video segments processed
     
     while cap.isOpened():
         ret, frame = cap.read()
@@ -27,7 +29,7 @@ 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))
@@ -43,11 +45,26 @@ def doo(video):
                 cap.read()  # Skip frames that have already been processed
             writer = None
             frame_count = 0
+            segment_count += 1
     
     cap.release()
     if writer is not None:
         writer.release()
-    return 'output.mp4'
+
+    # 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()