import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from PIL import Image

# Скачиваем веса из вашего репо на Hugging Face
REPO_ID = "shaquilledanil/yolo11-skinai"
FILENAME = "best.pt"
WEIGHT_PATH = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

# Загружаем YOLO-модель
model = YOLO(WEIGHT_PATH)

def yolo_inference(image: Image.Image):
    """
    Принимает PIL.Image, прогоняет через YOLO-модель, 
    возвращает кортеж (annotated_image, detections_info).
    """
    results = model.predict(source=image)

    # Подготовим изображение с нарисованными bboxes
    # Ultralytics в results[0].plot() вернёт PIL.Image с рамками
    annotated_image = results[0].plot() if len(results) else image

    # Соберём текстовую информацию о детекциях
    detections = []
    if len(results):
        boxes = results[0].boxes
        for box in boxes:
            cls_id = int(box.cls[0].item())
            conf = float(box.conf[0].item())
            xyxy = box.xyxy[0].tolist()
            class_name = results[0].names[cls_id]
            detections.append({
                "class_id": cls_id,
                "class_name": class_name,
                "confidence": round(conf, 3),
                "bbox": [round(x, 1) for x in xyxy]
            })

    return annotated_image, detections

# Создаём Gradio-интерфейс:
#  - inputs: одна картинка
#  - outputs: (1) картинка с боксами, (2) JSON (список детекций)
demo = gr.Interface(
    fn=yolo_inference,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Detected Objects"), 
        "json"
    ],
    title="YOLO SkinAI Demo",
    description="Загрузите изображение, чтобы увидеть результат детекции"
)

# Запуск (при локальном тесте):
if __name__ == "__main__":
    demo.launch()