prebackbone_infer.py

"""Prebackbone enrichment inference (A11_CA) — standalone, no ultralytics."""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import numpy as np
import torch
from PIL import Image

from a11_ca import build_prebackbone

PREBACKBONE_ONLY_NAME = "prebackbone_a11_ca.pt"


def _to_numpy_u8(arr) -> np.ndarray:
    """Canonical uint8 HWC in conda numpy (avoids ~/.local numpy vs torch/opencv)."""
    if isinstance(arr, Image.Image):
        arr = arr.convert("RGB")
        w, h = arr.size
        return np.frombuffer(arr.tobytes(), dtype=np.uint8).reshape((h, w, 3)).copy()
    raw = np.asarray(arr)
    if raw.ndim == 2:
        raw = np.stack([raw, raw, raw], axis=-1)
    elif raw.shape[-1] > 3:
        raw = raw[..., :3]
    return np.array(raw.tolist(), dtype=np.uint8, order="C")


def _here() -> Path:
    return Path(__file__).resolve().parent


def _prebackbone_only_path() -> Path:
    env = os.environ.get("PREBACKBONE_ONLY_WEIGHTS", "").strip()
    if env:
        return Path(env).expanduser()
    return _here() / "weights" / PREBACKBONE_ONLY_NAME


def _full_checkpoint_path() -> Path | None:
    env = os.environ.get("PREBACKBONE_FULL_CKPT", "").strip()
    if env:
        p = Path(env).expanduser()
        return p if p.exists() else None
    for candidate in (
        _here() / "weights" / "best.pt",
        _here().parent
        / "ultralytics"
        / "Proposed"
        / "yolo12_training"
        / "HRIPCB_Results"
        / "yolo12n_hripcb_200epochs_batch16"
        / "weights"
        / "best.pt",
    ):
        if candidate.exists():
            return candidate
    return None


def _download_hf_file(repo_id: str, filename: str) -> Path:
    from huggingface_hub import hf_hub_download

    dest_dir = _here() / "weights"
    dest_dir.mkdir(parents=True, exist_ok=True)
    return Path(hf_hub_download(repo_id=repo_id, filename=filename, local_dir=str(dest_dir)))


def _resolve_weights_path() -> Path:
    pb_only = _prebackbone_only_path()
    if pb_only.exists():
        return pb_only

    hf_repo = os.environ.get("HF_MODEL_REPO", "").strip()
    if hf_repo:
        try:
            return _download_hf_file(hf_repo, PREBACKBONE_ONLY_NAME)
        except Exception:
            pass
        env_weights = os.environ.get("PREBACKBONE_WEIGHTS", PREBACKBONE_ONLY_NAME)
        return _download_hf_file(hf_repo, env_weights)

    env = os.environ.get("PREBACKBONE_WEIGHTS", "").strip()
    if env and Path(env).expanduser().exists():
        return Path(env).expanduser()

    return pb_only


def _maybe_extract_from_full_ckpt(pb_only_path: Path) -> Path:
    if pb_only_path.exists():
        return pb_only_path
    full = _full_checkpoint_path()
    if full is None:
        return pb_only_path
    from extract_prebackbone_weights import extract

    print(f"[prebackbone] Extracting weights from {full} -> {pb_only_path}")
    return extract(full, pb_only_path)


def _filter_state_dict(state: dict, module: torch.nn.Module) -> dict:
    expected = set(module.state_dict().keys())
    filtered = {k: v for k, v in state.items() if k in expected}
    if len(filtered) < len(expected):
        missing = expected - set(filtered.keys())
        raise RuntimeError(f"Prebackbone weights missing keys: {sorted(missing)[:8]}...")
    return filtered


def _load_prebackbone_module(weights_path: Path, device: torch.device) -> torch.nn.Module:
    if not weights_path.exists():
        weights_path = _maybe_extract_from_full_ckpt(weights_path)

    if not weights_path.exists():
        raise FileNotFoundError(
            f"Prebackbone weights not found: {weights_path}\n"
            "Run: python extract_prebackbone_weights.py --ckpt weights/best.pt\n"
            "Or set PREBACKBONE_ONLY_WEIGHTS / HF_MODEL_REPO."
        )

    try:
        payload = torch.load(weights_path, map_location="cpu", weights_only=True)
    except TypeError:
        payload = torch.load(weights_path, map_location="cpu")

    if isinstance(payload, dict) and "state_dict" in payload:
        name = str(payload.get("prebackbone", "A11_CA")).upper()
        channels = int(payload.get("channels", 3))
        state = payload["state_dict"]
    else:
        name, channels, state = "A11_CA", 3, payload

    module = build_prebackbone(name, channels=channels)
    if module is None:
        raise RuntimeError(f"build_prebackbone({name}) returned None")
    state = _filter_state_dict(state, module)
    missing, unexpected = module.load_state_dict(state, strict=True)
    if missing or unexpected:
        raise RuntimeError(f"State dict mismatch: missing={missing}, unexpected={unexpected}")
    return module.to(device).eval()


def _load_image_rgb(image: str | Path | Image.Image | np.ndarray) -> np.ndarray:
    if isinstance(image, Image.Image):
        return _to_numpy_u8(image.convert("RGB"))
    if isinstance(image, np.ndarray):
        arr = image
        if arr.ndim == 2:
            return _to_numpy_u8(np.stack([arr, arr, arr], axis=-1))
        if arr.shape[2] == 4:
            return _to_numpy_u8(arr[..., :3])
        return _to_numpy_u8(arr[..., :3] if arr.shape[2] >= 3 else arr)
    path = Path(image)
    if not path.exists():
        raise FileNotFoundError(f"Unable to read image: {path}")
    return _to_numpy_u8(Image.open(path).convert("RGB"))


def _img_to_tensor_rgb(im_rgb: np.ndarray, device: torch.device) -> torch.Tensor:
    arr = np.ascontiguousarray(_to_numpy_u8(im_rgb), dtype=np.uint8)
    x = torch.tensor(arr, device=device, dtype=torch.float32)
    return x.permute(2, 0, 1).contiguous().unsqueeze(0) / 255.0


def _tensor_to_rgb_u8(x: torch.Tensor) -> np.ndarray:
    if x.ndim == 4:
        x = x[0]
    hwc = x.detach().float().clamp(0.0, 1.0).mul(255.0).round().byte().permute(1, 2, 0).cpu()
    return np.array(hwc.tolist(), dtype=np.uint8)


class PreBackboneEnricher:
    """Runs A11_CA prebackbone only (defect + golden -> enriched, same spatial size)."""

    def __init__(self, weights: str | Path | None = None, device: str | None = None):
        if device is None:
            device = os.environ.get("PREBACKBONE_DEVICE") or (
                "cuda" if torch.cuda.is_available() else "cpu"
            )
        self.device = torch.device(device)
        self.weights = Path(weights) if weights else _resolve_weights_path()
        self.prebackbone = _load_prebackbone_module(self.weights, self.device)

    @torch.inference_mode()
    def enrich(
        self,
        defect: str | Path | Image.Image | np.ndarray,
        reference: str | Path | Image.Image | np.ndarray,
        *,
        return_reference: bool = False,
    ) -> np.ndarray | tuple[np.ndarray, np.ndarray, np.ndarray]:
        defect_rgb = _load_image_rgb(defect)
        golden_rgb = _load_image_rgb(reference)

        if defect_rgb.shape != golden_rgb.shape:
            raise ValueError(
                f"Defect and reference must have the same shape (HxWxC), "
                f"got {defect_rgb.shape} vs {golden_rgb.shape}. "
                "Use pre-aligned pairs (e.g. training prebackbone_samples) with no extra resizing."
            )

        defect_t = _img_to_tensor_rgb(defect_rgb, self.device)
        golden_t = _img_to_tensor_rgb(golden_rgb, self.device)
        _ = self.prebackbone(defect_t, golden_t)
        dbg: dict[str, Any] = getattr(self.prebackbone, "_debug", {})
        enriched = dbg.get("enriched")
        if enriched is None:
            raise RuntimeError("Prebackbone did not populate _debug['enriched'].")

        enriched_rgb = _tensor_to_rgb_u8(enriched)
        if return_reference:
            return defect_rgb, golden_rgb, enriched_rgb
        return enriched_rgb


_enricher: PreBackboneEnricher | None = None


def get_enricher() -> PreBackboneEnricher:
    global _enricher
    if _enricher is None:
        _enricher = PreBackboneEnricher()
    return _enricher


def enrich_pair(
    defect: str | Path | Image.Image | np.ndarray,
    reference: str | Path | Image.Image | np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    return get_enricher().enrich(defect, reference, return_reference=True)  # type: ignore[return-value]