Update src/ai_processor.py

src/ai_processor.py

@@ -285,7 +285,6 @@ def estimate_px_per_cm_from_exif(pil_img: Image.Image, default_px_per_cm: float
 
 # ---------- Segmentation helpers ----------
 def _imagenet_norm(arr: np.ndarray) -> np.ndarray:
-    # expects RGB 0..255 -> float
     mean = np.array([123.675, 116.28, 103.53], dtype=np.float32)
     std  = np.array([58.395, 57.12, 57.375], dtype=np.float32)
     return (arr.astype(np.float32) - mean) / std
@@ -309,112 +308,166 @@ def _to_prob(pred: np.ndarray) -> np.ndarray:
         p = 1.0 / (1.0 + np.exp(-p))
     return p.astype(np.float32)
 
-# ---- Robust mask post-processing (for "proper" masking) ----
+# ---- Adaptive threshold + GrabCut grow ----
+def _adaptive_prob_threshold(p: np.ndarray) -> float:
+    """
+    Choose a threshold that avoids tiny blobs while not swallowing skin.
+    Try Otsu and the 90th percentile, clamp to [0.25, 0.65], pick by area heuristic.
+    """
+    p01 = np.clip(p.astype(np.float32), 0, 1)
+    p255 = (p01 * 255).astype(np.uint8)
+
+    ret_otsu, _ = cv2.threshold(p255, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    thr_otsu = float(np.clip(ret_otsu / 255.0, 0.25, 0.65))
+    thr_pctl = float(np.clip(np.percentile(p01, 90), 0.25, 0.65))
+
+    def area_frac(thr: float) -> float:
+        return float((p01 >= thr).sum()) / float(p01.size)
+
+    af_otsu = area_frac(thr_otsu)
+    af_pctl = area_frac(thr_pctl)
+
+    def score(af: float) -> float:
+        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:
+    """Grow from a confident core into low-contrast margins."""
+    h, w = bgr.shape[:2]
+    gc = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
+    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
+    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)
+    return np.where((gc == cv2.GC_FGD) | (gc == cv2.GC_PR_FGD), 1, 0).astype(np.uint8)
+
 def _fill_holes(mask01: np.ndarray) -> np.ndarray:
-    # Flood-fill from border, then invert
     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)
-    # Combine original with filled holes
     out = ((mask01 * 255) | m_inv) // 255
     return out.astype(np.uint8)
 
-# Global last debug dict (per-process) to attach into results
+def _clean_mask(mask01: np.ndarray) -> np.ndarray:
+    """Open → Close → Fill holes → Largest component (no dilation)."""
+    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=1)
+    mask01 = _fill_holes(mask01)
+    # Keep largest component only
+    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)
+    return (mask01 > 0).astype(np.uint8)
+
+# Global last debug dict (per-process)
 _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 → GrabCut grow → cleanup.
+    Fallback: KMeans-Lab.
     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])
+
             x = _preprocess_for_seg(image_bgr, (th, tw))
-            saw_roi = (cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) if SEG_EXPECTS_RGB else image_bgr)
+            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, image_bgr)
 
-            # Inference
             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)
 
+            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,
+            thr = _adaptive_prob_threshold(p)
+            core01 = (p >= thr).astype(np.uint8)
+            core_frac = float(core01.sum()) / float(core01.size)
+
+            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)
+
+            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={float(thr):.2f} core_frac={core_frac:.4f} final_frac={pos_frac:.4f}")
+
+            debug.update({
+                "used": "tf_model",
+                "reason": "ok",
+                "positive_fraction": pos_frac,
+                "thr": float(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:
@@ -427,17 +480,6 @@ def largest_component_mask(binary01: np.ndarray, min_area_px: int = 50) -> np.nd
     largest_idx = 1 + int(np.argmax(areas))
     return (labels == largest_idx).astype(np.uint8)
 
-def _clean_mask(mask01: np.ndarray) -> np.ndarray:
-    """Open→Close→Fill holes→Largest component."""
-    if mask01.dtype != np.uint8:
-        mask01 = mask01.astype(np.uint8)
-    k = np.ones((3, 3), np.uint8)
-    mask01 = cv2.morphologyEx(mask01, cv2.MORPH_OPEN, k, iterations=1)
-    mask01 = cv2.morphologyEx(mask01, cv2.MORPH_CLOSE, k, iterations=2)
-    mask01 = _fill_holes(mask01)
-    mask01 = largest_component_mask(mask01, min_area_px=30)
-    return (mask01 > 0).astype(np.uint8)
-
 def measure_min_area_rect(mask01: np.ndarray, px_per_cm: float) -> Tuple[float, float, Tuple]:
     contours, _ = cv2.findContours(mask01.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if not contours:
@@ -451,9 +493,23 @@ def measure_min_area_rect(mask01: np.ndarray, px_per_cm: float) -> Tuple[float,
     box = cv2.boxPoints(rect).astype(int)
     return length_cm, breadth_cm, (box, rect[0])
 
-def count_area_cm2(mask01: np.ndarray, px_per_cm: float) -> float:
-    px_count = float(mask01.astype(bool).sum())
-    return round(px_count / (max(px_per_cm, 1e-6) ** 2), 2)
+def area_cm2_from_contour(mask01: np.ndarray, px_per_cm: float) -> Tuple[float, Optional[np.ndarray]]:
+    """Area from largest polygon (sub-pixel); returns (area_cm2, contour)."""
+    m = (mask01 > 0).astype(np.uint8)
+    contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return 0.0, None
+    cnt = max(contours, key=cv2.contourArea)
+    poly_area_px2 = float(cv2.contourArea(cnt))
+    area_cm2 = round(poly_area_px2 / (max(px_per_cm, 1e-6) ** 2), 2)
+    return area_cm2, cnt
+
+def clamp_area_with_minrect(cnt: np.ndarray, px_per_cm: float, area_cm2_poly: float) -> float:
+    rect = cv2.minAreaRect(cnt)
+    (w_px, h_px) = rect[1]
+    rect_area_px2 = float(max(w_px, 0.0) * max(h_px, 0.0))
+    rect_area_cm2 = rect_area_px2 / (max(px_per_cm, 1e-6) ** 2)
+    return round(min(area_cm2_poly, rect_area_cm2 * 1.05), 2)
 
 def draw_measurement_overlay(
     base_bgr: np.ndarray,
@@ -464,16 +520,13 @@ def draw_measurement_overlay(
     thickness: int = 2
 ) -> np.ndarray:
     """
-    Draws:
-      1) Strong red mask overlay with white contour.
-      2) Min-area rectangle.
-      3) Two double-headed arrows:
-           - 'Length' along the longer side.
-           - 'Width'  along the shorter side.
+    1) Strong red mask overlay + white contour
+    2) Min-area rectangle
+    3) Double-headed arrows labeled Length/Width
     """
     overlay = base_bgr.copy()
 
-    # --- Strong overlay from mask (tinted red where mask==1) ---
+    # Mask tint
     mask255 = (mask01 * 255).astype(np.uint8)
     mask3 = cv2.merge([mask255, mask255, mask255])
     red = np.zeros_like(overlay); red[:] = (0, 0, 255)
@@ -481,7 +534,7 @@ def draw_measurement_overlay(
     tinted = cv2.addWeighted(overlay, 1 - alpha, red, alpha, 0)
     overlay = np.where(mask3 > 0, tinted, overlay)
 
-    # Draw wound contour
+    # Contour
     cnts, _ = cv2.findContours(mask255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if cnts:
         cv2.drawContours(overlay, cnts, -1, (255, 255, 255), 2)
@@ -490,19 +543,11 @@ def draw_measurement_overlay(
         cv2.polylines(overlay, [rect_box], True, (255, 255, 255), thickness)
         pts = rect_box.reshape(-1, 2)
 
-        def midpoint(a, b):
-            return (int((a[0] + b[0]) / 2), int((a[1] + b[1]) / 2))
-
-        # Edge lengths
+        def midpoint(a, b): return (int((a[0] + b[0]) / 2), int((a[1] + b[1]) / 2))
         e = [np.linalg.norm(pts[i] - pts[(i + 1) % 4]) for i in range(4)]
         long_edge_idx = int(np.argmax(e))
-        short_edge_idx = (long_edge_idx + 1) % 2  # 0/1 map for pairs below
-
-        # Midpoints of opposite edges for arrows
         mids = [midpoint(pts[i], pts[(i + 1) % 4]) for i in range(4)]
-        # Long side uses edges long_edge_idx and the opposite edge (i+2)
         long_pair = (long_edge_idx, (long_edge_idx + 2) % 4)
-        # Short side uses the other pair
         short_pair = ((long_edge_idx + 1) % 4, (long_edge_idx + 3) % 4)
 
         def draw_double_arrow(img, p1, p2):
@@ -516,7 +561,6 @@ def draw_measurement_overlay(
             cv2.putText(overlay, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 4, cv2.LINE_AA)
             cv2.putText(overlay, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
 
-        # Draw arrows and labels
         draw_double_arrow(overlay, mids[long_pair[0]], mids[long_pair[1]])
         draw_double_arrow(overlay, mids[short_pair[0]], mids[short_pair[1]])
         put_label(f"Length: {length_cm:.2f} cm", mids[long_pair[0]])
@@ -545,6 +589,11 @@ class AIProcessor:
         """
         try:
             px_per_cm, exif_meta = estimate_px_per_cm_from_exif(image_pil, DEFAULT_PX_PER_CM)
+            # Guardrails for calibration to avoid huge area blow-ups
+            px_per_cm = float(np.clip(px_per_cm, 20.0, 350.0))
+            if (exif_meta or {}).get("used") != "exif":
+                logging.warning(f"Calibration fallback used: px_per_cm={px_per_cm:.2f} (default). Prefer ruler/Aruco for accuracy.")
+
             image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
 
             # --- Detection ---
@@ -578,20 +627,23 @@ class AIProcessor:
             mask_u8_255, seg_debug = segment_wound(roi, ts, out_dir)
             mask01 = (mask_u8_255 > 127).astype(np.uint8)
 
-            # Robust post-processing to ensure "proper" masking
             if mask01.any():
                 mask01 = _clean_mask(mask01)
                 logging.debug(f"Mask postproc: px_after={int(mask01.sum())}")
 
-            # --- Measurement ---
+            # --- Measurement (accurate & conservative) ---
             if mask01.any():
                 length_cm, breadth_cm, (box_pts, _) = measure_min_area_rect(mask01, px_per_cm)
-                surface_area_cm2 = count_area_cm2(mask01, px_per_cm)
-                # Final annotated ROI with mask + arrows + labels
+                area_poly_cm2, largest_cnt = area_cm2_from_contour(mask01, px_per_cm)
+                if largest_cnt is not None:
+                    surface_area_cm2 = clamp_area_with_minrect(largest_cnt, px_per_cm, area_poly_cm2)
+                else:
+                    surface_area_cm2 = area_poly_cm2
+
                 anno_roi = draw_measurement_overlay(roi, mask01, box_pts, length_cm, breadth_cm)
                 segmentation_empty = False
             else:
-                # Graceful fallback if seg failed: use ROI box as bounds
+                # Fallback if seg failed: use ROI dimensions
                 h_px = max(0, y2 - y1); w_px = max(0, x2 - x1)
                 length_cm = round(max(h_px, w_px) / px_per_cm, 2)
                 breadth_cm = round(min(h_px, w_px) / px_per_cm, 2)
@@ -615,7 +667,7 @@ class AIProcessor:
             roi_mask_path = os.path.join(out_dir, f"roi_mask_{ts}.png")
             cv2.imwrite(roi_mask_path, (mask01 * 255).astype(np.uint8))
 
-            # ROI overlay (clear mask w/ white contour, no arrows)
+            # ROI overlay (mask tint + contour, without arrows)
             mask255 = (mask01 * 255).astype(np.uint8)
             mask3   = cv2.merge([mask255, mask255, mask255])
             red     = np.zeros_like(roi); red[:] = (0, 0, 255)
@@ -658,7 +710,7 @@ class AIProcessor:
                 "seg_used": seg_debug.get("used"),
                 "seg_reason": seg_debug.get("reason"),
                 "positive_fraction": round(float(seg_debug.get("positive_fraction", 0.0)), 6),
-                "threshold": seg_debug.get("threshold", SEG_THRESH),
+                "threshold": seg_debug.get("thr"),
                 "segmentation_empty": segmentation_empty,
                 "exif_px_per_cm": round(px_per_cm, 3),
             }
@@ -674,7 +726,7 @@ class AIProcessor:
                 "detection_confidence": float(results[0].boxes.conf[0].cpu().item())
                     if getattr(results[0].boxes, "conf", None) is not None else 0.0,
                 "detection_image_path": detection_path,
-                "segmentation_image_path": segmentation_path,
+                "segmentation_image_path": annotated_seg_path,
                 "segmentation_annotated_path": annotated_seg_path,
                 "segmentation_roi_path": segmentation_roi_path,
                 "roi_mask_path": roi_mask_path,
@@ -857,4 +909,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }