Update app.py

app.py

@@ -4,24 +4,36 @@ import gradio as gr
 
 def remove_green_screen(video_file, bg_image):
     cap = cv2.VideoCapture(video_file)
-    background = cv2.imread(bg_image)
+    bg = cv2.imread(bg_image)
+    img_array = []
     while True:
         ret, frame = cap.read()
         if not ret:
             break
         hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-        lower_green = np.array([35, 0, 0])
+        lower_green = np.array([30, 60, 60])
         upper_green = np.array([90, 255, 255])
         mask = cv2.inRange(hsv, lower_green, upper_green)
-        mask_inv = cv2.bitwise_not(mask)
-        fg = cv2.bitwise_and(frame, frame, mask=mask_inv)
-        bg = cv2.bitwise_and(background, background, mask=mask)
-        final = cv2.add(fg, bg)
-        cv2.imshow("output", final)
-        if cv2.waitKey(1) & 0xFF == ord("q"):
-            break
+        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((3,3),np.uint8))
+        mask = cv2.dilate(mask, np.ones((3,3),np.uint8), iterations=5)
+        mask = cv2.blur(mask, (30,30))
+        mask = np.stack([mask, mask, mask], axis=2)
+        frame = np.where(mask == np.array([255, 255, 255]), bg, frame)
+        img_array.append(frame)
     cap.release()
-    cv2.destroyAllWindows()
+    output_file = "output.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    fps = 30.0
+    height, width, layers = img_array[0].shape
+    size = (width,height)
+    out = cv2.VideoWriter(output_file, fourcc, fps, size)
+    for i in range(len(img_array)):
+        out.write(img_array[i])
+    out.release()
+    return output_file
+
+input_video = gr.inputs.Video(type="mp4")
+input_bg = gr.inputs.Image(label="Background Image")
+output_video = gr.outputs.Video(type="mp4")
 
-iface = gr.Interface(fn=remove_green_screen, inputs=["file", "image"], outputs=None, title="Green Screen Remover")
-iface.launch()
+gr.Interface(fn=remove_green_screen, inputs=[input_video, input_bg], outputs=output_video, title="Green Screen Remover").launch()