Update src/ai_processor.py

src/ai_processor.py

@@ -1,22 +1,20 @@
 # smartheal_ai_processor.py
-# Full, functional module with an always-present @spaces.GPU function (if `spaces` is importable)
-# and robust CPU fallbacks to avoid crashes when GPU isn't actually available yet.
-# + Automatic calibration (px/cm) and measurement overlay on segmentation.
+# Fully functional: auto-calibration from EXIF, mask-based measurements,
+# and annotated overlay with arrows+labels.
 
 import os
 import time
 import logging
 from datetime import datetime
-from typing import Optional, Dict, List, Tuple, Union
+from typing import Optional, Dict, List, Tuple
 
 import cv2
 import numpy as np
-from PIL import Image, TiffImagePlugin
+from PIL import Image, ImageOps
+from PIL.ExifTags import TAGS
 
-# =============== LOGGING ===============
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
-# =============== CONFIG ===============
 UPLOADS_DIR = "uploads"
 os.makedirs(UPLOADS_DIR, exist_ok=True)
 
@@ -24,15 +22,14 @@ HF_TOKEN = os.getenv("HF_TOKEN", None)
 YOLO_MODEL_PATH = "src/best.pt"
 SEG_MODEL_PATH = "src/segmentation_model.h5"   # optional
 GUIDELINE_PDFS = ["src/eHealth in Wound Care.pdf", "src/IWGDF Guideline.pdf", "src/evaluation.pdf"]
-DATASET_ID = "SmartHeal/wound-image-uploads"   # optional (requires HF_TOKEN)
-# Fallback px/cm if we cannot calibrate from EXIF
-DEFAULT_PIXELS_PER_CM = 38.0
+DATASET_ID = "SmartHeal/wound-image-uploads"
+DEFAULT_PX_PER_CM = 38.0  # fallback when we cannot calibrate
+PX_PER_CM_MIN, PX_PER_CM_MAX = 5.0, 1200.0  # sanity bounds
 
-# =============== CACHES ===============
 models_cache: Dict[str, object] = {}
 knowledge_base_cache: Dict[str, object] = {}
 
-# =============== Optional imports (lazy) ===============
+# ---------- Lazy imports ----------
 def _import_ultralytics():
     from ultralytics import YOLO
     return YOLO
@@ -63,7 +60,7 @@ def _import_hf_hub():
     from huggingface_hub import HfApi, HfFolder
     return HfApi, HfFolder
 
-# =============== Spaces GPU function (always defined if `spaces` import works) ===============
+# ---------- Spaces GPU function (always defined if `spaces` import works) ----------
 try:
     import spaces
 
@@ -75,39 +72,29 @@ try:
         image_pil: Image.Image,
         max_new_tokens: Optional[int] = None,
     ) -> str:
-        """
-        This function MUST exist at import time so Spaces Zero detects it.
-        It is guarded internally so if anything fails (no GPU yet, model load error),
-        it returns a warning and your pipeline will use the fallback report.
-        """
         try:
             import torch
             from transformers import pipeline
 
-            # Try to free cache; if no CUDA, this will raise and we return a warning.
             try:
                 if hasattr(torch, "cuda") and torch.cuda.is_available():
                     torch.cuda.empty_cache()
             except Exception:
                 pass
 
-            prompt = f"""
-You are a medical AI assistant. Analyze this wound image and patient data.
-
-Patient: {patient_info}
-Wound: {visual_results.get('wound_type', 'Unknown')} - {visual_results.get('length_cm', 0)}×{visual_results.get('breadth_cm', 0)} cm
-
-Provide a structured report with:
-1. Clinical Summary
-2. Treatment Recommendations
-3. Risk Assessment
-4. Monitoring Plan
-""".strip()
+            prompt = (
+                "You are a medical AI assistant. Analyze this wound image and patient data.\n\n"
+                f"Patient: {patient_info}\n"
+                f"Wound: {visual_results.get('wound_type', 'Unknown')} - "
+                f"{visual_results.get('length_cm', 0)}×{visual_results.get('breadth_cm', 0)} cm\n\n"
+                "Provide a structured report with:\n"
+                "1. Clinical Summary\n2. Treatment Recommendations\n3. Risk Assessment\n4. Monitoring Plan\n"
+            )
 
+            from transformers import pipeline
             pipe = pipeline(
                 "image-text-to-text",
                 model="google/medgemma-4b-it",
-                torch_dtype=getattr(torch, "bfloat16", None),
                 device_map="auto",
                 token=HF_TOKEN,
                 model_kwargs={"low_cpu_mem_usage": True, "use_cache": True},
@@ -129,7 +116,6 @@ Provide a structured report with:
             logging.info(f"✅ MedGemma finished in {time.time()-t0:.2f}s")
 
             if out and len(out) > 0:
-                # Defensive extraction (different transformers versions)
                 try:
                     return out[0]["generated_text"][-1].get("content", "").strip() or "⚠️ Empty response"
                 except Exception:
@@ -139,7 +125,6 @@ Provide a structured report with:
             logging.error(f"❌ MedGemma generation error: {e}")
             return "⚠️ GPU worker unavailable"
 except Exception:
-    # If `spaces` cannot be imported locally, expose a CPU-safe stub with same signature.
     def generate_medgemma_report(
         patient_info: str,
         visual_results: Dict,
@@ -149,7 +134,7 @@ except Exception:
     ) -> str:
         return "⚠️ GPU not available"
 
-# =============== Model init (CPU-safe) ===============
+# ---------- Initialize CPU models ----------
 def load_yolo_model():
     YOLO = _import_ultralytics()
     return YOLO(YOLO_MODEL_PATH)
@@ -213,7 +198,6 @@ def initialize_cpu_models() -> None:
 def setup_knowledge_base() -> None:
     if "vector_store" in knowledge_base_cache:
         return
-
     docs: List = []
     try:
         PyPDFLoader = _import_langchain_pdf()
@@ -241,119 +225,167 @@ def setup_knowledge_base() -> None:
         knowledge_base_cache["vector_store"] = None
         logging.warning("KB disabled (no docs or embeddings).")
 
-# Initialize on import so app is ready
 initialize_cpu_models()
 setup_knowledge_base()
 
-# =============== Utility: EXIF-based auto calibration ===============
-def _rational_to_float(val) -> Optional[float]:
+# ---------- Calibration helpers ----------
+def _exif_to_dict(pil_img: Image.Image) -> Dict[str, object]:
+    """Best-effort EXIF parse from PIL image."""
+    out = {}
     try:
-        if isinstance(val, TiffImagePlugin.IFDRational):
-            return float(val.numerator) / float(val.denominator or 1)
-        if isinstance(val, tuple) and len(val) == 2 and all(isinstance(x, (int, float)) for x in val):
-            # (num, den)
-            den = val[1] if val[1] else 1.0
-            return float(val[0]) / float(den)
+        exif = pil_img.getexif()
+        if not exif:
+            return out
+        for k, v in exif.items():
+            tag = TAGS.get(k, k)
+            out[tag] = v
+    except Exception:
+        pass
+    return out
+
+def _to_float(val) -> Optional[float]:
+    try:
+        if val is None:
+            return None
+        if isinstance(val, tuple) and len(val) == 2:
+            num, den = float(val[0]), float(val[1]) if float(val[1]) != 0 else 1.0
+            return num / den
         return float(val)
     except Exception:
         return None
 
-def _auto_pixels_per_cm_from_exif(image_pil: Image.Image) -> Tuple[float, str]:
+def _estimate_sensor_width_mm(f_mm: Optional[float], f35: Optional[float]) -> Optional[float]:
     """
-    Try several EXIF / info sources to estimate pixels-per-cm.
-    Return (px_per_cm, source_str).
-    NOTE: Many phones set DPI metadata arbitrarily; we clamp to a sensible range and
-    fall back to DEFAULT_PIXELS_PER_CM if values look bogus.
+    Use 35mm equivalent if present: sensor_width = 36 * f_mm / f35.
     """
-    # 1) PIL .info["dpi"]
-    try:
-        dpi_info = image_pil.info.get("dpi")
-        if isinstance(dpi_info, (tuple, list)) and len(dpi_info) >= 1:
-            xdpi = float(dpi_info[0]) if dpi_info[0] else None
-            if xdpi and 40 <= xdpi <= 1200:
-                ppcm = xdpi / 2.54
-                if 5 <= ppcm <= 500:
-                    return ppcm, "dpi_info"
-    except Exception:
-        pass
+    if f_mm and f35 and f35 > 0:
+        return 36.0 * f_mm / f35
+    return None
 
-    # 2) EXIF XResolution (282), YResolution (283), ResolutionUnit (296) [2 = inch, 3 = cm]
-    try:
-        exif = image_pil.getexif()
-        if exif:
-            xres = _rational_to_float(exif.get(282))  # XResolution
-            unit = int(exif.get(296) or 2)           # default to inches
-            if xres:
-                if unit == 3:  # per cm
-                    if 5 <= xres <= 500:
-                        return xres, "EXIF_XRes_cm"
-                else:          # per inch
-                    ppcm = xres / 2.54
-                    if 5 <= ppcm <= 500:
-                        return ppcm, "EXIF_XRes_in"
-    except Exception:
-        pass
-
-    # 3) Heuristic fallback
-    return DEFAULT_PIXELS_PER_CM, "default"
+def estimate_px_per_cm_from_exif(pil_img: Image.Image, default_px_per_cm: float = DEFAULT_PX_PER_CM) -> Tuple[float, Dict]:
+    """
+    Returns (px_per_cm, meta) using EXIF when available.
+    Formula: field_width_mm = sensor_width_mm * distance_mm / focal_mm
+             px_per_cm = image_width_px / (field_width_mm / 10)
+    """
+    meta = {"used": "default", "f_mm": None, "f35": None, "sensor_w_mm": None, "distance_m": None}
 
-# =============== Drawing helpers ===============
-def _draw_measurement_overlay(
+    try:
+        exif = _exif_to_dict(pil_img)
+        f_mm = _to_float(exif.get("FocalLength"))
+        f35 = _to_float(exif.get("FocalLengthIn35mmFilm") or exif.get("FocalLengthIn35mm"))
+        subj_dist_m = _to_float(exif.get("SubjectDistance"))
+        sensor_w_mm = _estimate_sensor_width_mm(f_mm, f35)
+
+        meta.update({"f_mm": f_mm, "f35": f35, "sensor_w_mm": sensor_w_mm, "distance_m": subj_dist_m})
+
+        if f_mm and sensor_w_mm and subj_dist_m and subj_dist_m > 0:
+            w_px = pil_img.width
+            field_w_mm = sensor_w_mm * (subj_dist_m * 1000.0) / f_mm
+            field_w_cm = field_w_mm / 10.0
+            px_per_cm = w_px / max(field_w_cm, 1e-6)
+
+            # sanity clamp
+            px_per_cm = float(np.clip(px_per_cm, PX_PER_CM_MIN, PX_PER_CM_MAX))
+            meta["used"] = "exif"
+            return px_per_cm, meta
+
+        # If EXIF partial but not enough to solve, keep default
+        return float(default_px_per_cm), meta
+    except Exception as e:
+        logging.warning(f"EXIF calibration failed: {e}")
+        return float(default_px_per_cm), meta
+
+# ---------- Mask processing + measurement ----------
+def largest_component_mask(binary: np.ndarray, min_area_px: int = 50) -> np.ndarray:
+    """Keep only the largest connected component in a binary mask."""
+    num, labels, stats, _ = cv2.connectedComponentsWithStats(binary.astype(np.uint8), connectivity=8)
+    if num <= 1:
+        return binary
+    # stats[:, cv2.CC_STAT_AREA]; skip label 0 (background)
+    areas = stats[1:, cv2.CC_STAT_AREA]
+    largest_idx = 1 + int(np.argmax(areas))
+    if areas.max() < min_area_px:
+        return binary
+    return (labels == largest_idx).astype(np.uint8)
+
+def measure_min_area_rect(mask: np.ndarray, px_per_cm: float) -> Tuple[float, float, Tuple]:
+    """
+    Compute oriented min-area rectangle on mask.
+    Returns (length_cm, breadth_cm, (box_points, center)).
+    """
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return 0.0, 0.0, (None, None)
+    cnt = max(contours, key=cv2.contourArea)
+    rect = cv2.minAreaRect(cnt)  # (center(x,y), (w,h), angle)
+    (w_px, h_px) = rect[1]
+    length_px, breadth_px = (max(w_px, h_px), min(w_px, h_px))
+    length_cm = round(length_px / max(px_per_cm, 1e-6), 2)
+    breadth_cm = round(breadth_px / max(px_per_cm, 1e-6), 2)
+    box = cv2.boxPoints(rect).astype(int)
+    return length_cm, breadth_cm, (box, rect[0])
+
+def count_area_cm2(mask: np.ndarray, px_per_cm: float) -> float:
+    px_count = float(mask.astype(bool).sum())
+    return round(px_count / (max(px_per_cm, 1e-6) ** 2), 2)
+
+def draw_measurement_overlay(
     base_bgr: np.ndarray,
-    rect_xywh: Tuple[int, int, int, int],
+    mask: np.ndarray,
+    rect_box: np.ndarray,
     length_cm: float,
     breadth_cm: float,
+    thickness: int = 2
 ) -> np.ndarray:
     """
-    Draw arrows for vertical (length) and horizontal (breadth) on top of base image.
-    rect_xywh is relative to base_bgr.
+    Draw semi-transparent mask + measurement arrows along the rectangle sides with labels.
     """
-    x, y, w, h = rect_xywh
-    img = base_bgr.copy()
-
-    # Colors (BGR) and styling
-    color = (255, 255, 255)   # white
-    shadow = (0, 0, 0)        # black outline
-    thickness = 2
-    font = cv2.FONT_HERSHEY_SIMPLEX
-
-    # --- Horizontal arrow (breadth) ---
-    y_mid = y + h // 2
-    x_left = x
-    x_right = x + w
-    # shadow line
-    cv2.arrowedLine(img, (x_left, y_mid+1), (x_right, y_mid+1), shadow, thickness+2, cv2.LINE_AA, tipLength=0.02)
-    # main line
-    cv2.arrowedLine(img, (x_left, y_mid), (x_right, y_mid), color, thickness, cv2.LINE_AA, tipLength=0.02)
-
-    # breadth label
-    label_b = f"{breadth_cm:.2f} cm"
-    (tw, th), _ = cv2.getTextSize(label_b, font, 0.7, 2)
-    tx = x + (w - tw) // 2
-    ty = y_mid - 8
-    cv2.putText(img, label_b, (tx+1, ty+1), font, 0.7, shadow, 3, cv2.LINE_AA)
-    cv2.putText(img, label_b, (tx, ty),     font, 0.7, color, 2, cv2.LINE_AA)
-
-    # --- Vertical arrow (length) ---
-    x_mid = x + w // 2
-    y_top = y
-    y_bottom = y + h
-    # shadow line
-    cv2.arrowedLine(img, (x_mid+1, y_top), (x_mid+1, y_bottom), shadow, thickness+2, cv2.LINE_AA, tipLength=0.02)
-    # main line
-    cv2.arrowedLine(img, (x_mid, y_top), (x_mid, y_bottom), color, thickness, cv2.LINE_AA, tipLength=0.02)
-
-    # length label
-    label_l = f"{length_cm:.2f} cm"
-    (tw2, th2), _ = cv2.getTextSize(label_l, font, 0.7, 2)
-    tx2 = x_mid - (tw2 // 2)
-    ty2 = y + th2 + 8
-    cv2.putText(img, label_l, (tx2+1, ty2+1), font, 0.7, shadow, 3, cv2.LINE_AA)
-    cv2.putText(img, label_l, (tx2, ty2),     font, 0.7, color, 2, cv2.LINE_AA)
-
-    return img
-
-# =============== AI PROCESSOR ===============
+    overlay = base_bgr.copy()
+    # red mask overlay
+    colored = np.zeros_like(base_bgr)
+    colored[:, :] = (0, 0, 255)
+    mask3 = np.dstack([mask * 255] * 3)
+    overlay = cv2.addWeighted(overlay, 1.0, (colored & mask3), 0.3, 0)
+
+    # draw rectangle
+    cv2.polylines(overlay, [rect_box], True, (255, 255, 255), thickness)
+
+    # pick the long side & short side arrows
+    # box points are in order; connect midpoints of opposite edges
+    pts = rect_box.reshape(-1, 2)
+    def midpoint(a, b): return ((a[0] + b[0]) // 2, (a[1] + b[1]) // 2)
+
+    # edges: (0-1,1-2,2-3,3-0)
+    mids = [midpoint(pts[i], pts[(i+1) % 4]) for i in range(4)]
+    # vector lengths
+    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)
+
+    # arrowed lines (white with black shadow)
+    def draw_arrow(p1, p2):
+        cv2.arrowedLine(overlay, p1, p2, (0, 0, 0), thickness + 2, tipLength=0.05)
+        cv2.arrowedLine(overlay, p2, p1, (0, 0, 0), thickness + 2, tipLength=0.05)
+        cv2.arrowedLine(overlay, p1, p2, (255, 255, 255), thickness, tipLength=0.05)
+        cv2.arrowedLine(overlay, p2, p1, (255, 255, 255), thickness, tipLength=0.05)
+
+    draw_arrow(mids[long_pair[0]], mids[long_pair[1]])
+    draw_arrow(mids[short_pair[0]], mids[short_pair[1]])
+
+    # labels near the midpoints
+    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)
+
+    put_label(f"{length_cm:.2f} cm", mids[long_pair[0]])
+    put_label(f"{breadth_cm:.2f} cm", mids[short_pair[0]])
+
+    return overlay
+
+# ---------- AI PROCESSOR ----------
 class AIProcessor:
     def __init__(self):
         self.models_cache = models_cache
@@ -368,24 +400,24 @@ class AIProcessor:
         return out_dir
 
     def perform_visual_analysis(self, image_pil: Image.Image) -> Dict:
-        """YOLO detect → (optional) Keras seg → (optional) HF classify → save visuals with measurement overlay."""
+        """
+        YOLO detect → segmentation → largest-component mask →
+        minAreaRect measurement (cm) using px/cm from EXIF when available →
+        save original, detection overlay, segmentation overlay, and annotated overlay.
+        """
         try:
-            image_rgb = image_pil.convert("RGB")
-            image_cv = cv2.cvtColor(np.array(image_rgb), cv2.COLOR_RGB2BGR)
+            # --- Auto calibration from EXIF (before any conversion that might drop EXIF) ---
+            px_per_cm, exif_meta = estimate_px_per_cm_from_exif(image_pil, DEFAULT_PX_PER_CM)
 
-            det = self.models_cache.get("det")
-            if det is None:
-                raise RuntimeError("YOLO model not loaded")
+            # Convert PIL to OpenCV BGR
+            image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
 
-            # ---------- Automatic calibration (px/cm) ----------
-            px_per_cm, calib_src = _auto_pixels_per_cm_from_exif(image_rgb)
-            # keep within reasonable range
-            if not (5.0 <= px_per_cm <= 500.0):
-                px_per_cm, calib_src = DEFAULT_PIXELS_PER_CM, "default"
-            logging.info(f"Calibration: {px_per_cm:.2f} px/cm (source={calib_src})")
+            # --- Detection (YOLO) ---
+            det_model = self.models_cache.get("det")
+            if det_model is None:
+                raise RuntimeError("YOLO model not loaded")
 
-            # YOLO on CPU
-            results = det.predict(image_cv, verbose=False, device="cpu")
+            results = det_model.predict(image_cv, verbose=False, device="cpu")
             if not results or not getattr(results[0], "boxes", None) or len(results[0].boxes) == 0:
                 raise ValueError("No wound could be detected.")
 
@@ -393,122 +425,122 @@ class AIProcessor:
             x1, y1, x2, y2 = [int(v) for v in box]
             x1, y1 = max(0, x1), max(0, y1)
             x2, y2 = min(image_cv.shape[1], x2), min(image_cv.shape[0], y2)
-            detected_region_cv = image_cv[y1:y2, x1:x2]
+            roi = image_cv[y1:y2, x1:x2].copy()
+            if roi.size == 0:
+                raise ValueError("Detected ROI is empty.")
 
-            # Optional segmentation
+            # --- Segmentation (optional but recommended) ---
             seg_model = self.models_cache.get("seg")
+            mask_resized = None
             length_cm = breadth_cm = surface_area_cm2 = 0.0
-            seg_path = None
 
-            rect_xywh_global = None  # for overlay on full image if seg missing
-
-            if seg_model is not None and detected_region_cv.size > 0:
+            if seg_model is not None:
                 try:
-                    input_size = seg_model.input_shape[1:3]
-                    resized = cv2.resize(detected_region_cv, (input_size[1], input_size[0]))
-                    mask_pred = seg_model.predict(np.expand_dims(resized / 255.0, 0), verbose=0)[0]
-                    mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
-
-                    contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-                    if contours:
-                        cnt = max(contours, key=cv2.contourArea)
-                        x, y, w, h = cv2.boundingRect(cnt)
-
-                        # Measurements using calibration
-                        length_cm = round(h / px_per_cm, 2)
-                        breadth_cm = round(w / px_per_cm, 2)
-                        surface_area_cm2 = round(cv2.contourArea(cnt) / (px_per_cm ** 2), 2)
-
-                        # Create segmentation overlay in the cropped region
-                        mask_resized = cv2.resize(
-                            mask_np * 255,
-                            (detected_region_cv.shape[1], detected_region_cv.shape[0]),
-                            interpolation=cv2.INTER_NEAREST,
-                        )
-                        overlay = detected_region_cv.copy()
-                        overlay[mask_resized > 127] = [0, 0, 255]  # red overlay
-                        seg_vis = cv2.addWeighted(detected_region_cv, 0.7, overlay, 0.3, 0)
-
-                        # Draw measurement arrows on seg_vis
-                        # Map rect from mask space -> cropped image space
-                        scale_x = detected_region_cv.shape[1] / float(input_size[1])
-                        scale_y = detected_region_cv.shape[0] / float(input_size[0])
-                        rect_xywh_cropped = (
-                            int(x * scale_x),
-                            int(y * scale_y),
-                            int(w * scale_x),
-                            int(h * scale_y),
-                        )
-                        seg_vis_meas = _draw_measurement_overlay(seg_vis, rect_xywh_cropped, length_cm, breadth_cm)
-
-                        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-                        out_dir = self._ensure_analysis_dir()
-                        seg_path = os.path.join(out_dir, f"segmentation_{ts}.png")
-                        cv2.imwrite(seg_path, seg_vis_meas)
-
-                        # Also store rect in full-image coordinates (if ever needed)
-                        rect_xywh_global = (
-                            x1 + rect_xywh_cropped[0],
-                            y1 + rect_xywh_cropped[1],
-                            rect_xywh_cropped[2],
-                            rect_xywh_cropped[3],
-                        )
+                    H, W = seg_model.input_shape[1:3]
+                    resized = cv2.resize(roi, (W, H))
+                    pred = seg_model.predict(np.expand_dims(resized / 255.0, 0), verbose=0)[0]
+                    raw_mask = pred[:, :, 0]
+
+                    # binarize + clean
+                    mask = (raw_mask > 0.5).astype(np.uint8)
+                    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=1)
+                    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8), iterations=1)
+                    mask = largest_component_mask(mask)
+
+                    # bring back to ROI size
+                    mask_resized = cv2.resize(mask * 255, (roi.shape[1], roi.shape[0]), interpolation=cv2.INTER_NEAREST)
+                    bin_mask_roi = (mask_resized > 127).astype(np.uint8)
+
+                    # measure with oriented rectangle (in ROI pixels)
+                    length_cm, breadth_cm, (box_pts, _) = measure_min_area_rect(bin_mask_roi, px_per_cm)
+                    surface_area_cm2 = count_area_cm2(bin_mask_roi, px_per_cm)
+
+                    # draw overlay with arrows/labels on ROI
+                    anno_roi = draw_measurement_overlay(roi, bin_mask_roi, box_pts, length_cm, breadth_cm)
+
                 except Exception as e:
-                    logging.warning(f"Segmentation skipped: {e}")
+                    logging.warning(f"Segmentation failed/partial: {e}")
+                    mask_resized = None
+                    anno_roi = roi.copy()
+            else:
+                # No segmentation → just draw detection box and keep defaults
+                anno_roi = roi.copy()
+
+            # --- Save all visualizations ---
+            out_dir = self._ensure_analysis_dir()
+            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+            # Original
+            original_path = os.path.join(out_dir, f"original_{ts}.png")
+            cv2.imwrite(original_path, image_cv)
 
-            # Optional classification
+            # Detection overlay (rectangle on full image)
+            det_vis = image_cv.copy()
+            cv2.rectangle(det_vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            detection_path = os.path.join(out_dir, f"detection_{ts}.png")
+            cv2.imwrite(detection_path, det_vis)
+
+            # Segmentation overlay (ROI pasted back into full frame for consistent display)
+            segmentation_path = None
+            annotated_seg_path = None
+            if mask_resized is not None:
+                # compose overlay on full image for "segmentation" view
+                seg_full = image_cv.copy()
+                roi_overlay = roi.copy()
+                red = np.zeros_like(roi_overlay); red[:] = (0, 0, 255)
+                alpha = 0.3
+                roi_overlay = cv2.addWeighted(roi_overlay, 1.0, red, alpha, 0, mask=mask_resized)
+                seg_full[y1:y2, x1:x2] = roi_overlay
+                segmentation_path = os.path.join(out_dir, f"segmentation_{ts}.png")
+                cv2.imwrite(segmentation_path, seg_full)
+
+                # annotated overlay (arrows+labels) placed back into full image
+                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")
+                cv2.imwrite(annotated_seg_path, anno_full)
+
+            # --- Optional classification ---
             wound_type = "Unknown"
             cls_pipe = self.models_cache.get("cls")
             if cls_pipe is not None:
                 try:
-                    detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
-                    preds = cls_pipe(detected_image_pil)
+                    preds = cls_pipe(Image.fromarray(cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)))
                     if preds:
                         wound_type = max(preds, key=lambda x: x.get("score", 0)).get("label", "Unknown")
                 except Exception as e:
                     logging.warning(f"Classification failed: {e}")
 
-            # Save detection & original
-            out_dir = self._ensure_analysis_dir()
-            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-            det_vis = image_cv.copy()
-            cv2.rectangle(det_vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            det_path = os.path.join(out_dir, f"detection_{ts}.png")
-            cv2.imwrite(det_path, det_vis)
-
-            original_path = os.path.join(out_dir, f"original_{ts}.png")
-            cv2.imwrite(original_path, image_cv)
-
             return {
                 "wound_type": wound_type,
-                "length_cm": float(length_cm),
-                "breadth_cm": float(breadth_cm),
-                "surface_area_cm2": float(surface_area_cm2),
-                "calibration_px_per_cm": float(px_per_cm),
-                "calibration_source": calib_src,
+                "length_cm": length_cm,
+                "breadth_cm": breadth_cm,
+                "surface_area_cm2": surface_area_cm2,
+                "px_per_cm": round(px_per_cm, 2),
+                "calibration_meta": exif_meta,  # for debugging/auditing
                 "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": det_path,
-                "segmentation_image_path": seg_path,   # <-- now includes arrow overlay if seg succeeded
+                    if getattr(results[0].boxes, "conf", None) is not None else 0.0,
+                "detection_image_path": detection_path,
+                "segmentation_image_path": segmentation_path,
+                "segmentation_annotated_path": annotated_seg_path,
                 "original_image_path": original_path,
             }
         except Exception as e:
-            logging.error(f"Visual analysis failed: {e}")
+            logging.error(f"Visual analysis failed: {e}", exc_info=True)
             raise
 
+    # ---------- Knowledge base and reporting stay unchanged ----------
     def query_guidelines(self, query: str) -> str:
-        """Query the (optional) guideline knowledge base."""
         try:
             vs = self.knowledge_base_cache.get("vector_store")
             if not vs:
                 return "Knowledge base is not available."
             try:
                 retriever = vs.as_retriever(search_kwargs={"k": 5})
-                docs = retriever.get_relevant_documents(query)  # LC >= 0.2
+                docs = retriever.get_relevant_documents(query)
             except Exception:
                 retriever = vs.as_retriever(search_kwargs={"k": 5})
-                docs = retriever.invoke(query)                 # older LC
+                docs = retriever.invoke(query)
             lines: List[str] = []
             for d in docs:
                 src = (d.metadata or {}).get("source", "N/A")
@@ -530,12 +562,13 @@ class AIProcessor:
 - **Dimensions**: {visual_results.get('length_cm', 0)} cm × {visual_results.get('breadth_cm', 0)} cm
 - **Surface Area**: {visual_results.get('surface_area_cm2', 0)} cm²
 - **Detection Confidence**: {visual_results.get('detection_confidence', 0):.1%}
-- **Calibration**: {visual_results.get('calibration_px_per_cm', 0)} px/cm (source: {visual_results.get('calibration_source','n/a')})
+- **Calibration**: {visual_results.get('px_per_cm','?')} px/cm ({(visual_results.get('calibration_meta') or {}).get('used','default')})
 
 ## 📊 Analysis Images
 - **Original**: {visual_results.get('original_image_path', 'N/A')}
 - **Detection**: {visual_results.get('detection_image_path', 'N/A')}
-- **Segmentation (with measurements)**: {visual_results.get('segmentation_image_path', 'N/A')}
+- **Segmentation**: {visual_results.get('segmentation_image_path', 'N/A')}
+- **Annotated**: {visual_results.get('segmentation_annotated_path', 'N/A')}
 
 ## 🎯 Clinical Summary
 Automated analysis provides quantitative measurements; verify via clinical examination.
@@ -563,7 +596,6 @@ Automated analysis provides quantitative measurements; verify via clinical exami
         image_pil: Image.Image,
         max_new_tokens: Optional[int] = None,
     ) -> str:
-        """Use GPU path when available, fallback otherwise."""
         try:
             report = generate_medgemma_report(
                 patient_info, visual_results, guideline_context, image_pil, max_new_tokens
@@ -577,7 +609,6 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             return self._generate_fallback_report(patient_info, visual_results, guideline_context)
 
     def save_and_commit_image(self, image_pil: Image.Image) -> str:
-        """Save locally and (optionally) upload to HF dataset."""
         try:
             os.makedirs(self.uploads_dir, exist_ok=True)
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -609,7 +640,6 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             return ""
 
     def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: Dict) -> Dict:
-        """End-to-end analysis."""
         try:
             saved_path = self.save_and_commit_image(image_pil)
             visual_results = self.perform_visual_analysis(image_pil)
@@ -656,7 +686,6 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             }
 
     def analyze_wound(self, image, questionnaire_data: Dict) -> Dict:
-        """Public entrypoint used by UI."""
         try:
             if isinstance(image, str):
                 if not os.path.exists(image):
@@ -679,4 +708,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }