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

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)
    
    writer = None
    frame_count = 0
    max_frames = 50  # Maximum number of frames to process at a time
    segment_count = 0  # Count of video segments processed
    
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # Total number of frames in the video
    processed_frames = 0  # Number of frames processed so far
    
    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(f'output_{segment_count}.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)
    
        out = remover.process(img, type='green')
        writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
        
        frame_count += 1
        processed_frames += 1
        
        # Print progress information
        print(f"Processing: {processed_frames}/{total_frames} frames")
        
        if frame_count >= max_frames:
            cap.release()
            writer.release()
            
            # Process the next segment
            cap = cv2.VideoCapture(video)
            for _ in range(frame_count):
                cap.read()  # Skip frames that have already been processed
            writer = None
            frame_count = 0
            segment_count += 1
    
    cap.release()
    if writer is not None:
        writer.release()

    # Concatenate all the video segments
    final_output = cv2.VideoWriter('final_output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)
    for i in range(segment_count):
        cap = cv2.VideoCapture(f'output_{i}.mp4')
        while cap.isOpened():
            ret, frame = cap.read()
            if ret is False:
                break
            final_output.write(frame)
        cap.release()
        os.remove(f'output_{i}.mp4')  # Delete the segment file
    final_output.release()

    return 'final_output.mp4'

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