app.py

import spaces
import gradio as gr
import cv2
import numpy as np
import gc  # For manual garbage collection

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)
    
    processed_frames = 0
    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))
        
        # Release GPU memory after processing a certain number of frames
        if processed_frames % 100 == 0:
            processed_frames = 0
            del out  # Delete the variable holding the processed frame
            writer.release()
            gc.collect()  # Perform manual garbage collection to release GPU memory
            print("Memory reset")
    
    cap.release()
    writer.release()
    return 'output.mp4'

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