src/models/inference.py

import os
import logging
from glob import glob

import matplotlib.pyplot as plt
import numpy as np
import torch

from models.preprocessing import process_image_pil, process_metadata_pad20
from models.model_loader import load_model, find_last_conv
from models.cam import GradCAMPlusPlus

logging.basicConfig(level=logging.INFO)


def _resolve_project_root() -> str:
    app_root = os.path.abspath(
        os.environ.get(
            "APP_ROOT",
            os.path.join(os.path.dirname(__file__), "..", ".."),
        )
    )
    return app_root


def _resolve_data_root(project_root: str) -> str:
    explicit_data_root = os.environ.get("DATA_ROOT")
    if explicit_data_root:
        return os.path.abspath(explicit_data_root)

    candidates = [
        "/data",
        os.path.join(project_root, "data"),
        "/app/data",
    ]

    for candidate in candidates:
        preprocess_dir = os.path.join(candidate, "preprocess_data")
        weights_dir = os.path.join(candidate, "weights")
        if os.path.exists(preprocess_dir) and os.path.exists(weights_dir):
            return candidate

    return os.path.join(project_root, "data")


DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
PROJECT_ROOT = _resolve_project_root()
DATA_ROOT = _resolve_data_root(PROJECT_ROOT)

CLASS_LIST = ["NEV", "BCC", "ACK", "SEK", "SCC", "MEL"]

ENCODER_DIR = os.path.join(DATA_ROOT, "preprocess_data")
MODEL_ROOT_PATTERNS = [
    os.path.join(
        DATA_ROOT,
        "weights",
        "TO_BE_USED",
        "*",
        "PAD-UFES-20",
        "*_weights",
        "*",
        "*",
        "model_*_with_one-hot-encoder_512_with_best_architecture",
    ),
    os.path.join(
        DATA_ROOT,
        "weights",
        "TO_BE_USED",
        "*",
        "model_*_with_one-hot-encoder_512_with_best_architecture",
    ),
]

PREFERRED_FOLD = 3
IMPLEMENTED_ATTENTION_MECHANISMS = {
    "no-metadata",
    "no-metadata-without-mlp",
    "concatenation",
    "crossattention",
    "weighted",
    "gfcam",
    "cross-weights-after-crossattention",
    "metablock",
    "only-with-att-intramodal+residual",
    "att-intramodal+residual",
    "att-intramodal+residual+cross-attention-metadados",
    "att-intramodal+residual+cross-attention-metadados+metablock",
    "att-intramodal+residual+cross-attention-metadados+att-intramodal+residual",
}

MODEL_CONFIGS = {}
MODEL_LABELS = {}
MODEL_CACHE = {}
DEFAULT_MODEL_KEY = None
_DISCOVERED = False


def _debug_paths() -> None:
    print(f"[inference] DEVICE={DEVICE}")
    print(f"[inference] PROJECT_ROOT={PROJECT_ROOT}")
    print(f"[inference] DATA_ROOT={DATA_ROOT}")
    print(f"[inference] DATA_ROOT exists? {os.path.exists(DATA_ROOT)}")
    print(f"[inference] ENCODER_DIR={ENCODER_DIR}")
    print(f"[inference] ENCODER_DIR exists? {os.path.exists(ENCODER_DIR)}")

    data_root = DATA_ROOT
    weights_root = os.path.join(data_root, "weights")
    to_be_used_root = os.path.join(weights_root, "TO_BE_USED")

    print(f"[inference] data root={data_root}")
    print(f"[inference] data root exists? {os.path.exists(data_root)}")
    print(f"[inference] weights root={weights_root}")
    print(f"[inference] weights root exists? {os.path.exists(weights_root)}")
    print(f"[inference] TO_BE_USED root={to_be_used_root}")
    print(f"[inference] TO_BE_USED root exists? {os.path.exists(to_be_used_root)}")

    for pattern in MODEL_ROOT_PATTERNS:
        print(f"[inference] MODEL_ROOT_PATTERN={pattern}")


def _parse_cnn_model_name(model_dir_name: str):
    prefix = "model_"
    suffix = "_with_one-hot-encoder_512_with_best_architecture"
    if not model_dir_name.startswith(prefix) or not model_dir_name.endswith(suffix):
        return None
    return model_dir_name[len(prefix):-len(suffix)]


def _find_fold_dir(model_root: str, cnn_model_name: str):
    fold_dirs = sorted(glob(os.path.join(model_root, f"{cnn_model_name}_fold_*")))
    if not fold_dirs:
        return None

    preferred = os.path.join(model_root, f"{cnn_model_name}_fold_{PREFERRED_FOLD}")
    if os.path.exists(os.path.join(preferred, "model.pth")):
        return preferred

    for fold_dir in fold_dirs:
        if os.path.exists(os.path.join(fold_dir, "model.pth")):
            return fold_dir

    return None


def _extract_fold_number(fold_dir: str) -> str:
    name = os.path.basename(fold_dir)
    return name.split("_fold_")[-1] if "_fold_" in name else "?"


def _extract_path_metadata(model_root: str):
    parts = os.path.normpath(model_root).split(os.sep)
    mechanism = os.path.basename(os.path.dirname(model_root))
    unfreeze_weights = "unfrozen_weights"
    num_heads = "8"

    if "PAD-UFES-20" in parts:
        idx = parts.index("PAD-UFES-20")
        if len(parts) > idx + 1:
            unfreeze_weights = parts[idx + 1]
        if len(parts) > idx + 2:
            num_heads = parts[idx + 2]
        if len(parts) > idx + 3:
            mechanism = parts[idx + 3]

    return mechanism, unfreeze_weights, num_heads


def _discover_models() -> None:
    global DEFAULT_MODEL_KEY

    print("[inference] starting model discovery...")
    MODEL_CONFIGS.clear()
    MODEL_LABELS.clear()
    DEFAULT_MODEL_KEY = None

    model_roots = sorted({
        path
        for pattern in MODEL_ROOT_PATTERNS
        for path in glob(pattern)
    })

    print(f"[inference] candidate model roots found: {len(model_roots)}")
    for idx, path in enumerate(model_roots[:20], start=1):
        print(f"[inference] candidate[{idx}] = {path}")
    if len(model_roots) > 20:
        print("[inference] ... additional candidates omitted from log ...")

    for model_root in model_roots:
        mechanism, unfreeze_weights, num_heads = _extract_path_metadata(model_root)
        model_dir_name = os.path.basename(model_root)
        cnn_model_name = _parse_cnn_model_name(model_dir_name)

        if cnn_model_name is None:
            print(f"[inference] skipping invalid model dir name: {model_dir_name}")
            continue

        fold_dir = _find_fold_dir(model_root, cnn_model_name)
        if fold_dir is None:
            print(f"[inference] no valid fold dir found for: {model_root}")
            continue

        model_path = os.path.join(fold_dir, "model.pth")
        if not os.path.exists(model_path):
            print(f"[inference] missing model.pth: {model_path}")
            continue

        fold_number = _extract_fold_number(fold_dir)
        model_key = f"{mechanism}|{cnn_model_name}|{unfreeze_weights}|{num_heads}|{fold_number}"
        supported = mechanism in IMPLEMENTED_ATTENTION_MECHANISMS
        support_tag = "" if supported else " | not-implemented"
        label = f"{mechanism} | {cnn_model_name} | fold {fold_number} | {unfreeze_weights}{support_tag}"

        MODEL_CONFIGS[model_key] = {
            "model_path": model_path,
            "attention_mecanism": mechanism,
            "cnn_model_name": cnn_model_name,
            "unfreeze_weights": unfreeze_weights,
            "num_heads": int(num_heads) if str(num_heads).isdigit() else 8,
            "supported": supported,
            "label": label,
        }
        MODEL_LABELS[model_key] = label
        print(f"[inference] registered model: {label}")

    preferred_defaults = [
        key for key, cfg in MODEL_CONFIGS.items()
        if cfg["attention_mecanism"] == "gfcam"
        and cfg["cnn_model_name"] == "densenet169"
        and cfg["unfreeze_weights"] == "unfrozen_weights"
        and cfg["model_path"].endswith(f"_fold_{PREFERRED_FOLD}/model.pth")
    ]

    if preferred_defaults:
        DEFAULT_MODEL_KEY = preferred_defaults[0]
    elif MODEL_CONFIGS:
        DEFAULT_MODEL_KEY = sorted(MODEL_CONFIGS.keys())[0]

    print(f"[inference] discovery done. models={len(MODEL_CONFIGS)}")
    print(f"[inference] DEFAULT_MODEL_KEY={DEFAULT_MODEL_KEY}")


def ensure_models_discovered() -> None:
    global _DISCOVERED
    if _DISCOVERED:
        return

    _debug_paths()
    _discover_models()
    _DISCOVERED = True


def get_available_model_choices():
    ensure_models_discovered()
    return [(MODEL_LABELS[key], key) for key in sorted(MODEL_LABELS.keys())]


def get_default_model_key():
    ensure_models_discovered()
    return DEFAULT_MODEL_KEY


def get_model_label(model_key):
    ensure_models_discovered()
    return MODEL_LABELS.get(model_key, "Unknown model")


def _get_model_and_cam(model_key):
    ensure_models_discovered()

    if not MODEL_CONFIGS:
        raise RuntimeError(
            f"No compatible model checkpoints were found in "
            f"{os.path.join(DATA_ROOT, 'weights')}. "
            f"Please verify that the model assets were uploaded to the Space."
        )

    if not os.path.exists(ENCODER_DIR):
        raise RuntimeError(
            f"Metadata encoder directory not found: {ENCODER_DIR}. "
            "Please verify that preprocess artifacts were uploaded to the Space."
        )

    if model_key is None:
        model_key = DEFAULT_MODEL_KEY

    if model_key not in MODEL_CONFIGS:
        raise RuntimeError(
            f"Selected model key is invalid: {model_key}. "
            "Please choose a valid model option in the interface."
        )

    cfg = MODEL_CONFIGS[model_key]
    if not cfg["supported"]:
        raise RuntimeError(
            f"The selected mechanism '{cfg['attention_mecanism']}' exists in TO_BE_USED, "
            "but is not implemented in the current inference model forward."
        )

    if model_key not in MODEL_CACHE:
        model_path = cfg["model_path"]
        print(f"[inference] loading model from: {model_path}")

        model = load_model(
            device=DEVICE,
            model_path=model_path,
            cnn_model_name=cfg["cnn_model_name"],
            attention_mecanism=cfg["attention_mecanism"],
            num_heads=cfg["num_heads"],
            unfreeze_weights=cfg["unfreeze_weights"],
        )

        print("[inference] locating final conv layer...")
        target_layer = find_last_conv(model.image_encoder)

        print("[inference] creating GradCAM++ object...")
        MODEL_CACHE[model_key] = (model, GradCAMPlusPlus(model, target_layer))
        print("[inference] model ready.")

    return MODEL_CACHE[model_key]


def run_inference(image_pil, metadata_text, model_key=None):
    model, cam = _get_model_and_cam(model_key)

    print("[inference] processing image...")
    image_tensor = process_image_pil(image_pil, DEVICE)

    print("[inference] processing metadata...")
    metadata_tensor = process_metadata_pad20(
        metadata_text,
        ENCODER_DIR,
        DEVICE
    )

    print("[inference] running forward pass...")
    with torch.no_grad():
        logits = model(image_tensor, metadata_tensor)
        probs = torch.softmax(logits, dim=1)

    pred_class = torch.argmax(probs, dim=1).item()
    confidence = probs[0, pred_class].item()

    print("[inference] generating heatmap...")
    heatmap = cam.generate(
        image_tensor,
        metadata_tensor,
        pred_class
    )

    image_np = np.array(image_pil.convert("RGB"))
    heatmap = np.asarray(heatmap, dtype=np.float32).squeeze()

    if heatmap.shape != image_np.shape[:2]:
        heatmap_tensor = torch.from_numpy(heatmap).unsqueeze(0).unsqueeze(0)
        heatmap = torch.nn.functional.interpolate(
            heatmap_tensor,
            size=image_np.shape[:2],
            mode="bilinear",
            align_corners=False,
        ).squeeze().cpu().numpy()

    heatmap = np.clip(heatmap, 0.0, 1.0)
    alpha_map = np.clip(heatmap * 0.6, 0.0, 0.6)

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.imshow(image_np)
    ax.imshow(heatmap, cmap="jet", alpha=alpha_map)
    ax.axis("off")

    title = f"{CLASS_LIST[pred_class]} | conf={confidence:.3f}"
    ax.set_title(title)

    fig.canvas.draw()
    result = np.array(fig.canvas.renderer.buffer_rgba())
    plt.close(fig)

    print(f"[inference] inference complete: {title}")
    return result, title