""" predictor.py ------------ Model loading and inference for the Crop vs No-Crop SegFormer-B2 app. Architecture is rebuilt from the `nvidia/mit-b2` backbone config with num_labels=2, then the fine-tuned `best.pt` state dict is loaded on top. Robust loader: auto-detects common key-naming differences, matches tensors by name AND shape, and reports detailed diagnostics (including whether the segmentation HEAD/classifier loaded) so an all-one-class prediction can be debugged precisely. """ from __future__ import annotations import logging import os import numpy as np import torch import torch.nn.functional as F from PIL import Image from torchvision import transforms from transformers import SegformerConfig, SegformerForSemanticSegmentation logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # # Configuration # --------------------------------------------------------------------------- # MODEL_NAME = os.environ.get("MODEL_NAME", "nvidia/mit-b2") NUM_LABELS = 2 INPUT_SIZE = int(os.environ.get("INPUT_SIZE", "512")) IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) WEIGHTS_FILENAME = os.environ.get("WEIGHTS_FILENAME", "best.pt") HF_REPO_ID = os.environ.get("HF_REPO_ID", "").strip() HF_TOKEN = os.environ.get("HF_TOKEN", "").strip() or None class ModelLoadError(Exception): """Raised when the model or its weights cannot be loaded.""" class InferenceError(Exception): """Raised when inference fails for a given image.""" def _resolve_weights_path() -> str: """Return a local path to best.pt (download from HF Hub if HF_REPO_ID set).""" if HF_REPO_ID: try: from huggingface_hub import hf_hub_download return hf_hub_download( repo_id=HF_REPO_ID, filename=WEIGHTS_FILENAME, token=HF_TOKEN ) except Exception as exc: # noqa: BLE001 raise ModelLoadError( f"Failed to download '{WEIGHTS_FILENAME}' from '{HF_REPO_ID}': {exc}" ) from exc here = os.path.join(os.path.dirname(__file__), WEIGHTS_FILENAME) if os.path.exists(here): return here if os.path.exists(WEIGHTS_FILENAME): return WEIGHTS_FILENAME raise ModelLoadError( f"Weights file '{WEIGHTS_FILENAME}' not found. Commit it to the repo " "or set the HF_REPO_ID environment variable." ) def _extract_state_dict(checkpoint): if isinstance(checkpoint, dict): for key in ("state_dict", "model_state_dict", "model"): inner = checkpoint.get(key) if isinstance(inner, dict): return inner return checkpoint def _strip_module(state_dict): cleaned = {} for key, value in state_dict.items(): new_key = key while new_key.startswith("module."): new_key = new_key[len("module."):] cleaned[new_key] = value return cleaned def _align_state_dict(ckpt_sd, model_sd): """ Match checkpoint tensors to model tensors by name (with a few common transforms) AND shape. Returns (aligned_dict, info) where info reports matched count, the model layers still MISSING weights, the checkpoint layers that were UNUSED, and whether the classifier head loaded. """ ckpt_sd = _strip_module(ckpt_sd) transforms_to_try = [ lambda k: k, # as-is lambda k: "segformer." + k, # add prefix lambda k: k[len("segformer."):] # strip prefix if k.startswith("segformer.") else None, lambda k: k[len("model."):] if k.startswith("model.") else None, ] best = {} used_ckpt_keys = set() for key, value in ckpt_sd.items(): for transform in transforms_to_try: new_key = transform(key) if ( new_key and new_key in model_sd and new_key not in best and tuple(model_sd[new_key].shape) == tuple(value.shape) ): best[new_key] = value used_ckpt_keys.add(key) break missing = sorted(set(model_sd) - set(best)) unexpected = sorted(set(ckpt_sd) - used_ckpt_keys) # The decode-head classifier is what turns features into crop/background. classifier_loaded = any("classifier" in k for k in best) info = { "model_total": len(model_sd), "ckpt_total": len(ckpt_sd), "matched": len(best), "missing": missing, "unexpected": unexpected, "classifier_loaded": classifier_loaded, } return best, info class CropSegmenter: """Loads the fine-tuned SegFormer-B2 model and performs inference.""" def __init__(self) -> None: self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu" ) logger.info("Using device: %s", self.device) self.last_logit_summary: dict | None = None # Build architecture from config (weights come from best.pt) try: seg_config = SegformerConfig.from_pretrained( MODEL_NAME, num_labels=NUM_LABELS ) self.model = SegformerForSemanticSegmentation(seg_config) except Exception as exc: # noqa: BLE001 try: self.model = SegformerForSemanticSegmentation.from_pretrained( MODEL_NAME, num_labels=NUM_LABELS, ignore_mismatched_sizes=True ) except Exception as exc2: # noqa: BLE001 raise ModelLoadError( f"Could not build SegFormer from '{MODEL_NAME}': {exc2}" ) from exc weights_path = _resolve_weights_path() try: try: checkpoint = torch.load(weights_path, map_location=self.device) except Exception: checkpoint = torch.load( weights_path, map_location=self.device, weights_only=False ) except Exception as exc: # noqa: BLE001 raise ModelLoadError( f"Failed to read weights file '{weights_path}': {exc}" ) from exc ckpt_sd = _extract_state_dict(checkpoint) aligned, self.load_info = _align_state_dict( ckpt_sd, self.model.state_dict() ) self.model.load_state_dict(aligned, strict=False) logger.info( "Weights matched %d/%d (ckpt had %d). classifier_loaded=%s", self.load_info["matched"], self.load_info["model_total"], self.load_info["ckpt_total"], self.load_info["classifier_loaded"], ) if self.load_info["matched"] == 0: raise ModelLoadError( "None of the weights in best.pt matched the SegFormer-B2 " "architecture. The checkpoint is likely from a different model." ) self.model.to(self.device) self.model.eval() def _build_preprocess(self, size: int, normalization: str): steps = [transforms.Resize((size, size)), transforms.ToTensor()] if normalization == "imagenet": steps.append(transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)) elif normalization == "minus_one_one": steps.append(transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))) # "zero_one" -> ToTensor only (pixels in [0, 1]) return transforms.Compose(steps) @torch.no_grad() def predict( self, image: Image.Image, input_size: int | None = None, swap_classes: bool = False, normalization: str = "imagenet", ) -> np.ndarray: """ Run segmentation. Returns a uint8 mask (H, W) in {0, 1} at the original resolution. Also stores per-call logit stats in self.last_logit_summary. """ try: size = int(input_size or INPUT_SIZE) if image.mode != "RGB": image = image.convert("RGB") orig_w, orig_h = image.size preprocess = self._build_preprocess(size, normalization) tensor = preprocess(image).unsqueeze(0).to(self.device) logits = self.model(pixel_values=tensor).logits # (1, 2, h/4, w/4) upsampled = F.interpolate( logits, size=(orig_h, orig_w), mode="bilinear", align_corners=False, ) class0 = upsampled[:, 0] class1 = upsampled[:, 1] self.last_logit_summary = { "frac_pixels_pred_crop": float((class1 > class0).float().mean()), "mean_gap_crop_minus_bg": float((class1 - class0).mean().item()), "logit_min": float(upsampled.min().item()), "logit_max": float(upsampled.max().item()), } mask = upsampled.argmax(dim=1).squeeze(0).to("cpu").numpy().astype(np.uint8) if swap_classes: mask = (1 - mask).astype(np.uint8) return mask except Exception as exc: # noqa: BLE001 raise InferenceError(f"Inference failed: {exc}") from exc