Upload run_gui.py

run_gui.py

@@ -0,0 +1,104 @@
+import torch
+from PIL import Image
+import matplotlib.pyplot as plt
+from torchvision import transforms
+import gradio as gr
+from models.birefnet import BiRefNet
+import io
+import tempfile  # Импортируем tempfile
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print("Используем устройство:", device)
+
+# Конфигурация моделей
+MODEL_CONFIG = {
+    "BiRefNet_HR": {
+        "repo": "ZhengPeng7/BiRefNet_HR",
+        "image_size": (2048, 2048)
+    },
+    "BiRefNet": {
+        "repo": "ZhengPeng7/BiRefNet",
+        "image_size": (1024, 1024)
+    },
+    "BiRefNet-matting": {
+        "repo": "ZhengPeng7/BiRefNet-matting",
+        "image_size": (1024, 1024)
+    },
+        "BiRefNet-portrait": {
+        "repo": "ZhengPeng7/BiRefNet-portrait",
+        "image_size": (1024, 1024)
+    },
+        "BiRefNet-HRSOD": {
+        "repo": "ZhengPeng7/BiRefNet-HRSOD",
+        "image_size": (1024, 1024)
+    },
+}
+
+# Кэш для загруженных моделей
+loaded_models = {}
+
+def load_model(model_name):
+    if model_name not in loaded_models:
+        print(f"Загрузка модели {model_name}...")
+        model = BiRefNet.from_pretrained(MODEL_CONFIG[model_name]["repo"])
+        model.to(device).eval()
+        if device == 'cuda':
+            model.half()  # FP16 для CUDA
+        loaded_models[model_name] = model
+    return loaded_models[model_name]
+
+def extract_object(image, model_name):
+    # Загрузка выбранной модели
+    model = load_model(model_name)
+    config = MODEL_CONFIG[model_name]
+
+    # Преобразование изображения
+    transform = transforms.Compose([
+        transforms.Resize(config["image_size"]),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+
+    input_tensor = transform(image).unsqueeze(0)
+    input_tensor = input_tensor.to(device)
+    if device == 'cuda':
+        input_tensor = input_tensor.half()
+
+    with torch.no_grad():
+        preds = model(input_tensor)[-1].sigmoid().cpu()
+
+    mask = transforms.ToPILImage()(preds[0].squeeze())
+    mask = mask.resize(image.size)
+
+    result = image.copy()
+    result.putalpha(mask)
+
+    # Save as PNG to a temporary file
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+        result.save(tmp_file)
+        temp_filepath = tmp_file.name # Get the temporary file path
+
+    return temp_filepath  # Return the filepath to Gradio
+
+# Создаем интерфейс с выбором модели
+iface = gr.Interface(
+    fn=extract_object,
+    inputs=[
+        gr.Image(type="pil", label="Входное изображение"),
+        gr.Dropdown(
+            choices=list(MODEL_CONFIG.keys()),
+            value="BiRefNet_HR",
+            label="Выбор модели"
+        )
+    ],
+    outputs=gr.Image(type="filepath", label="Результат"), # Output type is filepath
+    title="BiRefNet - Интерактивная сегментация",
+    description=(
+        "Выберите модель и загрузите изображение для сегментации. "
+        "Доступные модели: BiRefNet_HR (2048x2048), BiRefNet (1024x1024), BiRefNet-lite-2K (2048x2048)"
+    ),
+    allow_flagging="never" # Disable flagging, as it is not needed and causing issues
+)
+
+if __name__ == "__main__":
+    iface.launch(share=True)