import gradio as gr
import cv2
import numpy as np

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))
    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
        result = process_frame(frame, bg_image)
        out.write(result)
    cap.release()
    out.release()

def predict(input_video, bg_image):
    remove_green_screen(input_video, bg_image)
    return "output.mp4"

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()