model.py

import os
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    SegformerImageProcessor,
    SegformerForSemanticSegmentation,
)
import numpy as np
from ultralytics import YOLO
import matplotlib.pyplot as plt

# --- 本地模型路径配置 (确保您已将文件下载到以下路径) ---
CLS_MODEL_PATH = "models/vit_base"
SEG_MODEL_PATH = "models/segformer_b1"
DET_MODEL_PATH = "yolov8n.pt"  # YOLOv8 会自动管理下载，我们恢复到简单ID

# --- 1. 图像分类模型 (ViT) ---
try:
    # 从本地路径加载
    cls_processor = AutoImageProcessor.from_pretrained(CLS_MODEL_PATH, use_fast=True)
    cls_model = AutoModelForImageClassification.from_pretrained(CLS_MODEL_PATH)
    cls_model.eval()

    cls_labels = [cls_model.config.id2label[i] for i in range(cls_model.config.num_labels)]

    print("✅ Classification model (ViT) loaded from local path.")
except Exception as e:
    # 如果本地路径失败，尝试远程ID作为回退（仅在开发阶段有用）
    CLS_MODEL_ID = "google/vit-base-patch16-224"
    try:
        cls_processor = AutoImageProcessor.from_pretrained(CLS_MODEL_ID, use_fast=True)
        cls_model = AutoModelForImageClassification.from_pretrained(CLS_MODEL_ID)
        cls_model.eval()
        cls_labels = [cls_model.config.id2label[i] for i in range(cls_model.config.num_labels)]
        print("⚠️ Classification model loaded from remote Hub (Fallback).")
    except Exception as e_fb:
        print(f"❌ Error loading classification model: {e}")
        cls_model = None
        cls_labels = []


def cls_predict(input_img: Image.Image) -> dict:
    """图像分类预测函数"""
    if not cls_model: return {"Error": "Classification model not loaded."}

    inputs = cls_processor(images=input_img, return_tensors="pt")

    with torch.no_grad():
        outputs = cls_model(**inputs)
        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)

    top5_prob, top5_indices = torch.topk(probabilities, 5)

    results = {}
    for i in range(5):
        label = cls_labels[top5_indices[0][i].item()]
        prob = top5_prob[0][i].item()
        results[label] = round(prob, 4)

    return results


# --- 2. 语义分割模型 (SegFormer) ---
try:
    # 从本地路径加载
    seg_processor = SegformerImageProcessor.from_pretrained(SEG_MODEL_PATH)
    seg_model = SegformerForSemanticSegmentation.from_pretrained(
        SEG_MODEL_PATH,
        use_safetensors=False,  # 确保使用 pytorch_model.bin
        trust_remote_code=True
    )
    seg_model.eval()

    # 获取 SegFormer (ADE20K) 的类别标签
    seg_labels = [label for label in seg_model.config.id2label.values()]
    SEG_CLASSES = len(seg_labels)

    # 创建固定的颜色映射 (Matplotlib HSV Colormap)
    cmap = plt.cm.get_cmap('hsv', SEG_CLASSES)
    COLOR_MAP_float = cmap(np.arange(SEG_CLASSES))[:, :3]
    COLOR_MAP_int = (COLOR_MAP_float * 255).astype(np.uint8)

    COLOR_MAP_DICT = {}
    for i, label in enumerate(seg_labels):
        rgb_color = COLOR_MAP_int[i].tolist()
        hex_color = '#%02x%02x%02x' % tuple(rgb_color)
        COLOR_MAP_DICT[label] = hex_color

    COLOR_MAP = COLOR_MAP_int

    print("✅ Segmentation model (SegFormer) loaded from local path.")
except Exception as e:
    # 远程回退加载
    SEG_MODEL_ID = "nvidia/segformer-b1-finetuned-ade-512-512"
    try:
        seg_processor = SegformerImageProcessor.from_pretrained(SEG_MODEL_ID)
        seg_model = SegformerForSemanticSegmentation.from_pretrained(SEG_MODEL_ID, use_safetensors=False,
                                                                     trust_remote_code=True)
        seg_model.eval()
        seg_labels = [label for label in seg_model.config.id2label.values()]
        SEG_CLASSES = len(seg_labels)
        cmap = plt.cm.get_cmap('hsv', SEG_CLASSES)
        COLOR_MAP_float = cmap(np.arange(SEG_CLASSES))[:, :3]
        COLOR_MAP_int = (COLOR_MAP_float * 255).astype(np.uint8)
        COLOR_MAP_DICT = {}
        for i, label in enumerate(seg_labels):
            rgb_color = COLOR_MAP_int[i].tolist()
            hex_color = '#%02x%02x%02x' % tuple(rgb_color)
            COLOR_MAP_DICT[label] = hex_color
        COLOR_MAP = COLOR_MAP_int
        print("⚠️ Segmentation model loaded from remote Hub (Fallback).")
    except Exception as e_fb:
        print(f"❌ Error loading segmentation model: {e}")
        seg_model = None
        seg_labels = []
        COLOR_MAP_DICT = {"Error": "#000000"}
        SEG_CLASSES = 0


def seg_predict(input_img: Image.Image) -> Image.Image:
    """语义分割预测函数"""
    if not seg_model: return input_img

    original_size = input_img.size

    inputs = seg_processor(images=input_img, return_tensors="pt")

    with torch.no_grad():
        outputs = seg_model(**inputs)

    logits = outputs.logits.cpu()

    predicted_segmentation = F.interpolate(
        logits,
        size=original_size[::-1],
        mode='bilinear',
        align_corners=False
    ).argmax(dim=1)[0]

    # 可视化：创建彩色掩码
    segmentation_np = predicted_segmentation.numpy().astype(np.uint8)

    # 映射到颜色
    if SEG_CLASSES > 0:
        colored_mask = COLOR_MAP[segmentation_np % SEG_CLASSES]
    else:
        colored_mask = np.zeros((*segmentation_np.shape, 3), dtype=np.uint8)

    mask_img = Image.fromarray(colored_mask).convert("RGBA")

    # 叠加到原图上 (使用 50% 透明度)
    input_img = input_img.convert("RGBA")
    output_img = Image.blend(input_img, mask_img, alpha=0.5).convert("RGB")

    return output_img


# --- 3. 目标检测模型 (YOLOv8) ---
try:
    # 保持使用 ID，让 YOLOv8 自动管理下载和缓存
    det_model = YOLO(DET_MODEL_PATH)

    # 获取 YOLOv8 (COCO) 的类别标签
    det_labels_dict = det_model.names
    det_labels = {str(k): v for k, v in det_labels_dict.items()}

    print("✅ Detection model (YOLOv8n) loaded.")
except Exception as e:
    print(f"❌ Error loading detection model: {e}")
    det_model = None
    det_labels = {}


def det_predict(input_img: Image.Image, conf: float) -> Image.Image:
    """目标检测预测函数"""
    if not det_model: return input_img

    results = det_model.predict(
        source=input_img,
        conf=conf,
        iou=0.5,
        verbose=False
    )

    plotted_image_np = results[0].plot()
    return Image.fromarray(plotted_image_np)


# --- 导出标签列表供 Gradio 界面使用 ---
ALL_CLS_LABELS = cls_labels if 'cls_labels' in locals() else ["Classification labels not available."]
ALL_DET_LABELS = det_labels if 'det_labels' in locals() else {"Error": "Detection labels not available."}
ALL_SEG_COLOR_MAP = COLOR_MAP_DICT if 'COLOR_MAP_DICT' in locals() else {"Error": "#000000"}
ALL_SEG_LABELS = seg_labels if 'seg_labels' in locals() else ["Segmentation labels not available."]