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

@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)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    tmpname = random.randint(111111111, 999999999)
    tmp_video = f"/tmp/{tmpname}_video.mp4"
    tmp_audio = f"/tmp/{tmpname}_audio.aac"
    output_video = f"/tmp/{tmpname}_output.mp4"
    
    writer = cv2.VideoWriter(tmp_video, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    
    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")
            break
        
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame).convert('RGB')
        
        processed_frames += 1
        print(f"Processing: {processed_frames}")
        
        out = remover.process(img, type='green')
        writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_RGB2BGR))
    
    cap.release()
    writer.release()
    
    # 오디오 추출
    subprocess.call(['ffmpeg', '-i', video, '-vn', '-acodec', 'copy', tmp_audio])
    
    # 처리된 비디오와 오디오 결합
    subprocess.call(['ffmpeg', '-i', tmp_video, '-i', tmp_audio, '-c', 'copy', output_video])
    
    # 임시 파일 삭제
    os.remove(tmp_video)
    os.remove(tmp_audio)
    
    return output_video

# Gradio 인터페이스 정의
css = """footer { visibility: hidden; }"""

b_interface = gr.Interface(
    fn=doo,
    inputs=[
        gr.Video(label="Video"),
        gr.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",
    css=css,
    title="Video Background Removal",
    description="Select a video and mode to remove the background."
)

# Gradio 앱 실행
if __name__ == "__main__":
    b_interface.launch()