Update app.py

app.py

@@ -2,66 +2,72 @@ import gradio as gr
 import cv2
 import numpy as np
 
-def remove_green_screen(video):
-    def process_frame(frame):
-        # Convert to HSV color space
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-
-        # Define the range of green color
-        lower_green = np.array([25, 52, 72])
-        upper_green = np.array([102, 255, 255])
-
-        # Threshold the HSV image to get only green colors
-        mask = cv2.inRange(hsv, lower_green, upper_green)
-
-        # Bitwise-AND mask and original image
-        res = cv2.bitwise_and(frame, frame, mask=mask)
-
-        return res
-
-    # Load background image
-    bg_image = cv2.imread("background.jpg")
-
-    # Open video file
-    cap = cv2.VideoCapture(video)
-
-    # Get video codec and fps
-    fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
+def process_frame(frame, bg_image):
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
+    lower_green = np.array([45, 100, 50])
+    upper_green = np.array([75, 255, 255])
+    mask = cv2.inRange(hsv, lower_green, upper_green)
+    mask_inv = cv2.bitwise_not(mask)
+    bg = cv2.imread(bg_image)
+    bg = cv2.resize(bg, (frame.shape[1], frame.shape[0]))
+    fg = cv2.bitwise_and(frame, frame, mask=mask_inv)
+    bg = cv2.bitwise_and(bg, bg, mask=mask)
+    result = cv2.add(bg, fg)
+    return result
+
+def remove_green_screen(input_video, bg_image):
+    cap = cv2.VideoCapture(input_video.name)
+    codec = cv2.VideoWriter_fourcc(*"mp4v")
     fps = int(cap.get(cv2.CAP_PROP_FPS))
-
-    # Get video dimensions
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    # Create output video writer
-    out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))
-
-    frames = []
-
-    # Loop through video frames
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    out = cv2.VideoWriter("output.mp4", codec, fps, (frame_width, frame_height))
     while True:
         ret, frame = cap.read()
         if not ret:
             break
-
-        # Process frame
-        bg = cv2.imread(bg_image)
-        frames.append(process_frame(frame))
-
-    # Release video file and output video writer
+        result = process_frame(frame, bg_image)
+        out.write(result)
     cap.release()
     out.release()
 
-    # Write output video
-    out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))
-    for frame in frames:
-        out.write(frame)
-    out.release()
-
+def predict(input_video, bg_image):
+    remove_green_screen(input_video, bg_image)
     return "output.mp4"
 
-
-input_video = gr.inputs.Video(label="Input Video")
-outputs = gr.outputs.Video(label="Processed Video", type="mp4")
-
-gr.Interface(remove_green_screen, inputs=input_video, outputs=outputs).launch()
+inputs = [
+    gr.inputs.Video(label="Input Video"),
+    gr.inputs.Image(label="Background Image")
+]
+
+outputs = gr.outputs.Video(label="Processed Video", type="auto")
+
+title = "Green Screen Remover"
+description = "Upload a video and an image to use as the background to remove the green screen."
+article = "<p style='text-align: center'><a href='https://github.com/gradio-app/examples/blob/master/green_screen_removal.py'>This code</a> was made into an interactive interface using Gradio. See the full tutorial at <a href='https://blog.gradio.app/green-screen-removal-with-opencv/'>this blog post</a> on the Gradio blog.</p>"
+examples = [
+    [
+        "https://www.youtube.com/watch?v=clD6_yXKo2I",
+        "https://i.imgur.com/lxIhsG6.jpg"
+    ],
+    [
+        "https://www.youtube.com/watch?v=6DfZ6UOZi0A",
+        "https://i.imgur.com/6UaTvfo.jpg"
+    ],
+]
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=inputs,
+    outputs=outputs,
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+    analytics_enabled=False,
+    server_port=8000
+)
+
+if __name__ == '__main__':
+    iface.launch()