import gradio as gr from huggingface_hub import hf_hub_download from ultralytics import YOLO from PIL import Image, ImageDraw import torch import numpy as np # 1) Скачиваем веса (пример) REPO_ID = "shaquilledanil/yolo11-skinai" FILENAME = "best.pt" WEIGHT_PATH = hf_hub_download(repo_id=REPO_ID, filename=FILENAME) print("[INFO] Загружаем YOLO-модель (классификация) из:", WEIGHT_PATH) model = YOLO(WEIGHT_PATH) def yolo_classify_inference(img: Image.Image): """ Обрабатывает PIL-изображение через классификационную модель YOLO, возвращает (annotated_image, detections). """ if img is None: print("[ERROR] Пустое изображение.") return None, [] print(f"[INFO] Размер входного изображения: {img.size}") np_img = np.array(img) print(f"[INFO] numpy-данные: dtype={np_img.dtype}, shape={np_img.shape}") print("[INFO] Запускаем predict (Classification)...") results = model.predict(img) if not results or len(results) == 0: print("[INFO] Модель не вернула результатов.") return img, [] # Первый результат res = results[0] # Проверка, есть ли res.probs (классификация) if not hasattr(res, "probs") or res.probs is None: print("[INFO] Модель не вернула res.probs (не классификация?).") return img, [] # Извлекаем top1 top1_cls_id = int(res.probs.top1) top1_conf = float(res.probs.top1conf) class_name = res.names[top1_cls_id] print(f"[INFO] Top1: class={class_name}, conf={top1_conf:.2f}") # Рисуем надпись на изображении annotated = img.copy() draw = ImageDraw.Draw(annotated) text = f"{class_name} ({top1_conf:.2f})" draw.text((10, 10), text, fill="red") # Все вероятности arr = res.probs.data if isinstance(arr, torch.Tensor): arr = arr.cpu().numpy() detections = [] for i, conf in enumerate(arr): c = float(conf) cname = res.names[i] detections.append({ "class_id": i, "class_name": cname, "confidence": c }) return annotated, detections # 2) Создаём Gradio-интерфейс demo = gr.Interface( fn=yolo_classify_inference, inputs=gr.Image(type="pil"), outputs=[gr.Image(type="pil", label="Classification result"), "json"], title="YOLO SkinAI Classification Demo", description="Загрузите изображение для классификации" ) if __name__ == "__main__": # ВАЖНО: Отключаем очередь (queue=False) # Отключаем SSR (ssr=False или ssr_mode=False, в зависимости от версии Gradio) # Включаем api_open=True, чтобы /run/predict был доступен извне (REST API) demo.launch(share=True, ssr_mode=False, show_api = True) demo.queue(api_open = True)