Update src/ai_processor.py

src/ai_processor.py

@@ -1,7 +1,6 @@
 # smartheal_ai_processor.py
-# Fully functional: "segment like snippet" while preserving ALL original names.
-# You can keep using AIProcessor.perform_visual_analysis / analyze_wound / full_analysis_pipeline
-# exactly as before. A convenience AIProcessor.segment_like_snippet(...) is added.
+# Fully functional: robust segmentation + safe overlays + conditional GPU wrapper.
+# All original class/function names preserved. New helpers are additive.
 
 import os
 import time
@@ -9,18 +8,14 @@ import logging
 from datetime import datetime
 from typing import Optional, Dict, List, Tuple
 
+# --- quiet tokenizers fork warning (HF) ---
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
 import cv2
 import numpy as np
 from PIL import Image, ImageOps
 from PIL.ExifTags import TAGS
 
-try:
-    import gradio as gr
-except Exception:
-    class _GrErr(RuntimeError): ...
-    class gr:  # shim so `gr.Error` won’t crash if Gradio isn’t present
-        Error = _GrErr
-
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
 UPLOADS_DIR = "uploads"
@@ -44,7 +39,7 @@ def _import_ultralytics():
 
 def _import_tf_loader():
     import tensorflow as tf
-    tf.config.set_visible_devices([], "GPU")  # force CPU for TF
+    tf.config.set_visible_devices([], "GPU")  # force CPU for TF to avoid CUDA contention
     from tensorflow.keras.models import load_model
     return load_model
 
@@ -68,69 +63,90 @@ def _import_hf_hub():
     from huggingface_hub import HfApi, HfFolder
     return HfApi, HfFolder
 
-# ---------- Spaces GPU function (kept name/behavior) ----------
-try:
-    import spaces
-
-    @spaces.GPU(enable_queue=True, duration=90)
-    def generate_medgemma_report(
-        patient_info: str,
-        visual_results: Dict,
-        guideline_context: str,
-        image_pil: Image.Image,
-        max_new_tokens: Optional[int] = None,
-    ) -> str:
-        try:
-            import torch
-            from transformers import pipeline
-
+# ---------- Conditional Spaces GPU function ----------
+# Avoid scheduling a GPU worker when CUDA is not available (prevents cudaGetDeviceCount crash)
+def _cuda_available() -> bool:
+    try:
+        import torch
+        return bool(getattr(torch, "cuda", None)) and torch.cuda.is_available()
+    except Exception:
+        return False
+
+def _generate_medgemma_report_core(
+    patient_info: str,
+    visual_results: Dict,
+    guideline_context: str,
+    image_pil: Image.Image,
+    max_new_tokens: Optional[int] = None,
+) -> str:
+    try:
+        from transformers import pipeline
+        # Use CPU by default; if CUDA truly available, pipeline can still map automatically
+        pipe = pipeline(
+            "image-text-to-text",
+            model="google/medgemma-4b-it",
+            device_map="auto" if _cuda_available() else None,
+            token=HF_TOKEN,
+            model_kwargs={"low_cpu_mem_usage": True, "use_cache": True},
+        )
+
+        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"
+        )
+
+        messages = [{"role": "user", "content": [
+            {"type": "image", "image": image_pil},
+            {"type": "text",  "text": prompt},
+        ]}]
+
+        t0 = time.time()
+        out = pipe(
+            text=messages,
+            max_new_tokens=max_new_tokens or 800,
+            do_sample=False,
+            temperature=0.7,
+        )
+        logging.info(f"✅ MedGemma finished in {time.time()-t0:.2f}s")
+
+        if out and len(out) > 0:
             try:
-                if hasattr(torch, "cuda") and torch.cuda.is_available():
-                    torch.cuda.empty_cache()
+                return out[0]["generated_text"][-1].get("content", "").strip() or "⚠️ Empty response"
             except Exception:
-                pass
-
-            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"
-            )
-
-            pipe = pipeline(
-                "image-text-to-text",
-                model="google/medgemma-4b-it",
-                device_map="auto",
-                token=HF_TOKEN,
-                model_kwargs={"low_cpu_mem_usage": True, "use_cache": True},
-            )
-
-            messages = [{"role": "user", "content": [
-                {"type": "image", "image": image_pil},
-                {"type": "text",  "text": prompt},
-            ]}]
-
-            t0 = time.time()
-            out = pipe(
-                text=messages,
-                max_new_tokens=max_new_tokens or 800,
-                do_sample=False,
-                temperature=0.7,
-                pad_token_id=pipe.tokenizer.eos_token_id,
-            )
-            logging.info(f"✅ MedGemma finished in {time.time()-t0:.2f}s")
+                return (out[0].get("generated_text", "") or "").strip() or "⚠️ Empty response"
+        return "⚠️ No output generated"
+    except Exception as e:
+        logging.error(f"❌ MedGemma generation error: {e}")
+        return "⚠️ GPU/LLM worker unavailable"
 
-            if out and len(out) > 0:
-                try:
-                    return out[0]["generated_text"][-1].get("content", "").strip() or "⚠️ Empty response"
-                except Exception:
-                    return (out[0].get("generated_text", "") or "").strip() or "⚠️ Empty response"
-            return "⚠️ No output generated"
-        except Exception as e:
-            logging.error(f"❌ MedGemma generation error: {e}")
-            return "⚠️ GPU worker unavailable"
+# Preserve the SAME public function name.
+# Only decorate with @spaces.GPU if CUDA is truly available.
+try:
+    import spaces
+    if _cuda_available():
+        @spaces.GPU(enable_queue=True, duration=90)
+        def generate_medgemma_report(
+            patient_info: str,
+            visual_results: Dict,
+            guideline_context: str,
+            image_pil: Image.Image,
+            max_new_tokens: Optional[int] = None,
+        ) -> str:
+            return _generate_medgemma_report_core(patient_info, visual_results, guideline_context, image_pil, max_new_tokens)
+    else:
+        def generate_medgemma_report(
+            patient_info: str,
+            visual_results: Dict,
+            guideline_context: str,
+            image_pil: Image.Image,
+            max_new_tokens: Optional[int] = None,
+        ) -> str:
+            # no decorator -> no GPU worker init -> no cudaGetDeviceCount crash
+            return _generate_medgemma_report_core(patient_info, visual_results, guideline_context, image_pil, max_new_tokens)
 except Exception:
     def generate_medgemma_report(
         patient_info: str,
@@ -139,9 +155,9 @@ except Exception:
         image_pil: Image.Image,
         max_new_tokens: Optional[int] = None,
     ) -> str:
-        return "⚠️ GPU not available"
+        return _generate_medgemma_report_core(patient_info, visual_results, guideline_context, image_pil, max_new_tokens)
 
-# ---------- Initialize CPU models (same function names/behavior) ----------
+# ---------- Initialize CPU models ----------
 def load_yolo_model():
     YOLO = _import_ultralytics()
     return YOLO(YOLO_MODEL_PATH)
@@ -235,7 +251,7 @@ def setup_knowledge_base() -> None:
 initialize_cpu_models()
 setup_knowledge_base()
 
-# ---------- Calibration helpers (added, names unchanged elsewhere) ----------
+# ---------- Calibration helpers ----------
 def _exif_to_dict(pil_img: Image.Image) -> Dict[str, object]:
     out = {}
     try:
@@ -288,7 +304,52 @@ 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
 
-# ---------- Mask processing + measurement (helpers added) ----------
+# ---------- Segmentation helpers (additive; names preserved elsewhere) ----------
+def _get_seg_hw(seg_model) -> Tuple[int, int]:
+    shp = getattr(seg_model, "input_shape", None)
+    if shp and len(shp) >= 4:
+        return int(shp[1]), int(shp[2])
+    # try Keras .inputs shape
+    try:
+        shp = seg_model.inputs[0].shape
+        return int(shp[1]), int(shp[2])
+    except Exception:
+        pass
+    raise ValueError(f"Cannot infer (H,W) from segmentation model input shape: {shp}")
+
+def _to_prob(mask_pred: np.ndarray) -> np.ndarray:
+    m = np.array(mask_pred)
+    # squeeze batch/channel dims
+    while m.ndim > 2:
+        if m.shape[0] == 1:
+            m = np.squeeze(m, axis=0)
+        if m.ndim > 2 and m.shape[-1] == 1:
+            m = np.squeeze(m, axis=-1)
+        if m.ndim == 3 and m.shape[-1] > 1:
+            # pick the most active channel
+            ch = np.argmax(m.reshape(-1, m.shape[-1]).mean(0))
+            m = m[..., ch]
+        if m.ndim <= 2:
+            break
+    m = m.astype("float32")
+    # if looks like logits -> sigmoid
+    if m.max() > 1.5 or m.min() < -0.5:
+        m = 1.0 / (1.0 + np.exp(-m))
+    return np.clip(m, 0.0, 1.0)
+
+def _adaptive_threshold(prob: np.ndarray, hard: float = 0.5) -> np.ndarray:
+    if (prob >= hard).sum() > 0:
+        return (prob >= hard).astype("uint8")
+    # try Otsu
+    m8 = (np.clip(prob, 0, 1) * 255).astype("uint8")
+    try:
+        # we only need the threshold value _
+        _, _ = cv2.threshold(m8, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        return (m8 >= _).astype("uint8")
+    except Exception:
+        p = float(np.percentile(prob, 99.0))
+        return (prob >= max(0.2, min(0.9, p))).astype("uint8")
+
 def largest_component_mask(binary: np.ndarray, min_area_px: int = 50) -> np.ndarray:
     num, labels, stats, _ = cv2.connectedComponentsWithStats(binary.astype(np.uint8), connectivity=8)
     if num <= 1:
@@ -324,16 +385,13 @@ def draw_measurement_overlay(
     breadth_cm: float,
     thickness: int = 2
 ) -> np.ndarray:
-    """Safe overlay (no mask arg to addWeighted)."""
     overlay = base_bgr.copy()
-
-    # red mask overlay only where mask==1
+    # safe blend: blend once, then gate with mask (no mask kwarg!)
     colored = np.zeros_like(base_bgr); colored[:] = (0, 0, 255)
-    mask3 = np.dstack([mask01 * 255] * 3).astype(np.uint8)
     blended = cv2.addWeighted(overlay, 1.0, colored, 0.3, 0)
-    # keep blended only on mask
-    blended_masked = cv2.bitwise_and(blended, mask3)
-    bg = cv2.bitwise_and(overlay, cv2.bitwise_not(mask3))
+    m3 = np.dstack([mask01 * 255] * 3).astype("uint8")
+    blended_masked = cv2.bitwise_and(blended, m3)
+    bg = cv2.bitwise_and(overlay, cv2.bitwise_not(m3))
     overlay = cv2.add(bg, blended_masked)
 
     if rect_box is not None:
@@ -365,7 +423,7 @@ def draw_measurement_overlay(
         put_label(f"{breadth_cm:.2f} cm", mids[short_pair[0]])
     return overlay
 
-# ---------- AI PROCESSOR (ALL names preserved) ----------
+# ---------- AI PROCESSOR ----------
 class AIProcessor:
     def __init__(self):
         self.models_cache = models_cache
@@ -379,103 +437,26 @@ class AIProcessor:
         os.makedirs(out_dir, exist_ok=True)
         return out_dir
 
-    # NEW helper that mirrors your short snippet exactly (you can call or ignore)
-    def segment_like_snippet(self, image_pil: Image.Image) -> Tuple[Dict, Image.Image, Image.Image]:
+    def perform_visual_analysis(self, image_pil: Image.Image) -> Dict:
         """
-        Returns (visual_results, detected_image_pil, mask_pil) exactly like your snippet.
-        Uses EXIF-calibrated px/cm if available; otherwise DEFAULT_PX_PER_CM.
+        Detect → crop ROI → (optional) segment → cleanup → largest component →
+        oriented minAreaRect in cm (EXIF-calibrated) → save original/detect/seg/annotated.
         """
-        if image_pil is None:
-            raise gr.Error("No image provided.")
-
-        px_per_cm, _ = estimate_px_per_cm_from_exif(image_pil, DEFAULT_PX_PER_CM)
-
-        # Convert image
-        image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
-
-        # Detection
-        det_model = self.models_cache.get("det")
-        if det_model is None:
-            raise gr.Error("Detection model not loaded.")
-        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 gr.Error("No wound could be detected.")
-        box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
-        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]
-        if detected_region_cv.size == 0:
-            raise gr.Error("Detected ROI is empty.")
-
-        # Segmentation
-        seg_model = self.models_cache.get("seg")
-        mask_roi_01 = None
-        if seg_model is not None:
-            H, W = seg_model.input_shape[1:3]
-            resized = cv2.resize(detected_region_cv, (W, H))
-            pred = seg_model.predict(np.expand_dims(resized / 255.0, 0), verbose=0)[0]
-            raw = pred[:, :, 0]
-            mask = (raw > 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, min_area_px=50)
-            mask_roi_01 = cv2.resize(mask, (detected_region_cv.shape[1], detected_region_cv.shape[0]),
-                                     interpolation=cv2.INTER_NEAREST).astype(np.uint8)
-        else:
-            mask_roi_01 = np.zeros(detected_region_cv.shape[:2], dtype=np.uint8)
-
-        # Measurement (oriented rect)
-        if mask_roi_01.any():
-            length_cm, breadth_cm, _ = measure_min_area_rect(mask_roi_01, px_per_cm)
-            area_cm2 = count_area_cm2(mask_roi_01, px_per_cm)
-        else:
-            # fall back to detection box
-            h_px = max(0, y2 - y1)
-            w_px = max(0, x2 - x1)
-            length_cm, breadth_cm = round(h_px / px_per_cm, 2), round(w_px / px_per_cm, 2)
-            area_cm2 = round((h_px * w_px) / (px_per_cm ** 2), 2)
-
-        # Classification (optional)
-        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)
-                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}")
-                detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
-        else:
-            detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
-
-        visual_results = {
-            "wound_type": wound_type,
-            "length_cm": length_cm,
-            "breadth_cm": breadth_cm,
-            "surface_area_cm2": area_cm2
-        }
-        mask_pil = Image.fromarray((mask_roi_01 * 255).astype(np.uint8))
-        return visual_results, detected_image_pil, mask_pil
-
-    # ORIGINAL NAME preserved; inside it we follow the snippet-style flow and also save overlays
-    def perform_visual_analysis(self, image_pil: Image.Image) -> Dict:
         try:
             # --- Auto calibration from EXIF ---
             px_per_cm, exif_meta = estimate_px_per_cm_from_exif(image_pil, DEFAULT_PX_PER_CM)
 
-            # Convert image
+            # Convert PIL to OpenCV BGR
             image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
 
-            # --- Detection ---
+            # --- Detection (YOLO) ---
             det_model = self.models_cache.get("det")
             if det_model is None:
                 raise RuntimeError("YOLO model not loaded")
 
             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:
+                import gradio as gr  # local import to keep class name intact if gradio missing
                 raise gr.Error("No wound could be detected.")
 
             box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
@@ -484,39 +465,46 @@ class AIProcessor:
             x2, y2 = min(image_cv.shape[1], x2), min(image_cv.shape[0], y2)
             roi = image_cv[y1:y2, x1:x2].copy()
             if roi.size == 0:
+                import gradio as gr
                 raise gr.Error("Detected ROI is empty.")
 
-            # --- Segmentation (snippet style) ---
+            # --- Segmentation (robust) ---
             seg_model = self.models_cache.get("seg")
             mask_roi_01 = None
             if seg_model is not None:
-                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]
-                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)
-                mask_roi_01 = cv2.resize(mask, (roi.shape[1], roi.shape[0]), interpolation=cv2.INTER_NEAREST)
+                try:
+                    H, W = _get_seg_hw(seg_model)      # robust (H,W)
+                    resized = cv2.resize(roi, (W, H))  # cv2.resize expects (W,H)
+                    pred = seg_model.predict(np.expand_dims(resized / 255.0, 0), verbose=0)
+                    prob = _to_prob(pred)              # (H,W) in [0,1]
+                    binmask = _adaptive_threshold(prob, hard=0.5)
+                    # gentle cleanup + largest component
+                    binmask = cv2.morphologyEx(binmask, cv2.MORPH_OPEN, np.ones((3,3), np.uint8), iterations=1)
+                    binmask = cv2.morphologyEx(binmask, cv2.MORPH_CLOSE, np.ones((3,3), np.uint8), iterations=1)
+                    binmask = largest_component_mask(binmask, min_area_px=30)
+                    # back to ROI size {0,1}
+                    mask_roi_01 = cv2.resize(binmask, (roi.shape[1], roi.shape[0]), interpolation=cv2.INTER_NEAREST).astype(np.uint8)
+                    logging.info(f"seg prob stats: min={prob.min():.4f}, max={prob.max():.4f}, mean={prob.mean():.4f}; on={(mask_roi_01==1).sum()}")
+                except Exception as e:
+                    logging.warning(f"Segmentation failed: {e}")
+                    mask_roi_01 = None
             else:
-                mask_roi_01 = np.zeros(roi.shape[:2], dtype=np.uint8)
+                logging.info("Skipping segmentation (no model).")
 
-            # --- Measurement with oriented rect (better than boundingRect) ---
-            if mask_roi_01.any():
+            # --- Measurement ---
+            if mask_roi_01 is not None and mask_roi_01.any():
                 length_cm, breadth_cm, (box_pts, _) = measure_min_area_rect(mask_roi_01, px_per_cm)
                 surface_area_cm2 = count_area_cm2(mask_roi_01, px_per_cm)
                 anno_roi = draw_measurement_overlay(roi, mask_roi_01, box_pts, length_cm, breadth_cm)
             else:
-                # fallback to detection box if segmentation missing/empty
-                h_px = max(0, y2 - y1)
-                w_px = max(0, x2 - x1)
+                # fallback to detection-box cm
+                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()
 
-            # --- Save visuals ---
+            # --- Save visualizations ---
             out_dir = self._ensure_analysis_dir()
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
 
@@ -531,7 +519,7 @@ class AIProcessor:
             segmentation_path = None
             annotated_seg_path = None
             if mask_roi_01 is not None and mask_roi_01.any():
-                # Safe blending: blend once, then gate by mask
+                # safe masked blend (no mask kwarg to addWeighted)
                 seg_full = image_cv.copy()
                 roi_overlay = roi.copy()
                 red = np.zeros_like(roi_overlay); red[:] = (0, 0, 255)
@@ -551,7 +539,7 @@ class AIProcessor:
                 annotated_seg_path = os.path.join(out_dir, f"segmentation_annotated_{ts}.png")
                 cv2.imwrite(annotated_seg_path, anno_full)
 
-            # --- Classification (optional) ---
+            # --- Optional classification ---
             wound_type = "Unknown"
             cls_pipe = self.models_cache.get("cls")
             if cls_pipe is not None:
@@ -580,7 +568,7 @@ class AIProcessor:
             logging.error(f"Visual analysis failed: {e}", exc_info=True)
             raise
 
-    # ---------- Knowledge base and reporting (names preserved) ----------
+    # ---------- Knowledge base and reporting stay unchanged ----------
     def query_guidelines(self, query: str) -> str:
         try:
             vs = self.knowledge_base_cache.get("vector_store")