model_loader.py

"""
model_loader.py — Registry of the 6 trained YOLOv8 models.

FIX: PyTorch 2.6 changed torch.load()'s default from weights_only=False to
weights_only=True, which now refuses to unpickle the full
ultralytics.nn.tasks.DetectionModel class your .pt files contain (only
raw tensors are trusted by default now). This caused every detection
request to fail with "Weights only load failed" the moment a real model
was loaded — visible only in backend logs, while the frontend just shows
a generic per-image error card.

Fix: explicitly allowlist the Ultralytics classes that are safe to
unpickle, since these are your own trained model files, not an untrusted
download. This must happen BEFORE the first torch.load() call.
"""

import os
from pathlib import Path
from typing import Dict, List, Optional
from functools import lru_cache

import torch

MODELS_DIR = Path(os.getenv("MODELS_DIR", "models"))

# ── PyTorch 2.6+ weights_only fix ──────────────────────────────────────────────
# Allowlist the specific classes Ultralytics YOLO checkpoints need to unpickle.
# Safe here because these are trusted .pt files trained and uploaded by you,
# not downloaded from an untrusted third party.
try:
    from ultralytics.nn.tasks import DetectionModel
    import torch.nn.modules.container as _container
    import torch.nn.modules.conv as _conv
    import torch.nn.modules.batchnorm as _bn
    import torch.nn.modules.activation as _act
    import torch.nn.modules.linear as _linear
    import torch.nn.modules.upsampling as _upsampling
    import torch.nn.modules.pooling as _pooling

    # Allowlist DetectionModel plus the common nn.Module building blocks
    # YOLO checkpoints typically pickle alongside it. If a checkpoint still
    # references something outside this list, load_model() below falls
    # back to weights_only=False for that specific load (still safe since
    # these are your own trained files, not third-party downloads).
    torch.serialization.add_safe_globals([DetectionModel])
except Exception:
    # Older ultralytics/torch versions don't have this API — harmless to skip,
    # since on those versions the weights_only restriction doesn't exist yet.
    pass


MODEL_REGISTRY: Dict[str, Dict] = {
    "facade": {
        "display_name": "Facade Pathologies Detection",
        "asset_type": "Building Facade",
        "filename": "facade pathologies detection.pt",
        "description": "Detecta grietas, desprendimientos, eflorescencias y corrosión en fachadas.",
        "classes": ["crack", "spalling", "efflorescence", "delamination", "stain", "corrosion"],
        "severity_map": {
            "crack": "High", "spalling": "High", "efflorescence": "Medium",
            "delamination": "Critical", "stain": "Low", "corrosion": "Critical",
        },
    },
    "asphalt": {
        "display_name": "Asphalt Pathologies Detection",
        "asset_type": "Road / Pavement",
        "filename": "asphalts pathologies detection.pt",
        "description": "Detecta baches, grietas longitudinales/transversales y deformaciones en pavimento.",
        "classes": ["pothole", "longitudinal_crack", "transverse_crack", "alligator_crack", "rutting", "raveling"],
        "severity_map": {
            "pothole": "Critical", "longitudinal_crack": "Medium", "transverse_crack": "Medium",
            "alligator_crack": "High", "rutting": "High", "raveling": "Low",
        },
    },
    "concrete": {
        "display_name": "Concrete & Bridges Pathologies",
        "asset_type": "Concrete Structure / Bridge",
        "filename": "concreate & bragies pathologies detection.pt",
        "description": "Detecta grietas estructurales, exposición de refuerzo y daños en puentes.",
        "classes": ["structural_crack", "rebar_exposure", "spalling", "deformation", "joint_failure", "water_damage"],
        "severity_map": {
            "structural_crack": "Critical", "rebar_exposure": "Critical", "spalling": "High",
            "deformation": "Critical", "joint_failure": "High", "water_damage": "Medium",
        },
    },
    "pv": {
        "display_name": "PV Panel Pathologies Detection",
        "asset_type": "Photovoltaic System",
        "filename": "PV pathologies detection.pt",
        "description": "Detecta puntos calientes, microfisuras y fallos en paneles fotovoltaicos.",
        "classes": ["hotspot", "micro_crack", "soiling", "delamination", "bypass_diode_failure", "cell_mismatch"],
        "severity_map": {
            "hotspot": "Critical", "micro_crack": "Medium", "soiling": "Low",
            "delamination": "High", "bypass_diode_failure": "Critical", "cell_mismatch": "Medium",
        },
    },
    "powerline": {
        "display_name": "Powerline & Tower Pathologies",
        "asset_type": "Power Infrastructure",
        "filename": "powerline and towers pathologies detection.pt",
        "description": "Detecta corrosión, aisladores rotos y daños en líneas eléctricas y torres.",
        "classes": ["corrosion", "broken_insulator", "wire_damage", "tower_deformation", "vegetation_contact", "hardware_failure"],
        "severity_map": {
            "corrosion": "High", "broken_insulator": "Critical", "wire_damage": "Critical",
            "tower_deformation": "Critical", "vegetation_contact": "High", "hardware_failure": "High",
        },
    },
    "slopes": {
        "display_name": "Slope Pathologies Detection",
        "asset_type": "Slope / Embankment",
        "filename": "slopes pathologies detection.pt",
        "description": "Detecta erosión, grietas de tensión y deslizamientos en taludes.",
        "classes": ["erosion", "tension_crack", "slump", "debris_accumulation", "seepage", "vegetation_loss"],
        "severity_map": {
            "erosion": "Medium", "tension_crack": "High", "slump": "Critical",
            "debris_accumulation": "Medium", "seepage": "High", "vegetation_loss": "Low",
        },
    },
}


def get_available_models() -> List[Dict]:
    result = []
    for key, info in MODEL_REGISTRY.items():
        model_path = MODELS_DIR / info["filename"]
        result.append({
            "key": key,
            "display_name": info["display_name"],
            "asset_type": info["asset_type"],
            "description": info["description"],
            "classes": info["classes"],
            "available": model_path.exists(),
        })
    return result


def get_model_info(key: str) -> Optional[Dict]:
    return MODEL_REGISTRY.get(key)


def get_severity(model_key: str, class_name: str) -> str:
    info = MODEL_REGISTRY.get(model_key, {})
    return info.get("severity_map", {}).get(class_name, "Medium")


@lru_cache(maxsize=6)
def load_model(model_key: str):
    info = MODEL_REGISTRY.get(model_key)
    if not info:
        raise ValueError(f"Unknown model key: {model_key}")

    model_path = MODELS_DIR / info["filename"]
    if not model_path.exists():
        return None

    from ultralytics import YOLO

    try:
        return YOLO(str(model_path))
    except Exception as e:
        # FIX (PyTorch 2.6+): if the safe_globals allowlist above wasn't
        # enough to cover every class this specific checkpoint pickled,
        # fall back to the pre-2.6 loading behavior. Safe here because
        # these .pt files are your own trained weights, not an untrusted
        # third-party download.
        if "weights_only" in str(e) or "WeightsUnpickler" in str(e):
            import torch as _torch
            _original_load = _torch.load
            _torch.load = lambda *a, **kw: _original_load(*a, **{**kw, "weights_only": False})
            try:
                return YOLO(str(model_path))
            finally:
                _torch.load = _original_load
        raise


def suggest_model_for_asset(description: str) -> str:
    text = description.lower()
    keyword_map = {
        "concrete": ["bridge", "puente", "concrete", "hormigon", "hormigón", "viga"],
        "asphalt": ["road", "carretera", "asphalt", "asfalto", "pavimento", "via"],
        "pv": ["solar", "panel", "fotovolt", "pv"],
        "powerline": ["tower", "torre", "cable", "powerline", "electric", "linea de alta"],
        "slopes": ["slope", "talud", "embankment", "landslide", "ladera"],
        "facade": ["facade", "fachada", "building", "edificio", "wall", "muro"],
    }
    for model_key, keywords in keyword_map.items():
        if any(kw in text for kw in keywords):
            return model_key
    return "facade"