Update src/ai_processor.py

src/ai_processor.py

@@ -1223,11 +1223,10 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             visual = result.get("visual_analysis", {}) or {}
             px_per_cm = float(visual.get("px_per_cm", DEFAULT_PX_PER_CM))
     
-            # Attempt to load the ROI mask generated by the pipeline
+            # Attempt to load a mask
             roi_mask_path = visual.get("roi_mask_path")
             mask_img = None
     
-            # Use manual mask if provided
             if manual_mask_path:
                 try:
                     if os.path.exists(manual_mask_path):
@@ -1237,91 +1236,140 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 except Exception as e:
                     logging.warning(f"Failed to load manual mask: {e}")
             elif roi_mask_path and os.path.exists(roi_mask_path):
-                # Otherwise load the automatically generated ROI mask
                 try:
                     mask_img = Image.open(roi_mask_path)
                 except Exception as e:
                     logging.warning(f"Failed to load ROI mask for adjustment: {e}")
     
             if mask_img is not None:
-                # Convert to numpy for processing
                 mask_np = np.array(mask_img.convert("L"))
     
-                # If adjustment is requested and no manual mask override
+                # If adjustment requested and no manual override
                 if (manual_mask_path is None) and (abs(seg_adjust) >= 1e-5):
-                    # Determine the number of iterations based on percentage; roughly 5% increments
                     iter_count = max(1, int(round(abs(seg_adjust) / 5)))
                     kernel = np.ones((3, 3), np.uint8)
                     try:
                         if seg_adjust > 0:
-                            # Dilate (expand) mask
                             mask_np = cv2.dilate((mask_np > 127).astype(np.uint8), kernel, iterations=iter_count)
                         else:
-                            # Erode (shrink) mask
                             mask_np = cv2.erode((mask_np > 127).astype(np.uint8), kernel, iterations=iter_count)
                     except Exception as e:
                         logging.warning(f"Segmentation adjustment failed: {e}")
                 else:
-                    # If manual mask provided, binarize it directly
                     mask_np = (mask_np > 127).astype(np.uint8)
     
-                # Recalculate length, width and area using the adjusted or manual mask
+                # Recalculate metrics
                 try:
                     length_cm, breadth_cm, area_cm2 = self._refine_metrics_from_mask(mask_np, px_per_cm)
                     visual["length_cm"] = length_cm
                     visual["breadth_cm"] = breadth_cm
                     visual["surface_area_cm2"] = area_cm2
-                    # Indicate that segmentation was refined manually or adjusted
                     visual["segmentation_refined"] = bool(manual_mask_path) or (abs(seg_adjust) >= 1e-5)
                 except Exception as e:
                     logging.warning(f"Failed to recalculate metrics from mask: {e}")
     
-                # --- NEW: if a manual mask was supplied, create & store a manual overlay and update paths ---
+                # ------- Manual overlay with wound-only red + ARROWS -------
                 if manual_mask_path:
                     try:
-                        # Base image for overlay
                         base_rgb = np.array(image_pil.convert("RGB"))
                         base_bgr = cv2.cvtColor(base_rgb, cv2.COLOR_RGB2BGR)
                         h, w = base_bgr.shape[:2]
     
-                        # Ensure mask matches base size
                         if mask_np.shape[:2] != (h, w):
                             mask_np = cv2.resize(mask_np.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
     
-                        # Decide output directory
-                        out_dir = os.path.dirname(roi_mask_path or result.get("saved_image_path") or manual_mask_path)
-                        if not out_dir or not os.path.exists(out_dir):
-                            out_dir = os.getcwd()
+                        # If mask seems inverted (covers majority), flip it so 1 = wound
+                        wound_fraction = float(mask_np.mean())
+                        if wound_fraction > 0.5:
+                            mask_np = (1 - mask_np).astype(np.uint8)
     
+                        # Output dir
+                        out_dir = os.path.dirname(roi_mask_path or result.get("saved_image_path") or manual_mask_path) or os.getcwd()
                         ts = datetime.now().strftime("%Y%m%d_%H%M%S")
     
-                        # Save a clean binary manual mask (0/255)
+                        # Save clean binary manual mask
                         manual_mask_save = os.path.join(out_dir, f"manual_mask_{ts}.png")
                         cv2.imwrite(manual_mask_save, (mask_np * 255).astype(np.uint8))
     
-                        # Build red overlay with white contour
+                        # Base red overlay on wound only
                         red = np.zeros_like(base_bgr); red[:] = (0, 0, 255)
                         alpha = 0.55
                         tinted = cv2.addWeighted(base_bgr, 1 - alpha, red, alpha, 0)
-                        mask3 = cv2.merge([(mask_np * 255).astype(np.uint8)] * 3)
+                        mask255 = (mask_np * 255).astype(np.uint8)
+                        mask3   = cv2.merge([mask255, mask255, mask255])
                         overlay = np.where(mask3 > 0, tinted, base_bgr)
     
-                        cnts, _ = cv2.findContours((mask_np * 255).astype(np.uint8),
-                                                   cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-                        if cnts:
-                            cv2.drawContours(overlay, cnts, -1, (255, 255, 255), 2)
+                        # ---- Draw double-headed arrows + labels for Length & Width ----
+                        ys, xs = np.where(mask_np > 0)
+                        if xs.size and ys.size:
+                            x0, x1 = int(xs.min()), int(xs.max())
+                            y0, y1 = int(ys.min()), int(ys.max())
+                            w_px = x1 - x0 + 1
+                            h_px = y1 - y0 + 1
     
+                            # Compute cm from px (fallback-safe)
+                            def _px_to_cm(px):
+                                try:
+                                    return float(px) / float(px_per_cm if px_per_cm else DEFAULT_PX_PER_CM)
+                                except Exception:
+                                    return float(px)
+    
+                            L_px  = max(w_px, h_px)
+                            W_px  = min(w_px, h_px)
+                            L_cm  = _px_to_cm(L_px)
+                            W_cm  = _px_to_cm(W_px)
+    
+                            # Horizontal (center y) and vertical (center x) lines
+                            cy = (y0 + y1) // 2
+                            cx = (x0 + x1) // 2
+                            h_start, h_end = (x0, cy), (x1, cy)
+                            v_start, v_end = (cx, y0), (cx, y1)
+    
+                            # Helper: outlined arrowed line (black underlay + white line)
+                            def draw_double_headed(img, p1, p2, color_fg=(255,255,255), color_bg=(0,0,0), t_fg=3, t_bg=6):
+                                cv2.arrowedLine(img, p1, p2, color_bg, t_bg, tipLength=0.03)
+                                cv2.arrowedLine(img, p2, p1, color_bg, t_bg, tipLength=0.03)
+                                cv2.arrowedLine(img, p1, p2, color_fg, t_fg, tipLength=0.03)
+                                cv2.arrowedLine(img, p2, p1, color_fg, t_fg, tipLength=0.03)
+    
+                            # Draw both arrows
+                            draw_double_headed(overlay, h_start, h_end)
+                            draw_double_headed(overlay, v_start, v_end)
+    
+                            # Helper: outlined text
+                            def put_text_outlined(img, text, org, font=cv2.FONT_HERSHEY_SIMPLEX, scale=0.7,
+                                                  color_fg=(255,255,255), color_bg=(0,0,0), t_fg=2, t_bg=4):
+                                cv2.putText(img, text, org, font, scale, color_bg, t_bg, cv2.LINE_AA)
+                                cv2.putText(img, text, org, font, scale, color_fg, t_fg, cv2.LINE_AA)
+    
+                            # Decide which is Length vs Width for labels
+                            if w_px >= h_px:
+                                # horizontal is length
+                                put_text_outlined(overlay, f"Length: {L_cm:.2f} cm",
+                                                  (x0, max(25, cy - 10)))
+                                put_text_outlined(overlay, f"Width: {W_cm:.2f} cm",
+                                                  (max(5, cx + 10), y0 + 25))
+                            else:
+                                # vertical is length
+                                put_text_outlined(overlay, f"Length: {L_cm:.2f} cm",
+                                                  (max(5, cx + 10), cy))
+                                put_text_outlined(overlay, f"Width: {W_cm:.2f} cm",
+                                                  (x0, max(25, cy - 10)))
+    
+                        # Save overlay with arrows
                         manual_overlay_path = os.path.join(out_dir, f"segmentation_manual_{ts}.png")
                         cv2.imwrite(manual_overlay_path, overlay)
     
-                        # Update paths so the UI shows the MANUAL overlay/mask
+                        # Update paths so UI shows the manual overlay (with arrows)
                         visual["roi_mask_path"] = manual_mask_save
                         visual["segmentation_image_path"] = manual_overlay_path
-                        visual["segmentation_roi_path"] = manual_overlay_path  # alias some UIs read
+                        visual["segmentation_roi_path"] = manual_overlay_path
+                        visual["segmentation_annotated_path"] = manual_overlay_path
                         visual["segmentation_refined_type"] = "manual"
                         visual["manual_mask_used"] = True
                     except Exception as e:
                         logging.warning(f"Failed to generate manual segmentation overlay: {e}")
+                # ----------------------------------------------------------
     
             result["visual_analysis"] = visual
             return result