import spaces
import gradio as gr
import cv2
import numpy as np

from PIL import Image
from transparent_background import Remover

remover = Remover()

@spaces.GPU
def doo(video):
    cap = cv2.VideoCapture(video)
    fps = 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))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    writer = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
    
    batch_size = 30  # Process 30 frames at a time
    
    while cap.isOpened():
        frames = []
        for _ in range(batch_size):
            ret, frame = cap.read()
            if ret is False:
                break
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 
            img = Image.fromarray(frame).convert('RGB')
            out = remover.process(img, type='green')
            frames.append(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
        
        if not frames:
            break
            
        for frame in frames:
            writer.write(frame)
    
    cap.release()
    writer.release()
    return 'output.mp4'

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