import gradio as gr
from rembg import remove
from PIL import Image
import numpy as np
import cv2
import os

# Функция для удаления фона с изображения
def remove_bg_from_image(image):
    output = remove(image)
    return output

# Функция для удаления фона с видео
def remove_bg_from_video(video_path):
    # Открываем видео
    cap = cv2.VideoCapture(video_path)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    
    # Создаем временный файл для сохранения результата
    output_path = "output_video.mp4"
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height), isColor=True)
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Преобразуем кадр в изображение PIL
        frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        
        # Удаляем фон
        frame_no_bg = remove(frame_pil)
        
        # Преобразуем обратно в массив numpy
        frame_no_bg_np = np.array(frame_no_bg)
        
        # Записываем кадр в выходное видео
        out.write(cv2.cvtColor(frame_no_bg_np, cv2.COLOR_RGB2BGR))
    
    cap.release()
    out.release()
    
    return output_path

# Интерфейс Gradio
def process_input(input_type, input_data):
    if input_type == "Image":
        output = remove_bg_from_image(input_data)
        return output
    elif input_type == "Video":
        output = remove_bg_from_video(input_data)
        return output

# Создаем интерфейс Gradio
input_type = gr.Radio(["Image", "Video"], label="Input Type")
input_data = gr.File(label="Input File")  # Используем gr.File вместо gr.inputs.File

interface = gr.Interface(
    fn=process_input,
    inputs=[input_type, input_data],
    outputs=gr.File(label="Output File"),  # Используем gr.File вместо gr.outputs.File
    title="Background Removal",
    description="Remove background from images or videos using rembg."
)

# Запускаем интерфейс
interface.launch()