Create app.py

app.py

@@ -0,0 +1,69 @@
+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.inputs.File(label="Input File")
+
+interface = gr.Interface(
+    fn=process_input,
+    inputs=[input_type, input_data],
+    outputs=gr.outputs.File(label="Output File"),
+    title="Background Removal",
+    description="Remove background from images or videos using rembg."
+)
+
+# Запускаем интерфейс
+interface.launch()