Update app.py

app.py

@@ -12,26 +12,35 @@ remover = Remover()
 def doo(video):
     cap = cv2.VideoCapture(video)
     fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     
     writer = None
+    batch_size = 300
     
     while cap.isOpened():
-        ret, frame = cap.read()
-    
-        if ret is False:
+        frames = []
+        for _ in range(batch_size):
+            ret, frame = cap.read()
+            if ret is False:
+                break
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 
+            img = Image.fromarray(frame).convert('RGB')
+            frames.append(img)
+        
+        if not frames:
             break
             
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 
-        img = Image.fromarray(frame).convert('RGB')
-    
         if writer is None:
-            writer = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)
-    
-        out = remover.process(img, type='green')
-        writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
+            writer = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+        
+        for img in frames:
+            out = remover.process(img, type='green')
+            writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
     
     cap.release()
-    writer.release()
+    if writer:
+        writer.release()
     return 'output.mp4'
 
 iface = gr.Interface(fn=doo, inputs="video", outputs="video")