import cv2
import numpy as np
import gradio as gr

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