Update src/ai_processor.py

src/ai_processor.py

@@ -1,16 +1,17 @@
 # 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.
 
 import os
 import time
 import logging
 from datetime import datetime
-from typing import Optional, Dict, List
+from typing import Optional, Dict, List, Tuple, Union
 
 import cv2
 import numpy as np
-from PIL import Image
+from PIL import Image, TiffImagePlugin
 
 # =============== LOGGING ===============
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -24,7 +25,8 @@ 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)
-PIXELS_PER_CM = 38
+# Fallback px/cm if we cannot calibrate from EXIF
+DEFAULT_PIXELS_PER_CM = 38.0
 
 # =============== CACHES ===============
 models_cache: Dict[str, object] = {}
@@ -243,12 +245,119 @@ def setup_knowledge_base() -> None:
 initialize_cpu_models()
 setup_knowledge_base()
 
+# =============== Utility: EXIF-based auto calibration ===============
+def _rational_to_float(val) -> Optional[float]:
+    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)
+        return float(val)
+    except Exception:
+        return None
+
+def _auto_pixels_per_cm_from_exif(image_pil: Image.Image) -> Tuple[float, str]:
+    """
+    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.
+    """
+    # 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
+
+    # 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"
+
+# =============== Drawing helpers ===============
+def _draw_measurement_overlay(
+    base_bgr: np.ndarray,
+    rect_xywh: Tuple[int, int, int, int],
+    length_cm: float,
+    breadth_cm: float,
+) -> np.ndarray:
+    """
+    Draw arrows for vertical (length) and horizontal (breadth) on top of base image.
+    rect_xywh is relative to base_bgr.
+    """
+    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 ===============
 class AIProcessor:
     def __init__(self):
         self.models_cache = models_cache
         self.knowledge_base_cache = knowledge_base_cache
-        self.px_per_cm = PIXELS_PER_CM
         self.uploads_dir = UPLOADS_DIR
         self.dataset_id = DATASET_ID
         self.hf_token = HF_TOKEN
@@ -259,14 +368,22 @@ 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."""
+        """YOLO detect → (optional) Keras seg → (optional) HF classify → save visuals with measurement overlay."""
         try:
-            image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
+            image_rgb = image_pil.convert("RGB")
+            image_cv = cv2.cvtColor(np.array(image_rgb), cv2.COLOR_RGB2BGR)
 
             det = self.models_cache.get("det")
             if det is None:
                 raise RuntimeError("YOLO model not loaded")
 
+            # ---------- 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})")
+
             # YOLO on CPU
             results = det.predict(image_cv, verbose=False, device="cpu")
             if not results or not getattr(results[0], "boxes", None) or len(results[0].boxes) == 0:
@@ -280,8 +397,11 @@ class AIProcessor:
 
             # Optional segmentation
             seg_model = self.models_cache.get("seg")
-            length = breadth = area = 0.0
+            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:
                 try:
                     input_size = seg_model.input_shape[1:3]
@@ -293,24 +413,46 @@ class AIProcessor:
                     if contours:
                         cnt = max(contours, key=cv2.contourArea)
                         x, y, w, h = cv2.boundingRect(cnt)
-                        length = round(h / self.px_per_cm, 2)
-                        breadth = round(w / self.px_per_cm, 2)
-                        area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
 
-                        # overlay visualization
+                        # 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]
+                        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)
+                        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],
+                        )
                 except Exception as e:
                     logging.warning(f"Segmentation skipped: {e}")
 
@@ -339,14 +481,16 @@ class AIProcessor:
 
             return {
                 "wound_type": wound_type,
-                "length_cm": length,
-                "breadth_cm": breadth,
-                "surface_area_cm2": area,
+                "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,
                 "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,
+                "segmentation_image_path": seg_path,   # <-- now includes arrow overlay if seg succeeded
                 "original_image_path": original_path,
             }
         except Exception as e:
@@ -386,11 +530,12 @@ 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')})
 
 ## 📊 Analysis Images
 - **Original**: {visual_results.get('original_image_path', 'N/A')}
 - **Detection**: {visual_results.get('detection_image_path', 'N/A')}
-- **Segmentation**: {visual_results.get('segmentation_image_path', 'N/A')}
+- **Segmentation (with measurements)**: {visual_results.get('segmentation_image_path', 'N/A')}
 
 ## 🎯 Clinical Summary
 Automated analysis provides quantitative measurements; verify via clinical examination.
@@ -534,4 +679,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }