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(image):
    #img = Image.fromarray(image).convert('RGB') # read image
    #out = remover.process(img) # default setting - transparent background
    
    #out.save('output.png') # save result
    #return out

def doo(video):
    cap = cv2.VideoCapture(video) # video reader for input
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    writer = None
    
    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')
    
        if writer is None:
            writer = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size) # video writer for output
    
        out = remover.process(img, type='map') # same as image, except for 'rgba' which is not for video.
        writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
    
    cap.release()
    writer.release()
    return writer

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