Update src/ai_processor.py

src/ai_processor.py

@@ -232,6 +232,108 @@ initialize_cpu_models()
 setup_knowledge_base()
 
 # ---------- Calibration helpers ----------
+# ---- Adaptive thresholding for model prob map ----
+def _adaptive_prob_threshold(p: np.ndarray) -> float:
+    """
+    Pick a threshold that avoids tiny blobs while not swallowing skin.
+    Strategy:
+      - try Otsu on the prob map
+      - clamp to a reasonable band [0.25, 0.65]
+      - also consider percentile cut (p90) and take the "best" by area heuristic
+    """
+    p01 = np.clip(p.astype(np.float32), 0, 1)
+    p255 = (p01 * 255).astype(np.uint8)
+
+    # Otsu
+    _, thr_otsu = cv2.threshold(p255, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    thr_otsu = np.clip(thr_otsu / 255.0, 0.25, 0.65)
+
+    # Percentile (90th)
+    thr_pctl = float(np.clip(np.percentile(p01, 90), 0.25, 0.65))
+
+    # Prefer the threshold that yields an area fraction in [0.005..0.20]
+    def area_frac(thr):
+        return float((p01 >= thr).sum()) / float(p01.size)
+
+    af_otsu = area_frac(thr_otsu)
+    af_pctl = area_frac(thr_pctl)
+
+    # Score: closeness to a target area fraction (aim ~3–10%)
+    def score(af):
+        target_low, target_high = 0.03, 0.10
+        if af < target_low: return abs(af - target_low) * 3.0
+        if af > target_high: return abs(af - target_high) * 1.5
+        return 0.0
+
+    return thr_otsu if score(af_otsu) <= score(af_pctl) else thr_pctl
+
+
+def _grabcut_refine(bgr: np.ndarray, seed01: np.ndarray, iters: int = 3) -> np.ndarray:
+    """
+    Use OpenCV GrabCut to grow from a confident core into low-contrast margins.
+    seed01: 1=probable FG core, 0=unknown/other
+    """
+    h, w = bgr.shape[:2]
+    # Build GC mask: start with "unknown"
+    gc = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
+    # definite FG = dilated seed; probable FG = seed
+    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    seed_dil = cv2.dilate(seed01, k, iterations=1)
+    gc[seed01.astype(bool)] = cv2.GC_PR_FGD
+    gc[seed_dil.astype(bool)] = cv2.GC_FGD
+    # border is probable background
+    gc[0, :], gc[-1, :], gc[:, 0], gc[:, -1] = cv2.GC_BGD, cv2.GC_BGD, cv2.GC_BGD, cv2.GC_BGD
+
+    bgdModel = np.zeros((1, 65), np.float64)
+    fgdModel = np.zeros((1, 65), np.float64)
+    cv2.grabCut(bgr, gc, None, bgdModel, fgdModel, iters, cv2.GC_INIT_WITH_MASK)
+
+    # FG = definite or probable foreground
+    mask01 = np.where((gc == cv2.GC_FGD) | (gc == cv2.GC_PR_FGD), 1, 0).astype(np.uint8)
+    return mask01
+
+
+def _fill_holes(mask01: np.ndarray) -> np.ndarray:
+    h, w = mask01.shape[:2]
+    ff = np.zeros((h + 2, w + 2), np.uint8)
+    m = (mask01 * 255).astype(np.uint8).copy()
+    cv2.floodFill(m, ff, (0, 0), 255)
+    m_inv = cv2.bitwise_not(m)
+    out = ((mask01 * 255) | m_inv) // 255
+    return out.astype(np.uint8)
+
+
+def _clean_mask(mask01: np.ndarray) -> np.ndarray:
+    """Open → Close → Fill holes → Largest component → light smooth."""
+    mask01 = (mask01 > 0).astype(np.uint8)
+    k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+    k5 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    mask01 = cv2.morphologyEx(mask01, cv2.MORPH_OPEN, k3, iterations=1)
+    mask01 = cv2.morphologyEx(mask01, cv2.MORPH_CLOSE, k5, iterations=2)
+    mask01 = _fill_holes(mask01)
+
+    # keep largest component
+    num, labels, stats, _ = cv2.connectedComponentsWithStats(mask01, 8)
+    if num > 1:
+        areas = stats[1:, cv2.CC_STAT_AREA]
+        if areas.size:
+            largest_idx = 1 + int(np.argmax(areas))
+            mask01 = (labels == largest_idx).astype(np.uint8)
+
+    # tiny masks → gentle grow (distance transform based)
+    area = int(mask01.sum())
+    if area > 0:
+        grow = 1 if area < 2000 else 0
+        if grow:
+            k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+            mask01 = cv2.dilate(mask01, k, iterations=1)
+
+    return (mask01 > 0).astype(np.uint8)
+
+
+
+
+
 def _exif_to_dict(pil_img: Image.Image) -> Dict[str, object]:
     out = {}
     try:
@@ -326,95 +428,102 @@ _last_seg_debug: Dict[str, object] = {}
 
 def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndarray, Dict[str, object]]:
     """
-    Attempts TF segmentation first; falls back to KMeans if needed.
+    TF model → adaptive threshold on prob → (optional) GrabCut grow → cleanup.
+    Falls back to KMeans-Lab when model missing/fails.
     Returns (mask_uint8_0_255, debug_dict)
     """
-    global _last_seg_debug
-    _last_seg_debug = {}
+    debug = {"used": None, "reason": None, "positive_fraction": 0.0,
+             "thr": None, "heatmap_path": None, "roi_seen_by_model": None}
 
     seg_model = models_cache.get("seg", None)
-    used = "fallback_kmeans"
-    reason = "no_model"
-    heatmap_path = None
-    saw_roi_path = None
 
+    # --- Model path ---
     if seg_model is not None:
         try:
             ishape = getattr(seg_model, "input_shape", None)
             if not ishape or len(ishape) < 4:
                 raise ValueError(f"Bad seg input_shape: {ishape}")
             th, tw = int(ishape[1]), int(ishape[2])
+
+            # preprocess
             x = _preprocess_for_seg(image_bgr, (th, tw))
-            saw_roi = (cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) if SEG_EXPECTS_RGB else image_bgr)
+            rgb_for_view = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+            roi_seen_path = None
             if SMARTHEAL_DEBUG:
-                saw_roi_path = os.path.join(out_dir, f"roi_for_seg_{ts}.png")
-                cv2.imwrite(saw_roi_path, (cv2.cvtColor(saw_roi, cv2.COLOR_RGB2BGR) if SEG_EXPECTS_RGB else saw_roi))
+                roi_seen_path = os.path.join(out_dir, f"roi_for_seg_{ts}.png")
+                cv2.imwrite(roi_seen_path, cv2.cvtColor(rgb_for_view, cv2.COLOR_RGB2BGR))
 
-            # Inference
+            # predict → prob map back to ROI size
             pred = seg_model.predict(x, verbose=0)
-            if isinstance(pred, (list, tuple)):
-                pred = pred[0]
-            p = _to_prob(pred)  # HxW
-            p = cv2.resize(p, (image_bgr.shape[1], image_bgr.shape[0]))  # back to ROI size
-
-            # Debug stats
-            pmin, pmax, pmean = float(p.min()), float(p.max()), float(p.mean())
-            _log_kv("SEG_PROB_STATS", {"min": pmin, "max": pmax, "mean": pmean})
+            if isinstance(pred, (list, tuple)): pred = pred[0]
+            p = _to_prob(pred)
+            p = cv2.resize(p, (image_bgr.shape[1], image_bgr.shape[0]), interpolation=cv2.INTER_LINEAR)
 
+            # visualization (optional)
+            heatmap_path = None
             if SMARTHEAL_DEBUG:
                 hm = (np.clip(p, 0, 1) * 255).astype(np.uint8)
                 heat = cv2.applyColorMap(hm, cv2.COLORMAP_JET)
                 heatmap_path = os.path.join(out_dir, f"seg_pred_heatmap_{ts}.png")
                 cv2.imwrite(heatmap_path, heat)
 
-            # Threshold
-            thr = SEG_THRESH
-            mask = (p >= thr).astype(np.uint8)  # 0/1
-            pos = int(mask.sum())
-            frac = pos / float(mask.size)
-            logging.info(f"SegModel USED | thr={thr} pos_px={pos} pos_frac={frac:.4f} ex_rgb={SEG_EXPECTS_RGB} norm={SEG_NORM}")
-
-            used = "tf_model"
-            reason = "ok"
-
-            _last_seg_debug = {
-                "used": used,
-                "reason": reason,
-                "input_shape": ishape,
-                "prob_min": pmin, "prob_max": pmax, "prob_mean": pmean,
-                "threshold": thr,
-                "positive_fraction": frac,
+            # --- Adaptive threshold ---
+            thr = _adaptive_prob_threshold(p)
+            core01 = (p >= thr).astype(np.uint8)
+            core_frac = float(core01.sum()) / float(core01.size)
+
+            # If still too tiny, try a gentler threshold
+            if core_frac < 0.005:
+                thr2 = max(thr - 0.10, 0.15)
+                core01 = (p >= thr2).astype(np.uint8)
+                thr = thr2
+                core_frac = float(core01.sum()) / float(core01.size)
+
+            # --- Grow with GrabCut (only if some core exists) ---
+            if core01.any():
+                gc01 = _grabcut_refine(image_bgr, core01, iters=3)
+                mask01 = _clean_mask(gc01)
+            else:
+                mask01 = np.zeros(core01.shape, np.uint8)
+
+            pos_frac = float(mask01.sum()) / float(mask01.size)
+            logging.info(f"SegModel USED | thr={thr:.2f} core_frac={core_frac:.4f} final_frac={pos_frac:.4f}")
+
+            debug.update({
+                "used": "tf_model",
+                "reason": "ok",
+                "positive_fraction": pos_frac,
+                "thr": thr,
                 "heatmap_path": heatmap_path,
-                "roi_seen_by_model": saw_roi_path,
-            }
-            return (mask * 255).astype(np.uint8), _last_seg_debug
+                "roi_seen_by_model": roi_seen_path
+            })
+            return (mask01 * 255).astype(np.uint8), debug
 
         except Exception as e:
-            reason = f"model_failed: {e}"
-            logging.warning(f"⚠️ Segmentation model prediction failed → fallback. Reason: {e}")
+            logging.warning(f"⚠️ Segmentation model failed → fallback. Reason: {e}")
+            debug.update({"used": "fallback_kmeans", "reason": f"model_failed: {e}"})
 
-    # --- Fallback: KMeans (k=2), pick 'reddest' cluster in Lab a* ---
+    # --- Fallback: KMeans in Lab (reddest cluster as wound) ---
     Z = image_bgr.reshape((-1, 3)).astype(np.float32)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
     _, labels, centers = cv2.kmeans(Z, 2, None, criteria, 5, cv2.KMEANS_PP_CENTERS)
     centers_u8 = centers.astype(np.uint8).reshape(1, 2, 3)
     centers_lab = cv2.cvtColor(centers_u8, cv2.COLOR_BGR2LAB)[0]
-    wound_idx = int(np.argmax(centers_lab[:, 1]))  # maximize a* (redness)
-    mask = (labels.reshape(image_bgr.shape[:2]) == wound_idx).astype(np.uint8)
-
-    pos = int(mask.sum()); frac = pos / float(mask.size)
-    logging.info(f"KMeans USED | pos_px={pos} pos_frac={frac:.4f}")
-
-    _last_seg_debug = {
-        "used": used,
-        "reason": reason,
-        "kmeans_centers_bgr": centers.tolist(),
-        "kmeans_centers_lab": centers_lab.astype(float).tolist(),
-        "positive_fraction": frac,
-        "heatmap_path": heatmap_path,
-        "roi_seen_by_model": saw_roi_path,
-    }
-    return (mask * 255).astype(np.uint8), _last_seg_debug
+    wound_idx = int(np.argmax(centers_lab[:, 1]))  # maximize a* (red)
+    mask01 = (labels.reshape(image_bgr.shape[:2]) == wound_idx).astype(np.uint8)
+    mask01 = _clean_mask(mask01)
+
+    pos_frac = float(mask01.sum()) / float(mask01.size)
+    logging.info(f"KMeans USED | final_frac={pos_frac:.4f}")
+
+    debug.update({
+        "used": "fallback_kmeans",
+        "reason": debug.get("reason") or "no_model",
+        "positive_fraction": pos_frac,
+        "thr": None
+    })
+    return (mask01 * 255).astype(np.uint8), debug
+
 
 # ---------- Measurement + overlay helpers ----------
 def largest_component_mask(binary01: np.ndarray, min_area_px: int = 50) -> np.ndarray: