Update src/ai_processor.py

src/ai_processor.py

@@ -232,117 +232,6 @@ initialize_cpu_models()
 setup_knowledge_base()
 
 # ---------- Calibration helpers ----------
-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 → use the returned scalar threshold (ret), NOT the image
-    ret_otsu, _dst = cv2.threshold(p255, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-    thr_otsu = float(np.clip(ret_otsu / 255.0, 0.25, 0.65))
-
-    # Percentile (90th)
-    thr_pctl = float(np.clip(np.percentile(p01, 90), 0.25, 0.65))
-
-    # Area fraction helper
-    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)
-
-    # Score: prefer ~3–10% coverage
-    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
-
-
-    # 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:
@@ -396,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
@@ -420,25 +308,80 @@ 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]]:
     """
-    TF model → adaptive threshold on prob → (optional) GrabCut grow → cleanup.
-    Falls back to KMeans-Lab when model missing/fails.
+    TF model → adaptive threshold on prob → GrabCut grow → cleanup.
+    Fallback: KMeans-Lab.
     Returns (mask_uint8_0_255, debug_dict)
     """
     debug = {"used": None, "reason": None, "positive_fraction": 0.0,
@@ -454,21 +397,17 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
                 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))
-            rgb_for_view = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
             roi_seen_path = None
             if SMARTHEAL_DEBUG:
                 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))
+                cv2.imwrite(roi_seen_path, image_bgr)
 
-            # 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)
             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)
@@ -476,19 +415,16 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
                 heatmap_path = os.path.join(out_dir, f"seg_pred_heatmap_{ts}.png")
                 cv2.imwrite(heatmap_path, heat)
 
-            # --- 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)
@@ -496,13 +432,13 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
                 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}")
+            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": thr,
+                "thr": float(thr),
                 "heatmap_path": heatmap_path,
                 "roi_seen_by_model": roi_seen_path
             })
@@ -533,7 +469,6 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
     })
     return (mask01 * 255).astype(np.uint8), debug
 
-
 # ---------- Measurement + overlay helpers ----------
 def largest_component_mask(binary01: np.ndarray, min_area_px: int = 50) -> np.ndarray:
     num, labels, stats, _ = cv2.connectedComponentsWithStats(binary01.astype(np.uint8), connectivity=8)
@@ -545,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:
@@ -569,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,
@@ -582,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)
@@ -599,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)
@@ -608,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):
@@ -634,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]])
@@ -663,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 ---
@@ -696,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)
@@ -733,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)
@@ -776,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),
             }
@@ -983,4 +917,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }