import gradio as gr
import cv2
import numpy as np
from PIL import Image
from transparent_background import Remover

remover = Remover(mode='fast')  # Custom setting

def doo(video):
    cap = cv2.VideoCapture(video)  # Video reader for input
    fps = cap.get(cv2.CAP_PROP_FPS)

    processed_frames = []  # List to store processed frames

    while cap.isOpened():
        ret, frame = cap.read()  # Read video

        if ret is False:
            break

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame).convert('RGB')

        # Process the frame using the transparent-background model
        out = remover.process(img, type='map')  # Same as image, except for 'rgba'

        # Convert the processed frame back to a NumPy array
        processed_frame = np.array(out)

        # Ensure the processed frame has shape (height, width, 3)
        if processed_frame.shape[2] != 3:
            raise ValueError("Processed frame does not have 3 channels (RGB)")

        # Append the processed frame to the list
        processed_frames.append(processed_frame)

    cap.release()

    # Return the processed frames (as a list of NumPy arrays) and the original video's frame rate
    return processed_frames, fps

iface = gr.Interface(fn=doo, inputs="video", outputs="video")
iface.launch()