""" backend/inference.py ==================== Model loading + preprocessing + prediction for the CityVision segmentation API. The trained checkpoints in ``backend/model/`` are plain ``state_dict``s produced by the ``src/training_*.py`` scripts, so we rebuild the **exact same** model classes from ``src/`` to load them, and we apply the **same** ImageNet preprocessing as ``src/dataloader.py`` so inference matches training. """ import os import sys from typing import Optional, Tuple import numpy as np import torch import torch.nn.functional as F from PIL import Image # Make the shared cityvision package importable. _PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) if _PROJECT_ROOT not in sys.path: sys.path.insert(0, _PROJECT_ROOT) # Single source of truth for classes/palette and model architectures. from cityvision.constants import CLASS_NAMES, NUM_CLASSES, PALETTE # noqa: E402 from cityvision.models import build_model # noqa: E402,F401 (re-exported) # ============================================================ # Constants # ============================================================ IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) # Inference resolution (H, W) — same as training. IMG_SIZE = (512, 1024) MODEL_DIR = os.path.join(_PROJECT_ROOT, "backend", "model") # architecture (as logged in MLflow) -> checkpoint filename CKPT_BY_ARCH = { "UNet": "unet_best.pth", "ResNet34-UNet": "resnet_best.pth", "ResNet50-UNet": "resnet50_best.pth", "VGG16-UNet": "vgg_best.pth", "SegNet": "segnet_best.pth", } # Fallback preference when we cannot ask MLflow which run scored best. _ARCH_PRIORITY = ["ResNet50-UNet", "ResNet34-UNet", "VGG16-UNet", "UNet", "SegNet"] # ============================================================ # Checkpoint selection # ============================================================ def _arch_from_filename(fname: str) -> Optional[str]: for arch, f in CKPT_BY_ARCH.items(): if f == os.path.basename(fname): return arch return None def _best_from_mlflow() -> Optional[Tuple[str, str, float]]: """Ask the MLflow backend for the best run whose checkpoint exists. Best-effort.""" db = os.path.join(_PROJECT_ROOT, "mlflow.db") if not os.path.exists(db): return None try: import mlflow from mlflow.tracking import MlflowClient mlflow.set_tracking_uri(f"sqlite:///{db}") client = MlflowClient() exp = client.get_experiment_by_name("urban-segmentation") if exp is None: return None for run in client.search_runs( [exp.experiment_id], order_by=["metrics.best_val_mIoU DESC"] ): arch = run.data.params.get("architecture") miou = run.data.metrics.get("best_val_mIoU") if arch in CKPT_BY_ARCH and miou is not None: ckpt = os.path.join(MODEL_DIR, CKPT_BY_ARCH[arch]) if os.path.exists(ckpt): return arch, ckpt, float(miou) except Exception: return None return None def select_checkpoint() -> Tuple[str, str, Optional[float]]: """ Decide which model to serve, in priority order: 1. CITYVISION_CHECKPOINT env var (explicit .pth path) 2. CITYVISION_ARCH env var (architecture name) 3. best run recorded in mlflow.db (if present) 4. first available checkpoint by quality priority Returns (architecture, checkpoint_path, best_val_mIoU or None). """ env_ckpt = os.environ.get("CITYVISION_CHECKPOINT") if env_ckpt: if not os.path.exists(env_ckpt): raise FileNotFoundError(f"CITYVISION_CHECKPOINT not found: {env_ckpt}") arch = os.environ.get("CITYVISION_ARCH") or _arch_from_filename(env_ckpt) if arch is None: raise ValueError( "Could not infer architecture from checkpoint name; " "set CITYVISION_ARCH." ) return arch, env_ckpt, None env_arch = os.environ.get("CITYVISION_ARCH") if env_arch: if env_arch not in CKPT_BY_ARCH: raise ValueError(f"Unknown CITYVISION_ARCH={env_arch!r}.") ckpt = os.path.join(MODEL_DIR, CKPT_BY_ARCH[env_arch]) if not os.path.exists(ckpt): raise FileNotFoundError(f"No checkpoint for {env_arch}: {ckpt}") return env_arch, ckpt, None best = _best_from_mlflow() if best is not None: return best for arch in _ARCH_PRIORITY: ckpt = os.path.join(MODEL_DIR, CKPT_BY_ARCH[arch]) if os.path.exists(ckpt): return arch, ckpt, None raise RuntimeError(f"No model checkpoints found in {MODEL_DIR}.") # ============================================================ # Segmenter — load once, predict many # ============================================================ class Segmenter: """Loads the selected model once and runs segmentation on PIL images.""" def __init__(self, device: Optional[str] = None): self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.arch, self.checkpoint_path, self.val_miou = select_checkpoint() self.model = build_model(self.arch).to(self.device) state_dict = torch.load( self.checkpoint_path, map_location=self.device, weights_only=True ) self.model.load_state_dict(state_dict) self.model.eval() self._mean = IMAGENET_MEAN.to(self.device) self._std = IMAGENET_STD.to(self.device) # --- preprocessing --- def _preprocess(self, image: Image.Image) -> torch.Tensor: img = image.convert("RGB").resize((IMG_SIZE[1], IMG_SIZE[0]), Image.BILINEAR) arr = torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0) arr = arr.permute(2, 0, 1).unsqueeze(0).to(self.device) # (1,3,H,W) return (arr - self._mean) / self._std @torch.no_grad() def predict(self, image: Image.Image) -> np.ndarray: """ Returns the predicted class-index mask (H, W) at the ORIGINAL image size. """ orig_w, orig_h = image.size x = self._preprocess(image) logits = self.model(x) # (1, C, h, w) # Upsample logits to the original resolution, then argmax. logits = F.interpolate( logits, size=(orig_h, orig_w), mode="bilinear", align_corners=False ) return logits.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.uint8) # --- output renderings --- @staticmethod def colorize(mask: np.ndarray) -> Image.Image: """Class-index mask (H, W) -> colour RGB PIL image.""" return Image.fromarray(PALETTE[mask], mode="RGB") @staticmethod def overlay( image: Image.Image, mask: np.ndarray, alpha: float = 0.5 ) -> Image.Image: """Blend the colour mask over the original image.""" base = image.convert("RGB") color = Image.fromarray(PALETTE[mask], mode="RGB").resize( base.size, Image.NEAREST ) return Image.blend(base, color, alpha) @staticmethod def rle_encode(mask: np.ndarray) -> list: """Row-major run-length encoding of the mask. Returns a list of [class_id, run_length] pairs. Reconstruct with: flat = np.concatenate([np.full(n, c) for c, n in runs]) mask = flat.reshape(height, width) Far smaller than the dense grid for segmentation masks (large contiguous regions of the same class). """ flat = mask.reshape(-1) if flat.size == 0: return [] # run boundaries = positions where the value changes change = np.nonzero(np.diff(flat))[0] + 1 starts = np.concatenate([[0], change]) ends = np.concatenate([change, [flat.size]]) return [[int(flat[s]), int(e - s)] for s, e in zip(starts, ends)] @staticmethod def class_summary(mask: np.ndarray) -> list: """Per-class pixel coverage of the prediction, sorted by area desc.""" total = mask.size ids, counts = np.unique(mask, return_counts=True) out = [] for cls, cnt in zip(ids.tolist(), counts.tolist()): out.append( { "class_id": int(cls), "class_name": CLASS_NAMES[cls], "color": [int(c) for c in PALETTE[cls]], "pixel_count": int(cnt), "percentage": round(100.0 * cnt / total, 2), } ) out.sort(key=lambda d: d["pixel_count"], reverse=True) return out def info(self) -> dict: return { "architecture": self.arch, "checkpoint": os.path.basename(self.checkpoint_path), "val_mIoU": self.val_miou, "device": self.device, "num_classes": NUM_CLASSES, "classes": CLASS_NAMES, "palette": [[int(c) for c in PALETTE[i]] for i in range(NUM_CLASSES)], "input_size": {"height": IMG_SIZE[0], "width": IMG_SIZE[1]}, }