Update src/ai_processor.py

src/ai_processor.py

@@ -292,29 +292,35 @@ def _imagenet_norm(arr: np.ndarray) -> np.ndarray:
 
 def _preprocess_for_seg(bgr_roi: np.ndarray, target_hw: Tuple[int, int]) -> np.ndarray:
     H, W = target_hw
-    # Resize first
     resized = cv2.resize(bgr_roi, (W, H), interpolation=cv2.INTER_LINEAR)
-    # Convert to RGB if required
     if SEG_EXPECTS_RGB:
         resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
-    # Normalize
     if SEG_NORM.lower() == "imagenet":
         x = _imagenet_norm(resized)
     else:
         x = resized.astype(np.float32) / 255.0
-    # Add batch dim
     x = np.expand_dims(x, axis=0)  # (1,H,W,3)
     return x
 
 def _to_prob(pred: np.ndarray) -> np.ndarray:
-    # Pred could be (1,H,W,1), (H,W,1), (1,H,W), (H,W), or logits
     p = np.squeeze(pred)
-    # If values look like logits, apply sigmoid
     pmin, pmax = float(p.min()), float(p.max())
     if pmax > 1.0 or pmin < 0.0:
         p = 1.0 / (1.0 + np.exp(-p))
     return p.astype(np.float32)
 
+# ---- Robust mask post-processing (for "proper" masking) ----
+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
 _last_seg_debug: Dict[str, object] = {}
 
@@ -356,7 +362,6 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
             _log_kv("SEG_PROB_STATS", {"min": pmin, "max": pmax, "mean": pmean})
 
             if SMARTHEAL_DEBUG:
-                # save heatmap (0..255)
                 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")
@@ -364,8 +369,8 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
 
             # Threshold
             thr = SEG_THRESH
-            mask = (p >= thr).astype(np.uint8) * 255
-            pos = int((mask > 0).sum())
+            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}")
 
@@ -382,7 +387,7 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
                 "heatmap_path": heatmap_path,
                 "roi_seen_by_model": saw_roi_path,
             }
-            return mask.astype(np.uint8), _last_seg_debug
+            return (mask * 255).astype(np.uint8), _last_seg_debug
 
         except Exception as e:
             reason = f"model_failed: {e}"
@@ -395,9 +400,9 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
     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) * 255
+    mask = (labels.reshape(image_bgr.shape[:2]) == wound_idx).astype(np.uint8)
 
-    pos = int((mask > 0).sum()); frac = pos / float(mask.size)
+    pos = int(mask.sum()); frac = pos / float(mask.size)
     logging.info(f"KMeans USED | pos_px={pos} pos_frac={frac:.4f}")
 
     _last_seg_debug = {
@@ -409,7 +414,7 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
         "heatmap_path": heatmap_path,
         "roi_seen_by_model": saw_roi_path,
     }
-    return mask.astype(np.uint8), _last_seg_debug
+    return (mask * 255).astype(np.uint8), _last_seg_debug
 
 # ---------- Measurement + overlay helpers ----------
 def largest_component_mask(binary01: np.ndarray, min_area_px: int = 50) -> np.ndarray:
@@ -422,6 +427,17 @@ 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:
@@ -447,45 +463,65 @@ def draw_measurement_overlay(
     breadth_cm: float,
     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.
+    """
     overlay = base_bgr.copy()
 
-    # Strong overlay + contour
+    # --- Strong overlay from mask (tinted red where mask==1) ---
+    mask255 = (mask01 * 255).astype(np.uint8)
+    mask3 = cv2.merge([mask255, mask255, mask255])
     red = np.zeros_like(overlay); red[:] = (0, 0, 255)
     alpha = 0.55
     tinted = cv2.addWeighted(overlay, 1 - alpha, red, alpha, 0)
-    m3 = cv2.merge([mask01 * 255] * 3).astype("uint8")
-    overlay = np.where(m3 > 0, tinted, overlay)
+    overlay = np.where(mask3 > 0, tinted, overlay)
 
-    # Draw contour
-    cnts, _ = cv2.findContours((mask01 * 255).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # Draw wound contour
+    cnts, _ = cv2.findContours(mask255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if cnts:
         cv2.drawContours(overlay, cnts, -1, (255, 255, 255), 2)
 
     if rect_box is not None:
         cv2.polylines(overlay, [rect_box], True, (255, 255, 255), thickness)
         pts = rect_box.reshape(-1, 2)
-        def midpoint(a, b): return ((a[0] + b[0]) // 2, (a[1] + b[1]) // 2)
-        mids = [midpoint(pts[i], pts[(i+1) % 4]) for i in range(4)]
-        e_lens = [np.linalg.norm(pts[i] - pts[(i+1) % 4]) for i in range(4)]
-        long_pair = (0, 2) if e_lens[0] + e_lens[2] >= e_lens[1] + e_lens[3] else (1, 3)
-        short_pair = (1, 3) if long_pair == (0, 2) else (0, 2)
 
-        def draw_arrow(img, p1, p2):
+        def midpoint(a, b):
+            return (int((a[0] + b[0]) / 2), int((a[1] + b[1]) / 2))
+
+        # Edge lengths
+        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):
             cv2.arrowedLine(img, p1, p2, (0, 0, 0), thickness + 2, tipLength=0.05)
             cv2.arrowedLine(img, p2, p1, (0, 0, 0), thickness + 2, tipLength=0.05)
             cv2.arrowedLine(img, p1, p2, (255, 255, 255), thickness, tipLength=0.05)
             cv2.arrowedLine(img, p2, p1, (255, 255, 255), thickness, tipLength=0.05)
 
-        def put_label(text, org):
-            cv2.putText(overlay, text, (org[0] + 4, org[1] - 4),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 4, cv2.LINE_AA)
-            cv2.putText(overlay, text, (org[0] + 4, org[1] - 4),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
+        def put_label(text, anchor):
+            org = (anchor[0] + 6, anchor[1] - 6)
+            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_arrow(overlay, mids[long_pair[0]], mids[long_pair[1]])
-        draw_arrow(overlay, mids[short_pair[0]], mids[short_pair[1]])
+        # 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]])
-        put_label(f"Breadth: {breadth_cm:.2f} cm", mids[short_pair[0]])
+        put_label(f"Width:  {breadth_cm:.2f} cm", mids[short_pair[0]])
+
     return overlay
 
 # ---------- AI PROCESSOR ----------
@@ -504,7 +540,7 @@ class AIProcessor:
 
     def perform_visual_analysis(self, image_pil: Image.Image) -> Dict:
         """
-        YOLO detect → crop ROI → segment_wound(ROI) → largest component →
+        YOLO detect → crop ROI → segment_wound(ROI) → clean mask →
         minAreaRect measurement (cm) using EXIF px/cm → save outputs.
         """
         try:
@@ -542,24 +578,23 @@ class AIProcessor:
             mask_u8_255, seg_debug = segment_wound(roi, ts, out_dir)
             mask01 = (mask_u8_255 > 127).astype(np.uint8)
 
-            # Post-processing + metrics
+            # Robust post-processing to ensure "proper" masking
             if mask01.any():
-                mask_before = mask01.sum()
-                mask01 = cv2.morphologyEx(mask01, cv2.MORPH_OPEN, np.ones((3,3), np.uint8), iterations=1)
-                mask01 = cv2.morphologyEx(mask01, cv2.MORPH_CLOSE, np.ones((3,3), np.uint8), iterations=1)
-                mask01 = largest_component_mask(mask01, min_area_px=30)
-                logging.debug(f"Mask postproc: px_before={mask_before} px_after={int(mask01.sum())}")
+                mask01 = _clean_mask(mask01)
+                logging.debug(f"Mask postproc: px_after={int(mask01.sum())}")
 
             # --- Measurement ---
             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
                 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
                 h_px = max(0, y2 - y1); w_px = max(0, x2 - x1)
-                length_cm = round(h_px / px_per_cm, 2)
-                breadth_cm = round(w_px / px_per_cm, 2)
+                length_cm = round(max(h_px, w_px) / px_per_cm, 2)
+                breadth_cm = round(min(h_px, w_px) / px_per_cm, 2)
                 surface_area_cm2 = round((h_px * w_px) / (px_per_cm ** 2), 2)
                 anno_roi = roi.copy()
                 cv2.rectangle(anno_roi, (2, 2), (anno_roi.shape[1]-3, anno_roi.shape[0]-3), (0, 0, 255), 3)
@@ -580,7 +615,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 (very clear)
+            # ROI overlay (clear mask w/ white contour, no arrows)
             mask255 = (mask01 * 255).astype(np.uint8)
             mask3   = cv2.merge([mask255, mask255, mask255])
             red     = np.zeros_like(roi); red[:] = (0, 0, 255)
@@ -591,7 +626,7 @@ class AIProcessor:
                 cnts, _ = cv2.findContours(mask255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                 cv2.drawContours(roi_overlay, cnts, -1, (255, 255, 255), 2)
             else:
-                roi_overlay = anno_roi  # already marked X
+                roi_overlay = anno_roi
 
             seg_full = image_cv.copy()
             seg_full[y1:y2, x1:x2] = roi_overlay
@@ -601,6 +636,7 @@ class AIProcessor:
             segmentation_roi_path = os.path.join(out_dir, f"segmentation_roi_{ts}.png")
             cv2.imwrite(segmentation_roi_path, roi_overlay)
 
+            # Annotated (mask + arrows + labels) in full-frame
             anno_full = image_cv.copy()
             anno_full[y1:y2, x1:x2] = anno_roi
             annotated_seg_path = os.path.join(out_dir, f"segmentation_annotated_{ts}.png")