|
import spaces |
|
import gradio as gr |
|
import cv2 |
|
import numpy as np |
|
|
|
from PIL import Image |
|
from transparent_background import Remover |
|
|
|
remover = Remover() |
|
processed_frames = 0 |
|
|
|
@spaces.GPU |
|
def doo(video): |
|
cap = cv2.VideoCapture(video) |
|
fps = cap.get(cv2.CAP_PROP_FPS) |
|
|
|
writer = None |
|
|
|
while cap.isOpened(): |
|
ret, frame = cap.read() |
|
|
|
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) |
|
|
|
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 'output.mp4' |
|
|
|
iface = gr.Interface(fn=doo, inputs="video", outputs="video") |
|
iface.launch() |
