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

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