import gradio as gr
import numpy as np
from PIL import Image
import base64
import io
from medsam2_model import MedSAM2

# Загрузка модели
model = MedSAM2("MedSAM2_pretrain_10ep_b1_AMD-SD_sam2_hiera_t.pth")

def roboflow_predict(data):
    # Извлекаем base64 изображение
    image_b64 = data.get("image")
    box = data.get("box", [])

    if not image_b64 or not box:
        return {"error": "Missing image or box"}

    # Преобразуем base64 в PIL
    image_bytes = base64.b64decode(image_b64.split(",")[-1])
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    image_np = np.array(image)

    # Предсказание
    mask = model.predict(image_np, box)
    mask_img = Image.fromarray(mask * 255)

    # Кодируем маску обратно в base64
    buffer = io.BytesIO()
    mask_img.save(buffer, format="PNG")
    mask_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

    return {
        "predictions": [
            {"mask": f"data:image/png;base64,{mask_b64}"}
        ]
    }

# Интерфейс Gradio с API
demo = gr.Interface(
    fn=roboflow_predict,
    inputs=gr.JSON(label="Roboflow-style Input"),
    outputs=gr.JSON(label="Segmentation Output"),
    title="MedSAM2 for CVAT",
    description="Gradio API для CVAT Models. Передавайте base64 изображение и координаты прямоугольника."
)

if __name__ == "__main__":
    demo.launch()