Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import os
 
 # ---- Bind to Spaces port/address & disable usage stats ----
@@ -5,13 +6,18 @@ os.environ["STREAMLIT_SERVER_PORT"] = os.getenv("PORT", "7860")
 os.environ["STREAMLIT_SERVER_ADDRESS"] = "0.0.0.0"
 os.environ.setdefault("STREAMLIT_BROWSER_GATHERUSAGESTATS", "false")
 
-# ---- Use /tmp (always writable) for HOME and caches ----
+# ---- Writable dirs on Spaces ----
 BASE_DIR = os.getenv("APP_WRITE_DIR", "/tmp")
 os.environ.setdefault("HOME", BASE_DIR)
 os.environ.setdefault("XDG_CACHE_HOME", f"{BASE_DIR}/.cache")
 os.environ.setdefault("XDG_CONFIG_HOME", f"{BASE_DIR}/.config")
 os.environ.setdefault("HF_HOME", f"{BASE_DIR}/.cache/huggingface")
-for p in [f"{BASE_DIR}/.cache", f"{BASE_DIR}/.config", f"{BASE_DIR}/.cache/huggingface", f"{BASE_DIR}/.streamlit"]:
+for p in [
+    f"{BASE_DIR}/.cache",
+    f"{BASE_DIR}/.config",
+    f"{BASE_DIR}/.cache/huggingface",
+    f"{BASE_DIR}/.streamlit",
+]:
     os.makedirs(p, exist_ok=True)
 
 import streamlit as st
@@ -20,7 +26,7 @@ st.set_page_config(
     layout="wide",
 )
 
-import sys, gc
+import gc
 import torch
 import numpy as np
 from PIL import Image, ImageFilter
@@ -31,7 +37,7 @@ from transformers import (
 )
 from huggingface_hub import hf_hub_download
 
-# ---- Make CPU predictable/light ----
+# ---- Keep CPU predictable/light ----
 try:
     torch.set_num_threads(1)
     torch.set_num_interop_threads(1)
@@ -87,17 +93,17 @@ if os.path.exists(logo_path):
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Preferred class IDs (fallback to auto-detect if no segments match)
-PREF_CANOPY_ID   = os.getenv("CANOPY_CLASS_ID")
-PREF_DISEASE_ID  = os.getenv("DISEASE_CLASS_ID")
-PREF_CANOPY_ID   = int(PREF_CANOPY_ID)  if PREF_CANOPY_ID  is not None else None
-PREF_DISEASE_ID  = int(PREF_DISEASE_ID) if PREF_DISEASE_ID is not None else None
+PREF_CANOPY_ID  = os.getenv("CANOPY_CLASS_ID")
+PREF_DISEASE_ID = os.getenv("DISEASE_CLASS_ID")
+PREF_CANOPY_ID  = int(PREF_CANOPY_ID)  if PREF_CANOPY_ID  is not None else None
+PREF_DISEASE_ID = int(PREF_DISEASE_ID) if PREF_DISEASE_ID is not None else None
 
 # ---- Model loader helpers ----
 def _load_state(weights_path: str):
     if weights_path.endswith(".safetensors"):
         from safetensors.torch import load_file as safe_load_file
         return safe_load_file(weights_path)
+    # .bin
     return torch.load(weights_path, map_location="cpu", weights_only=True)
 
 def _load_maskformer(config_path: str, weights_path: str):
@@ -122,7 +128,7 @@ def load_segmentation_models():
 canopy_model, disease_model, processor, SEG_LOAD_ERR = load_segmentation_models()
 SEG_READY = (canopy_model is not None) and (disease_model is not None)
 
-# ---- Lazy GPT-2 generator (optional) ----
+# ---- Optional text generator (kept lazy) ----
 @st.cache_resource
 def get_text_generator():
     try:
@@ -131,7 +137,7 @@ def get_text_generator():
     except Exception as e:
         return None, f"{type(e).__name__}: {e}"
 
-# ---- Inference helpers ----
+# ---- Inference & mask helpers ----
 @torch.inference_mode()
 def run_inference(model, image):
     inputs = processor(image, return_tensors="pt", do_resize=False)
@@ -144,18 +150,17 @@ def run_inference(model, image):
         mask_threshold=0.5,
         overlap_mask_area_threshold=0.5,
     )[0]
-    return results  # dict with 'segmentation' and 'segments_info'
+    return results
 
 def _segmap_to_numpy(seg_map):
-    """Ensure segmentation map is NumPy."""
     if isinstance(seg_map, torch.Tensor):
         return seg_map.detach().cpu().numpy()
     return np.array(seg_map)
 
-def union_mask_with_fallback(results, preferred_label_id=None):
+def union_mask_with_fallback(results, preferred_label_id=None, prefer_smallest=False):
     """
-    Build a boolean mask. Try preferred_label_id first; if it yields no pixels,
-    choose the label_id with the largest total area.
+    Build a boolean mask. If preferred_label_id is provided and present, use it.
+    Otherwise choose by area: largest by default, or smallest if prefer_smallest=True.
     Returns (mask_bool, chosen_label_id).
     """
     seg_map = _segmap_to_numpy(results["segmentation"])
@@ -172,13 +177,13 @@ def union_mask_with_fallback(results, preferred_label_id=None):
                 hit = True
         return m, hit
 
-    # 1) Try preferred id (if provided)
+    # Preferred first
     if preferred_label_id is not None:
         m, hit = union_for(preferred_label_id)
         if hit:
             return m, int(preferred_label_id)
 
-    # 2) Auto-pick: label with largest area
+    # Compute area per label_id
     area_per_label = {}
     for info in infos:
         sid = int(info["id"])
@@ -186,80 +191,111 @@ def union_mask_with_fallback(results, preferred_label_id=None):
         area_per_label[lid] = area_per_label.get(lid, 0) + int((seg_map == sid).sum())
     if not area_per_label:
         return np.zeros(seg_map.shape, dtype=bool), None
-    best_lid = max(area_per_label, key=area_per_label.get)
+
+    # Choose by area
+    best_lid = (min if prefer_smallest else max)(area_per_label, key=area_per_label.get)
     m, _ = union_for(best_lid)
     return m, int(best_lid)
 
-# ---------- Soft, true-shape overlay (no square boxes) ----------
-def soft_overlay(image_np: np.ndarray,
-                 canopy_mask: np.ndarray,
-                 disease_mask: np.ndarray,
-                 alpha_canopy: float = 0.30,
-                 alpha_disease: float = 0.50,
-                 sigma: float = 1.8) -> Image.Image:
-    """
-    Feathered compositing with no outline, so shapes follow the mask exactly.
-    """
-    out = image_np.astype(np.float32)
-
-    def _mask_to_soft_alpha(mask_bool: np.ndarray, sigma_val: float) -> np.ndarray:
-        if not mask_bool.any():
-            return np.zeros_like(mask_bool, dtype=np.float32)
-        m = Image.fromarray((mask_bool.astype(np.uint8) * 255))
-        m = m.filter(ImageFilter.GaussianBlur(radius=max(0.1, float(sigma_val))))
-        return (np.asarray(m, dtype=np.float32) / 255.0)
-
-    # Canopy first (orange), then disease (blue) on top
-    ca = _mask_to_soft_alpha(canopy_mask, sigma) * alpha_canopy
-    if ca.any():
-        for ch, col in enumerate([255, 165, 0]):
-            out[..., ch] = out[..., ch] * (1.0 - ca) + col * ca
-
-    da = _mask_to_soft_alpha(disease_mask, sigma) * alpha_disease
-    if da.any():
-        for ch, col in enumerate([0, 0, 255]):
-            out[..., ch] = out[..., ch] * (1.0 - da) + col * da
-
-    out = np.clip(out, 0, 255).astype(np.uint8)
-    return Image.fromarray(out)
-# ---------------------------------------------------------------
+def smooth_bool_mask(mask_bool, radius=1.1):
+    """Slight Gaussian blur + threshold to de-jag mask edges."""
+    if mask_bool.dtype != np.bool_:
+        mask_bool = mask_bool.astype(bool)
+    im = Image.fromarray((mask_bool * 255).astype(np.uint8))
+    im = im.filter(ImageFilter.GaussianBlur(radius=radius))
+    arr = np.array(im) > 127
+    return arr
+
+def erode_4n(mask):
+    """Simple 4-neighbour erosion (approx) without SciPy."""
+    up    = np.roll(mask, -1, axis=0)
+    down  = np.roll(mask,  1, axis=0)
+    left  = np.roll(mask, -1, axis=1)
+    right = np.roll(mask,  1, axis=1)
+    return mask & up & down & left & right
+
+def mask_outline(mask):
+    """1px outline from a boolean mask."""
+    er = erode_4n(mask)
+    return mask & (~er)
+
+def to_color(mask_bool, image_np, color):
+    rgb = np.zeros_like(image_np)
+    rgb[mask_bool] = color
+    return rgb
+
+def create_overlay(image_np, canopy_bool, disease_bool, alpha_canopy=0.35, alpha_disease=0.45):
+    canopy_rgb  = to_color(canopy_bool,  image_np, [255, 165,   0])  # orange
+    disease_rgb = to_color(disease_bool, image_np, [  0,   0, 255])  # blue
+
+    # Blend canopy first, then disease (so disease remains visible over canopy)
+    overlay1 = (image_np * (1 - alpha_canopy) + canopy_rgb * alpha_canopy).astype(np.uint8)
+    overlay2 = (overlay1  * (1 - alpha_disease) + disease_rgb * alpha_disease).astype(np.uint8)
+
+    # Thin outline around disease to improve visibility
+    edges = mask_outline(disease_bool)
+    overlay2[edges] = [ 10,  74, 255]  
+
+    return Image.fromarray(overlay2)
 
 def process_image(image: Image.Image):
-    # Slightly larger default improves mask smoothness
-    MAX_SIDE = int(os.getenv("MAX_SIDE", "640"))  # was 512
+    MAX_SIDE = int(os.getenv("MAX_SIDE", "768"))
     im = image.copy()
     im.thumbnail((MAX_SIDE, MAX_SIDE), Image.LANCZOS)
-
     image_np = np.array(im)
 
     # Run BOTH models
     canopy_results  = run_inference(canopy_model,  im)
     disease_results = run_inference(disease_model, im)
 
-    # Build masks, using preferred label ids, with fallbacks to largest-area label
-    canopy_mask_bool, chosen_canopy_id   = union_mask_with_fallback(canopy_results,  PREF_CANOPY_ID)
-    disease_mask_bool, chosen_disease_id = union_mask_with_fallback(disease_results, PREF_DISEASE_ID)
+    # Canopy: prefer pinned id or (fallback) largest area
+    canopy_mask, chosen_canopy_id = union_mask_with_fallback(
+        canopy_results, PREF_CANOPY_ID, prefer_smallest=False
+    )
+    canopy_mask = smooth_bool_mask(canopy_mask, radius=1.0)
+
+    # Disease: prefer pinned id or (fallback) smallest area (spots)
+    disease_mask, chosen_disease_id = union_mask_with_fallback(
+        disease_results, PREF_DISEASE_ID, prefer_smallest=True
+    )
+    disease_mask = smooth_bool_mask(disease_mask, radius=0.8)
 
     # If canopy not found, bail gracefully
-    if canopy_mask_bool.sum() == 0:
+    if canopy_mask.sum() == 0:
         return im, Image.fromarray(image_np), 0.0
 
-    # Only count disease inside canopy
-    disease_in_canopy = disease_mask_bool & canopy_mask_bool
-    canopy_pixels  = int(canopy_mask_bool.sum())
+    # Only count disease INSIDE canopy
+    disease_in_canopy = disease_mask & canopy_mask
+
+    frac = disease_in_canopy.sum() / max(1, canopy_mask.sum())
+    if frac > 0.6:
+        # Re-pick a smaller label (2nd smallest by area)
+        seg_map = _segmap_to_numpy(disease_results["segmentation"])
+        infos   = disease_results.get("segments_info", [])
+        area_per_label = {}
+        for info in infos:
+            sid = int(info["id"]); lid = int(info["label_id"])
+            area_per_label.setdefault(lid, 0)
+            area_per_label[lid] += int((seg_map == sid).sum())
+        if len(area_per_label) > 1:
+            small_to_large = sorted(area_per_label.items(), key=lambda kv: kv[1])
+            for lid, _ in small_to_large:
+                if PREF_DISEASE_ID is not None and lid == PREF_DISEASE_ID:
+                    continue
+                m = np.zeros_like(seg_map, dtype=bool)
+                for info in infos:
+                    if int(info["label_id"]) == lid:
+                        m |= (seg_map == int(info["id"]))
+                if m.sum() > 0:
+                    disease_mask = smooth_bool_mask(m, radius=0.8)
+                    disease_in_canopy = disease_mask & canopy_mask
+                    break
+
+    canopy_pixels  = int(canopy_mask.sum())
     disease_pixels = int(disease_in_canopy.sum())
     disease_pct = round((disease_pixels / canopy_pixels) * 100, 2) if canopy_pixels > 0 else 0.0
 
-    # Soft, more transparent overlay 
-    overlay = soft_overlay(
-        image_np,
-        canopy_mask=canopy_mask_bool,
-        disease_mask=disease_in_canopy,
-        alpha_canopy=0.30,
-        alpha_disease=0.50,
-        sigma=1.8
-    )
-
+    overlay = create_overlay(image_np, canopy_mask, disease_in_canopy)
     return im, overlay, disease_pct
 
 # ---- UI ----
@@ -284,7 +320,6 @@ user_description = st.sidebar.text_area(
 )
 
 if uploaded_file:
-    from PIL import Image  # local import to keep startup snappy
     image = Image.open(uploaded_file).convert("RGB")
     st.write("**Running Segmentation..**")
 
@@ -297,9 +332,9 @@ if uploaded_file:
             small, overlay_image, disease_percentage = process_image(image)
             col1, col2 = st.columns([1, 1])
             with col1:
-                st.image(small, caption="Original Image", width=350)
+                st.image(small, caption="Original Image", width=450)
             with col2:
-                st.image(overlay_image, caption="Segmented Image (canopy=orange, disease=blue)", width=350)
+                st.image(overlay_image, caption="Segmented Image (canopy=orange, disease=blue)", width=450)
         except Exception as e:
             st.warning(f"Segmentation failed gracefully: {type(e).__name__}: {e}")
             st.image(image, caption="Original Image", width=350)