Upload 2 files

app.py

@@ -1,10 +1,9 @@
 """
 app.py
 ------
-Streamlit front-end for the Crop vs No-Crop SegFormer-B2 demo.
-
-Includes a diagnostics panel and Advanced controls (inference size,
-swap crop/background) to debug an all-background (0% coverage) prediction.
+Streamlit front-end for the Crop vs No-Crop SegFormer-B2 demo, with a
+detailed diagnostics panel and Advanced controls (inference size, normalization,
+swap crop/background) to debug an all-one-class prediction.
 """
 
 import logging
@@ -28,6 +27,12 @@ st.set_page_config(
     layout="wide",
 )
 
+NORMALIZATIONS = {
+    "ImageNet (mean/std)": "imagenet",
+    "Simple 0–1 (just /255)": "zero_one",
+    "-1 to 1": "minus_one_one",
+}
+
 
 @st.cache_resource(show_spinner="Loading segmentation model (first run only)…")
 def get_segmenter() -> CropSegmenter:
@@ -35,7 +40,6 @@ def get_segmenter() -> CropSegmenter:
 
 
 def render_sidebar() -> dict:
-    """Render the sidebar and return the selected Advanced options."""
     st.sidebar.header("ℹ️ Model Information")
     st.sidebar.markdown(
         f"""
@@ -46,9 +50,6 @@ def render_sidebar() -> dict:
 """
     )
 
-    st.sidebar.header("🖼️ Supported Formats")
-    st.sidebar.markdown("- JPG\n- JPEG\n- PNG")
-
     st.sidebar.header("🎨 Class Definitions")
     st.sidebar.markdown(
         """
@@ -62,36 +63,50 @@ def render_sidebar() -> dict:
     st.sidebar.header("🛠️ Advanced (debugging)")
     sizes = [256, 384, 512, 640]
     default_index = sizes.index(INPUT_SIZE) if INPUT_SIZE in sizes else 2
-    input_size = st.sidebar.selectbox(
-        "Inference size", sizes, index=default_index,
-        help="Use the same size you trained at if you know it.",
+    input_size = st.sidebar.selectbox("Inference size", sizes, index=default_index)
+    norm_label = st.sidebar.selectbox(
+        "Normalization", list(NORMALIZATIONS.keys()), index=0,
+        help="Must match how images were normalized during training.",
     )
     swap_classes = st.sidebar.checkbox(
-        "Swap crop / background",
-        value=False,
-        help="Tick this if your crop class was label 0 during training.",
+        "Swap crop / background", value=False,
+        help="Tick if your crop class was label 0 during training.",
     )
-    return {"input_size": input_size, "swap_classes": swap_classes}
+    return {
+        "input_size": input_size,
+        "normalization": NORMALIZATIONS[norm_label],
+        "swap_classes": swap_classes,
+    }
 
 
 def render_diagnostics(segmenter: CropSegmenter) -> None:
     info = segmenter.load_info
     matched, total = info["matched"], info["model_total"]
-    with st.sidebar.expander("🔍 Weight-loading diagnostics", expanded=(matched < total)):
+    head_ok = info["classifier_loaded"]
+    needs_attention = (matched < total) or (not head_ok)
+
+    with st.sidebar.expander("🔍 Weight-loading diagnostics", expanded=needs_attention):
         st.write(f"**Matched:** {matched} / {total} model tensors")
         st.write(f"**Checkpoint tensors:** {info['ckpt_total']}")
-        if matched == total:
+        st.write(f"**Classifier head loaded:** {'✅ yes' if head_ok else '❌ NO'}")
+
+        if not head_ok:
+            st.error(
+                "The segmentation head did NOT load — this is why every pixel "
+                "comes out the same class. Copy the two lists below and send "
+                "them to your developer."
+            )
+        elif matched == total:
             st.success("All weights loaded correctly.")
-        elif matched == 0:
-            st.error("No weights matched — wrong checkpoint/architecture.")
         else:
-            st.warning(
-                f"{total - matched} model tensors did NOT receive weights. "
-                "If results look wrong, paste the missing keys below to your "
-                "developer."
-            )
-            if info["missing"]:
-                st.code("\n".join(info["missing"][:20]))
+            st.warning(f"{total - matched} model tensors did not get weights.")
+
+        if info["missing"]:
+            st.caption("Model layers still MISSING weights:")
+            st.code("\n".join(info["missing"][:30]))
+        if info["unexpected"]:
+            st.caption("Checkpoint layers NOT used (unexpected names):")
+            st.code("\n".join(info["unexpected"][:30]))
 
 
 def main() -> None:
@@ -139,6 +154,7 @@ def main() -> None:
                 image,
                 input_size=options["input_size"],
                 swap_classes=options["swap_classes"],
+                normalization=options["normalization"],
             )
             color_mask = utils.create_color_mask(mask)
             overlay = utils.create_overlay(image, mask)
@@ -161,11 +177,29 @@ def main() -> None:
 
     crop_ratio = float((mask == 1).mean()) * 100.0
     st.metric("Estimated crop coverage", f"{crop_ratio:.1f}%")
-    if crop_ratio == 0.0:
+
+    # Show the raw model behaviour — confirms whether it is discriminating.
+    summary = segmenter.last_logit_summary
+    if summary:
+        with st.expander("📈 Model output details (for debugging)"):
+            st.write(
+                f"Pixels predicted as crop (before swap): "
+                f"{summary['frac_pixels_pred_crop'] * 100:.1f}%"
+            )
+            st.write(
+                f"Mean logit gap (crop − background): "
+                f"{summary['mean_gap_crop_minus_bg']:.3f}  "
+                f"(positive favors crop, negative favors background)"
+            )
+            st.write(
+                f"Logit range: {summary['logit_min']:.2f} … {summary['logit_max']:.2f}"
+            )
+
+    if crop_ratio in (0.0, 100.0):
         st.warning(
-            "0% crop detected. Open '🔍 Weight-loading diagnostics' in the "
-            "sidebar to check the weights loaded, and try the 'Swap crop / "
-            "background' toggle or a different inference size."
+            "Every pixel got the same class. Open '🔍 Weight-loading "
+            "diagnostics' to check the classifier head loaded, and try a "
+            "different **Normalization** in the sidebar."
         )
 
     png_bytes = utils.image_to_png_bytes(combined)

predictor.py

@@ -6,12 +6,10 @@ 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.
 
-This version includes a ROBUST weight loader that:
-  * auto-detects common key-naming differences (module./model./segformer.)
-  * matches tensors by name AND shape
-  * reports exactly how many weights matched (self.load_info)
-
-so an all-background prediction (0% coverage) can be diagnosed quickly.
+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
@@ -49,17 +47,12 @@ class InferenceError(Exception):
     """Raised when inference fails for a given image."""
 
 
-# --------------------------------------------------------------------------- #
-# Weights location
-# --------------------------------------------------------------------------- #
 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
 
-            logger.info("Downloading '%s' from HF repo '%s'",
-                        WEIGHTS_FILENAME, HF_REPO_ID)
             return hf_hub_download(
                 repo_id=HF_REPO_ID, filename=WEIGHTS_FILENAME, token=HF_TOKEN
             )
@@ -79,11 +72,7 @@ def _resolve_weights_path() -> str:
     )
 
 
-# --------------------------------------------------------------------------- #
-# Robust state-dict alignment
-# --------------------------------------------------------------------------- #
 def _extract_state_dict(checkpoint):
-    """Pull the actual weight dict out of common checkpoint wrappers."""
     if isinstance(checkpoint, dict):
         for key in ("state_dict", "model_state_dict", "model"):
             inner = checkpoint.get(key)
@@ -93,7 +82,6 @@ def _extract_state_dict(checkpoint):
 
 
 def _strip_module(state_dict):
-    """Remove repeated 'module.' prefixes (from DataParallel/DDP checkpoints)."""
     cleaned = {}
     for key, value in state_dict.items():
         new_key = key
@@ -105,14 +93,13 @@ def _strip_module(state_dict):
 
 def _align_state_dict(ckpt_sd, model_sd):
     """
-    Try a few key transforms and keep the subset of checkpoint tensors that
-    match the model by name AND shape. Returns (aligned_dict, info).
+    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)
 
-    # Each transform is tried per-key (not globally), so checkpoints that mix
-    # prefixes — e.g. backbone keys needing a 'segformer.' prefix while head
-    # keys are already correct — still align fully.
     transforms_to_try = [
         lambda k: k,                                              # as-is
         lambda k: "segformer." + k,                              # add prefix
@@ -122,6 +109,7 @@ def _align_state_dict(ckpt_sd, model_sd):
     ]
 
     best = {}
+    used_ckpt_keys = set()
     for key, value in ckpt_sd.items():
         for transform in transforms_to_try:
             new_key = transform(key)
@@ -132,21 +120,25 @@ def _align_state_dict(ckpt_sd, model_sd):
                 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
 
 
-# --------------------------------------------------------------------------- #
-# Segmenter
-# --------------------------------------------------------------------------- #
 class CropSegmenter:
     """Loads the fine-tuned SegFormer-B2 model and performs inference."""
 
@@ -155,8 +147,9 @@ class CropSegmenter:
             "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 (small download; weights come from best.pt)
+        # Build architecture from config (weights come from best.pt)
         try:
             seg_config = SegformerConfig.from_pretrained(
                 MODEL_NAME, num_labels=NUM_LABELS
@@ -172,7 +165,6 @@ class CropSegmenter:
                     f"Could not build SegFormer from '{MODEL_NAME}': {exc2}"
                 ) from exc
 
-        # Read the checkpoint (handle both new weights_only default and old files)
         weights_path = _resolve_weights_path()
         try:
             try:
@@ -193,34 +185,42 @@ class CropSegmenter:
         self.model.load_state_dict(aligned, strict=False)
 
         logger.info(
-            "Weights matched %d/%d model tensors (checkpoint had %d).",
+            "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 may be from a different model."
+                "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
-        image resolution.
-
-        input_size: square size used for inference (default INPUT_SIZE).
-        swap_classes: flip class 0<->1 (use if crop was label 0 in training).
+        resolution. Also stores per-call logit stats in self.last_logit_summary.
         """
         try:
             size = int(input_size or INPUT_SIZE)
@@ -228,13 +228,7 @@ class CropSegmenter:
                 image = image.convert("RGB")
             orig_w, orig_h = image.size
 
-            preprocess = transforms.Compose(
-                [
-                    transforms.Resize((size, size)),
-                    transforms.ToTensor(),
-                    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
-                ]
-            )
+            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)
@@ -242,8 +236,17 @@ class CropSegmenter:
                 logits, size=(orig_h, orig_w),
                 mode="bilinear", align_corners=False,
             )
-            mask = upsampled.argmax(dim=1).squeeze(0).to("cpu").numpy().astype(np.uint8)
 
+            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