Update src/ai_processor.py

src/ai_processor.py

@@ -1,11 +1,6 @@
 # smartheal_ai_processor.py
-# Preserves ALL original class/function names.
-# What you get:
-#  - Uses your segmentation_model.h5 first; clean KMeans fallback if it fails/missing
-#  - Safe overlay (no 'mask' kwarg with addWeighted)
-#  - Always writes a segmentation view (so it never looks like the plain original)
-#  - CPU by default; optional VLM (MedGemma) is OFF unless SMARTHEAL_ENABLE_VLM=1
-#  - Optional @spaces.GPU **stub** (no queue) to satisfy Spaces startup without touching CUDA
+# Verbose, instrumented version — preserves public class/function names
+# Turn on deep logging: export LOGLEVEL=DEBUG SMARTHEAL_DEBUG=1
 
 import os
 import time
@@ -13,31 +8,36 @@ import logging
 from datetime import datetime
 from typing import Optional, Dict, List, Tuple
 
-# Quiet tokenizers; default to CPU for safety on ZeroGPU/Spaces
+# ---- Environment defaults ----
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+LOGLEVEL = os.getenv("LOGLEVEL", "INFO").upper()
+SMARTHEAL_DEBUG = os.getenv("SMARTHEAL_DEBUG", "0") == "1"
 
 import cv2
 import numpy as np
 from PIL import Image
 from PIL.ExifTags import TAGS
 
-# --- Optional: register a harmless @spaces.GPU-decorated stub to silence startup warning ---
+# --- Logging config ---
+logging.basicConfig(
+    level=getattr(logging, LOGLEVEL, logging.INFO),
+    format="%(asctime)s - %(levelname)s - %(message)s",
+)
+
+def _log_kv(prefix: str, kv: Dict):
+    logging.debug(prefix + " | " + " | ".join(f"{k}={v}" for k, v in kv.items()))
+
+# --- Optional Spaces GPU stub (harmless) ---
 try:
     import spaces as _spaces
-
-    @_spaces.GPU(enable_queue=False)  # not queued -> won't start a ZeroGPU worker
+    @_spaces.GPU(enable_queue=False)
     def smartheal_gpu_stub(ping: int = 0) -> str:
-        """No-op so Spaces detects at least one @spaces.GPU function without touching CUDA."""
         return "ready"
-
-    logging.info("Registered @spaces.GPU stub (enable_queue=False); startup detector satisfied.")
-except Exception as _e:
-    # It's fine if 'spaces' isn't available locally.
+    logging.info("Registered @spaces.GPU stub (enable_queue=False).")
+except Exception:
     pass
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-
 UPLOADS_DIR = "uploads"
 os.makedirs(UPLOADS_DIR, exist_ok=True)
 
@@ -49,6 +49,11 @@ DATASET_ID = "SmartHeal/wound-image-uploads"
 DEFAULT_PX_PER_CM = 38.0
 PX_PER_CM_MIN, PX_PER_CM_MAX = 5.0, 1200.0
 
+# Segmentation preprocessing knobs
+SEG_EXPECTS_RGB = os.getenv("SEG_EXPECTS_RGB", "1") == "1"  # most TF models trained on RGB
+SEG_NORM = os.getenv("SEG_NORM", "0to1")                    # "0to1" | "imagenet"
+SEG_THRESH = float(os.getenv("SEG_THRESH", "0.5"))
+
 models_cache: Dict[str, object] = {}
 knowledge_base_cache: Dict[str, object] = {}
 
@@ -86,7 +91,7 @@ def _import_hf_hub():
     from huggingface_hub import HfApi, HfFolder
     return HfApi, HfFolder
 
-# ---------- LLM report (OFF by default; enable with SMARTHEAL_ENABLE_VLM=1) ----------
+# ---------- VLM (disabled by default) ----------
 def generate_medgemma_report(
     patient_info: str,
     visual_results: Dict,
@@ -94,24 +99,18 @@ def generate_medgemma_report(
     image_pil: Image.Image,
     max_new_tokens: Optional[int] = None,
 ) -> str:
-    """
-    CPU-only MedGemma call (safe). Disabled by default to avoid env mismatches.
-    Set SMARTHEAL_ENABLE_VLM=1 to try loading the model.
-    """
     if os.getenv("SMARTHEAL_ENABLE_VLM", "0") != "1":
         return "⚠️ VLM disabled"
-
     try:
         from transformers import pipeline
         pipe = pipeline(
             task="image-text-to-text",
             model="google/medgemma-4b-it",
-            device_map=None,               # CPU
+            device_map=None,
             token=HF_TOKEN,
             trust_remote_code=True,
-            model_kwargs={"low_cpu_mem_usage": True},  # avoid 'use_cache' arg mismatch
+            model_kwargs={"low_cpu_mem_usage": True},
         )
-
         prompt = (
             "You are a medical AI assistant. Analyze this wound image and patient data.\n\n"
             f"Patient: {patient_info}\n"
@@ -120,12 +119,10 @@ def generate_medgemma_report(
             "Provide a structured report with:\n"
             "1. Clinical Summary\n2. Treatment Recommendations\n3. Risk Assessment\n4. Monitoring Plan\n"
         )
-
         messages = [{"role": "user", "content": [
             {"type": "image", "image": image_pil},
             {"type": "text",  "text": prompt},
         ]}]
-
         out = pipe(text=messages, max_new_tokens=max_new_tokens or 600, do_sample=False, temperature=0.7)
         if out and len(out) > 0:
             try:
@@ -174,7 +171,10 @@ def initialize_cpu_models() -> None:
         try:
             if os.path.exists(SEG_MODEL_PATH):
                 models_cache["seg"] = load_segmentation_model()
-                logging.info("✅ Segmentation model loaded (CPU)")
+                m = models_cache["seg"]
+                ishape = getattr(m, "input_shape", None)
+                oshape = getattr(m, "output_shape", None)
+                logging.info(f"✅ Segmentation model loaded (CPU) | input_shape={ishape} output_shape={oshape}")
             else:
                 models_cache["seg"] = None
                 logging.warning("Segmentation model file missing; skipping.")
@@ -283,47 +283,133 @@ def estimate_px_per_cm_from_exif(pil_img: Image.Image, default_px_per_cm: float
     except Exception:
         return float(default_px_per_cm), meta
 
-# ---------- Segmentation (model-first, KMeans fallback) ----------
-def segment_wound(image: np.ndarray) -> np.ndarray:
+# ---------- 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
+
+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)
+
+# Global last debug dict (per-process) to attach into results
+_last_seg_debug: Dict[str, object] = {}
+
+def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndarray, Dict[str, object]]:
     """
-    Segments wound from a preprocessed ROI image, with a fallback to KMeans if the model fails.
-    Returns a mask in 0..255 (uint8), same HxW as input image.
+    Attempts TF segmentation first; falls back to KMeans if needed.
+    Returns (mask_uint8_0_255, debug_dict)
     """
-    segmentation_model = models_cache.get("seg", None)
+    global _last_seg_debug
+    _last_seg_debug = {}
+
+    seg_model = models_cache.get("seg", None)
+    used = "fallback_kmeans"
+    reason = "no_model"
+    heatmap_path = None
+    saw_roi_path = None
 
-    if segmentation_model is not None:
+    if seg_model is not None:
         try:
-            input_shape = getattr(segmentation_model, "input_shape", None)
-            if input_shape is None or len(input_shape) < 3:
-                raise ValueError(f"Bad seg input_shape: {input_shape}")
-            H, W = int(input_shape[1]), int(input_shape[2])  # (None,H,W,C)
-
-            resized = cv2.resize(image, (W, H))                    # (W,H)
-            norm = np.expand_dims(resized / 255.0, axis=0)         # (1,H,W,3)
-            prediction = segmentation_model.predict(norm, verbose=0)
-
-            # Handle models with multiple outputs
-            if isinstance(prediction, (list, tuple)):
-                prediction = prediction[0]
-            # squeeze batch dim if present
-            prediction = prediction[0] if getattr(prediction, "ndim", 0) >= 3 else prediction
-
-            pred2d = np.squeeze(prediction)                        # (H,W) or (H,W,1)->(H,W)
-            mask_prob = cv2.resize(pred2d, (image.shape[1], image.shape[0]))
-            mask = (mask_prob >= 0.5).astype(np.uint8) * 255
-            return mask.astype(np.uint8)
+            ishape = getattr(seg_model, "input_shape", None)
+            if not ishape or len(ishape) < 4:
+                raise ValueError(f"Bad seg input_shape: {ishape}")
+            th, tw = int(ishape[1]), int(ishape[2])
+            x = _preprocess_for_seg(image_bgr, (th, tw))
+            saw_roi = (cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) if SEG_EXPECTS_RGB else image_bgr)
+            if SMARTHEAL_DEBUG:
+                saw_roi_path = os.path.join(out_dir, f"roi_for_seg_{ts}.png")
+                cv2.imwrite(saw_roi_path, (cv2.cvtColor(saw_roi, cv2.COLOR_RGB2BGR) if SEG_EXPECTS_RGB else saw_roi))
+
+            # Inference
+            pred = seg_model.predict(x, verbose=0)
+            if isinstance(pred, (list, tuple)):
+                pred = pred[0]
+            p = _to_prob(pred)  # HxW
+            p = cv2.resize(p, (image_bgr.shape[1], image_bgr.shape[0]))  # back to ROI size
+
+            # Debug stats
+            pmin, pmax, pmean = float(p.min()), float(p.max()), float(p.mean())
+            _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")
+                cv2.imwrite(heatmap_path, heat)
+
+            # Threshold
+            thr = SEG_THRESH
+            mask = (p >= thr).astype(np.uint8) * 255
+            pos = int((mask > 0).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}")
+
+            used = "tf_model"
+            reason = "ok"
+
+            _last_seg_debug = {
+                "used": used,
+                "reason": reason,
+                "input_shape": ishape,
+                "prob_min": pmin, "prob_max": pmax, "prob_mean": pmean,
+                "threshold": thr,
+                "positive_fraction": frac,
+                "heatmap_path": heatmap_path,
+                "roi_seen_by_model": saw_roi_path,
+            }
+            return mask.astype(np.uint8), _last_seg_debug
+
         except Exception as e:
-            logging.warning(f"⚠️ Segmentation model prediction failed: {e}. Falling back to KMeans.")
+            reason = f"model_failed: {e}"
+            logging.warning(f"⚠️ Segmentation model prediction failed → fallback. Reason: {e}")
 
-    # --- Fallback: color clustering (KMeans, k=2), pick 'reddest' cluster in Lab a* ---
-    Z = image.reshape((-1, 3)).astype(np.float32)
+    # --- Fallback: KMeans (k=2), pick 'reddest' cluster in Lab a* ---
+    Z = image_bgr.reshape((-1, 3)).astype(np.float32)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
     _, labels, centers = cv2.kmeans(Z, 2, None, criteria, 5, cv2.KMEANS_PP_CENTERS)
     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]))  # a* channel (redness)
-    mask = (labels.reshape(image.shape[:2]) == wound_idx).astype(np.uint8) * 255
-    return mask.astype(np.uint8)
+    wound_idx = int(np.argmax(centers_lab[:, 1]))  # maximize a* (redness)
+    mask = (labels.reshape(image_bgr.shape[:2]) == wound_idx).astype(np.uint8) * 255
+
+    pos = int((mask > 0).sum()); frac = pos / float(mask.size)
+    logging.info(f"KMeans USED | pos_px={pos} pos_frac={frac:.4f}")
+
+    _last_seg_debug = {
+        "used": used,
+        "reason": reason,
+        "kmeans_centers_bgr": centers.tolist(),
+        "kmeans_centers_lab": centers_lab.astype(float).tolist(),
+        "positive_fraction": frac,
+        "heatmap_path": heatmap_path,
+        "roi_seen_by_model": saw_roi_path,
+    }
+    return mask.astype(np.uint8), _last_seg_debug
 
 # ---------- Measurement + overlay helpers ----------
 def largest_component_mask(binary01: np.ndarray, min_area_px: int = 50) -> np.ndarray:
@@ -362,16 +448,21 @@ def draw_measurement_overlay(
     thickness: int = 2
 ) -> np.ndarray:
     overlay = base_bgr.copy()
-    # Safe masked blend (OpenCV addWeighted has no 'mask' kwarg)
+
+    # Strong overlay + contour
     red = np.zeros_like(overlay); red[:] = (0, 0, 255)
-    blended = cv2.addWeighted(overlay, 1.0, red, 0.3, 0)
-    m3 = np.dstack([mask01 * 255] * 3).astype("uint8")
-    overlay = cv2.add(cv2.bitwise_and(overlay, cv2.bitwise_not(m3)),
-                      cv2.bitwise_and(blended, m3))
+    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)
+
+    # Draw contour
+    cnts, _ = cv2.findContours((mask01 * 255).astype(np.uint8), 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)]
@@ -444,32 +535,40 @@ class AIProcessor:
                 except Exception:
                     raise RuntimeError("Detected ROI is empty.")
 
+            out_dir = self._ensure_analysis_dir()
+            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+
             # --- Segmentation (model-first + KMeans fallback) ---
-            mask_u8_255 = segment_wound(roi)               # 0..255
+            mask_u8_255, seg_debug = segment_wound(roi, ts, out_dir)
             mask01 = (mask_u8_255 > 127).astype(np.uint8)
+
+            # Post-processing + metrics
             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())}")
 
             # --- 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)
                 anno_roi = draw_measurement_overlay(roi, mask01, box_pts, length_cm, breadth_cm)
+                segmentation_empty = False
             else:
-                # fallback to detection box if segmentation is empty
                 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)
                 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)
+                cv2.line(anno_roi, (0, 0), (anno_roi.shape[1]-1, anno_roi.shape[0]-1), (0, 0, 255), 2)
+                cv2.line(anno_roi, (anno_roi.shape[1]-1, 0), (0, anno_roi.shape[0]-1), (0, 0, 255), 2)
                 box_pts = None
+                segmentation_empty = True
 
-            # --- Save visualizations (ALWAYS create a segmentation image) ---
-            out_dir = self._ensure_analysis_dir()
-            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-
+            # --- Save visualizations ---
             original_path = os.path.join(out_dir, f"original_{ts}.png")
             cv2.imwrite(original_path, image_cv)
 
@@ -478,27 +577,30 @@ class AIProcessor:
             detection_path = os.path.join(out_dir, f"detection_{ts}.png")
             cv2.imwrite(detection_path, det_vis)
 
-            # Save ROI mask image (helps debug)
             roi_mask_path = os.path.join(out_dir, f"roi_mask_{ts}.png")
             cv2.imwrite(roi_mask_path, (mask01 * 255).astype(np.uint8))
 
-            # Segmentation overlay (paste back to full image). If mask empty, tint ROI red so it's NOT identical to original.
-            seg_full = image_cv.copy()
-            red = np.zeros_like(roi); red[:] = (0, 0, 255)
-            if mask01.any():
-                blended = cv2.addWeighted(roi, 1.0, red, 0.30, 0)
-                m3 = np.dstack([mask01 * 255] * 3).astype("uint8")
-                roi_overlay = cv2.add(cv2.bitwise_and(roi, cv2.bitwise_not(m3)),
-                                      cv2.bitwise_and(blended, m3))
+            # ROI overlay (very clear)
+            mask255 = (mask01 * 255).astype(np.uint8)
+            mask3   = cv2.merge([mask255, mask255, mask255])
+            red     = np.zeros_like(roi); red[:] = (0, 0, 255)
+            alpha   = 0.55
+            tinted  = cv2.addWeighted(roi, 1 - alpha, red, alpha, 0)
+            if mask255.any():
+                roi_overlay = np.where(mask3 > 0, tinted, roi)
+                cnts, _ = cv2.findContours(mask255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                cv2.drawContours(roi_overlay, cnts, -1, (255, 255, 255), 2)
             else:
-                # No mask → light red tint over the ROI to make the "segmentation" view visually distinct.
-                roi_overlay = cv2.addWeighted(roi, 0.75, red, 0.25, 0)
+                roi_overlay = anno_roi  # already marked X
 
+            seg_full = image_cv.copy()
             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 (arrows + labels)
+            segmentation_roi_path = os.path.join(out_dir, f"segmentation_roi_{ts}.png")
+            cv2.imwrite(segmentation_roi_path, roi_overlay)
+
             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")
@@ -515,6 +617,17 @@ class AIProcessor:
                 except Exception as e:
                     logging.warning(f"Classification failed: {e}")
 
+            # Log end-of-seg summary
+            seg_summary = {
+                "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),
+                "segmentation_empty": segmentation_empty,
+                "exif_px_per_cm": round(px_per_cm, 3),
+            }
+            _log_kv("SEG_SUMMARY", seg_summary)
+
             return {
                 "wound_type": wound_type,
                 "length_cm": length_cm,
@@ -525,16 +638,19 @@ class AIProcessor:
                 "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": detection_path,
-                "segmentation_image_path": segmentation_path,            # always present
+                "segmentation_image_path": segmentation_path,
                 "segmentation_annotated_path": annotated_seg_path,
-                "roi_mask_path": roi_mask_path,                          # helpful for debugging
+                "segmentation_roi_path": segmentation_roi_path,
+                "roi_mask_path": roi_mask_path,
+                "segmentation_empty": segmentation_empty,
+                "segmentation_debug": seg_debug,
                 "original_image_path": original_path,
             }
         except Exception as e:
             logging.error(f"Visual analysis failed: {e}", exc_info=True)
             raise
 
-    # ---------- Knowledge base + reporting (unchanged names) ----------
+    # ---------- Knowledge base + reporting ----------
     def query_guidelines(self, query: str) -> str:
         try:
             vs = self.knowledge_base_cache.get("vector_store")