import spaces
import gradio as gr
import cv2
import numpy as np
import time
import random
import torch
from PIL import Image
from transparent_background import Remover


@spaces.GPU  # 修改装饰器，去掉括号
def doo(video, mode):
    if(mode == 'Fast'):
        remover = Remover(mode='fast')
    else:
        remover = Remover()
    cap = cv2.VideoCapture(video)
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    writer = None
    tmpname = random.randint(111111111, 999999999)
    processed_frames = 0
    start_time = time.time()
    
    while cap.isOpened():
        ret, frame = cap.read()
    
        if ret is False:
            break

        if time.time() - start_time >= 55:
            print("GPU Timeout is coming")
            cap.release()
            writer.release()
            return str(tmpname) + '.mp4'
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 
        img = Image.fromarray(frame).convert('RGB')
    
        if writer is None:
            writer = cv2.VideoWriter(str(tmpname) + '.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)

        processed_frames += 1
        print(f"Processing: {processed_frames}")
        out = remover.process(img, type='green')
        writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
    
    cap.release()
    writer.release()
    return str(tmpname) + '.mp4'

title = "🎞️Video Background Removal tool🎥"

description = r"""Please note that if your video file is long (or having high amount of frames) the result may be shorter than input video because of the GPU timeout.<br>
In this case consider trying Fast mode."""

examples = [['./input.mp4'],]

iface = gr.Interface(
    fn=doo, 
    inputs=["video", gr.components.Radio(['Normal', 'Fast'], label='Select mode', value='Normal', info= 'Normal is more accurate, but takes longer. | Fast has lower accuracy so the process will be faster.')], 
    outputs="video", examples=examples, title=title, description=description
)
iface.launch()