Update app.py

app.py

@@ -1,39 +1,46 @@
 import cv2
 import numpy as np
 import gradio as gr
+from tqdm import tqdm
+
 
 def remove_green_screen(video_file, bg_image):
-    cap = cv2.VideoCapture(video_file)
-    bg = cv2.imread(bg_image)
-    img_array = []
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
+
+    def process_frame(frame):
+        # remove green screen
         hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-        lower_green = np.array([30, 60, 60])
-        upper_green = np.array([90, 255, 255])
-        mask = cv2.inRange(hsv, lower_green, upper_green)
-        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)
+        mask = cv2.inRange(hsv, (36, 25, 25), (86, 255, 255))
+        mask = cv2.merge((mask, mask, mask))
+        masked = cv2.bitwise_and(frame, mask)
+
+        # add background
+        bg = cv2.imread(bg_image)
+        bg = cv2.resize(bg, (frame.shape[1], frame.shape[0]))
+        masked[np.where((masked == [0, 0, 0]).all(axis=2))] = bg[np.where((masked == [0, 0, 0]).all(axis=2))]
+
+        return masked
+
+    cap = cv2.VideoCapture(video_file)
+    frames = []
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    for i in tqdm(range(frame_count)):
+        success, frame = cap.read()
+        if not success:
+            continue
+        frames.append(process_frame(frame))
+
     cap.release()
-    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")
-
-gr.Interface(fn=remove_green_screen, inputs=[input_video, input_bg], outputs=output_video, title="Green Screen Remover").launch()
+
+    return np.array(frames)
+
+
+iface = gr.Interface(
+    fn=remove_green_screen,
+    inputs=[
+        gr.inputs.Video(type="mp4", label="Input video"),
+        gr.inputs.Image(label="Background image")
+    ],
+    outputs=gr.outputs.Video(label="Processed video")
+)
+
+iface.launch()