Update app.py

app.py

@@ -14,9 +14,11 @@ def doo(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))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
-    writer = None
-    batch_size = 30
+    writer = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+    
+    batch_size = 30  # Process 30 frames at a time
     
     while cap.isOpened():
         frames = []
@@ -26,21 +28,17 @@ def doo(video):
                 break
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 
             img = Image.fromarray(frame).convert('RGB')
-            frames.append(img)
+            out = remover.process(img, type='green')
+            frames.append(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
         
         if not frames:
             break
             
-        if writer is None:
-            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))
+        for frame in frames:
+            writer.write(frame)
     
     cap.release()
-    if writer:
-        writer.release()
+    writer.release()
     return 'output.mp4'
 
 iface = gr.Interface(fn=doo, inputs="video", outputs="video")