"""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]