Update src/ai_processor.py

src/ai_processor.py

@@ -1,10 +1,9 @@
 # smartheal_ai_processor.py
 # Preserves ALL original class/function names.
-# Changes:
-#  - Adds segment_wound(image) with your logic (+ KMeans fallback)
-#  - perform_visual_analysis() now calls segment_wound() for mask
-#  - Safe overlay (no mask kwarg in addWeighted)
-#  - Conditional @spaces.GPU to avoid cudaGetDeviceCount crash
+# Same logic as your Colab run:
+#  - Uses segmentation_model.h5 if present (fallback to KMeans)
+#  - Safe overlay (no 'mask' kwarg in addWeighted)
+#  - CPU-only by default (no CUDA probe). Optional Spaces GPU is opt-in.
 
 import os
 import time
@@ -14,12 +13,31 @@ from typing import Optional, Dict, List, Tuple
 
 # Quiet HF tokenizers fork warning
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+# Default to CPU-only to match Colab logic
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
 
 import cv2
 import numpy as np
 from PIL import Image
 from PIL.ExifTags import TAGS
-import spaces
+
+# --- Optional Spaces GPU (explicit opt-in) ---
+ENABLE_SPACES_GPU = os.getenv("ENABLE_SPACES_GPU", "0") == "1"
+ALLOW_CUDA_PROBE = os.getenv("ALLOW_CUDA_PROBE", "0") == "1"  # leave "0" for ZeroGPU safety
+
+try:
+    import spaces as _spaces
+except Exception:
+    _spaces = None
+
+def _cuda_available() -> bool:
+    if not ALLOW_CUDA_PROBE:
+        return False
+    try:
+        import torch
+        return bool(getattr(torch, "cuda", None)) and torch.cuda.is_available()
+    except Exception:
+        return False
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
@@ -44,7 +62,10 @@ def _import_ultralytics():
 
 def _import_tf_loader():
     import tensorflow as tf
-    tf.config.set_visible_devices([], "GPU")  # force TF CPU
+    try:
+        tf.config.set_visible_devices([], "GPU")  # force TF CPU
+    except Exception:
+        pass
     from tensorflow.keras.models import load_model
     return load_model
 
@@ -68,16 +89,8 @@ def _import_hf_hub():
     from huggingface_hub import HfApi, HfFolder
     return HfApi, HfFolder
 
-# ---------- Conditional Spaces GPU wrapper ----------
-def _cuda_available() -> bool:
-    try:
-        import torch
-        return bool(getattr(torch, "cuda", None)) and torch.cuda.is_available()
-    except Exception:
-        return False
-        
-@spaces.GPU(enable_queue=True, duration=90)
-def _generate_medgemma_report_core(
+# ---------- LLM report: CPU by default; optional Spaces GPU if enabled ----------
+def _generate_medgemma_report_cpu(
     patient_info: str,
     visual_results: Dict,
     guideline_context: str,
@@ -89,7 +102,7 @@ def _generate_medgemma_report_core(
         pipe = pipeline(
             "image-text-to-text",
             model="google/medgemma-4b-it",
-            device_map="auto" if _cuda_available() else None,
+            device_map=None,  # CPU
             token=HF_TOKEN,
             model_kwargs={"low_cpu_mem_usage": True, "use_cache": True},
         )
@@ -127,28 +140,9 @@ def _generate_medgemma_report_core(
         logging.error(f"❌ MedGemma generation error: {e}")
         return "⚠️ GPU/LLM worker unavailable"
 
-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:
-            return _generate_medgemma_report_core(patient_info, visual_results, guideline_context, image_pil, max_new_tokens)
-except Exception:
+# Optional GPU path if you *explicitly* enable it and the env supports it
+if ENABLE_SPACES_GPU and _spaces is not None:
+    @_spaces.GPU(enable_queue=True, duration=90)
     def generate_medgemma_report(
         patient_info: str,
         visual_results: Dict,
@@ -156,7 +150,53 @@ except Exception:
         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)
+        # Even here, avoid probing CUDA unless allowed; device_map="auto" if we trust the env
+        try:
+            from transformers import pipeline
+            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},
+            ]}]
+            out = pipe(
+                text=messages,
+                max_new_tokens=max_new_tokens or 800,
+                do_sample=False,
+                temperature=0.7,
+            )
+            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 (GPU path) error: {e}")
+            return _generate_medgemma_report_cpu(patient_info, visual_results, guideline_context, image_pil, max_new_tokens)
+else:
+    # CPU default (Colab-like behavior)
+    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_cpu(patient_info, visual_results, guideline_context, image_pil, max_new_tokens)
 
 # ---------- Initialize CPU models ----------
 def load_yolo_model():
@@ -304,7 +344,7 @@ 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
 
-# ---------- Your requested segmentation logic ----------
+# ---------- Segmentation (model-first, KMeans fallback) ----------
 def segment_wound(image: np.ndarray) -> np.ndarray:
     """
     Segments wound from a preprocessed ROI image, with a fallback to KMeans if the model fails.
@@ -314,39 +354,35 @@ def segment_wound(image: np.ndarray) -> np.ndarray:
 
     if segmentation_model is not None:
         try:
-            input_size = getattr(segmentation_model, "input_shape", None)
-            if input_size is None or len(input_size) < 3:
-                raise ValueError(f"Bad seg input_shape: {input_size}")
-            H, W = int(input_size[1]), int(input_size[2])  # (None,H,W,C)
+            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))                    # cv2 takes (W,H)
+            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):
+            if isinstance(prediction, (list, tuple)):
                 prediction = prediction[0]
             # squeeze batch dim if present
-            prediction = prediction[0] if prediction.ndim >= 3 else prediction
+            prediction = prediction[0] if getattr(prediction, "ndim", 0) >= 3 else prediction
 
-            # prediction can be (H,W,1) or (H,W)
-            pred2d = prediction.squeeze()
-            mask_prob = cv2.resize(pred2d, (image.shape[1], image.shape[0]))  # back to ROI size
+            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
-            if mask.max() == 0:
-                logging.info("Seg model returned empty mask at 0.5 — keeping as-is (KMeans fallback will handle if needed).")
             return mask.astype(np.uint8)
         except Exception as e:
             logging.warning(f"⚠️ Segmentation model prediction failed: {e}. Falling back to KMeans.")
 
-    # --- Fallback: color clustering (KMeans, k=2) ---
+    # --- Fallback: color clustering (KMeans, k=2), pick 'reddest' cluster in Lab a* ---
     Z = image.reshape((-1, 3)).astype(np.float32)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
-    _K = 2
-    _, labels, centers = cv2.kmeans(Z, _K, None, criteria, 5, cv2.KMEANS_PP_CENTERS)
-    centers = centers.astype(np.uint8).reshape(1, _K, 3)
-    centers_lab = cv2.cvtColor(centers, cv2.COLOR_BGR2LAB)[0]
-    wound_idx = int(np.argmax(centers_lab[:, 1]))  # reddest cluster (a* channel)
+    _, 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)
 
@@ -387,6 +423,7 @@ def draw_measurement_overlay(
     thickness: int = 2
 ) -> np.ndarray:
     overlay = base_bgr.copy()
+    # Safe masked blend (OpenCV addWeighted has no 'mask' kwarg)
     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")
@@ -417,8 +454,8 @@ def draw_measurement_overlay(
 
         draw_arrow(overlay, mids[long_pair[0]], mids[long_pair[1]])
         draw_arrow(overlay, mids[short_pair[0]], mids[short_pair[1]])
-        put_label(f"{length_cm:.2f} cm", mids[long_pair[0]])
-        put_label(f"{breadth_cm:.2f} cm", mids[short_pair[0]])
+        put_label(f"Length: {length_cm:.2f} cm", mids[long_pair[0]])
+        put_label(f"Breadth: {breadth_cm:.2f} cm", mids[short_pair[0]])
     return overlay
 
 # ---------- AI PROCESSOR ----------
@@ -449,9 +486,12 @@ class AIProcessor:
             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
-                raise gr.Error("No wound could be detected.")
+            if (not results) or (not getattr(results[0], "boxes", None)) or (len(results[0].boxes) == 0):
+                try:
+                    import gradio as gr
+                    raise gr.Error("No wound could be detected.")
+                except Exception:
+                    raise RuntimeError("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]
@@ -459,16 +499,19 @@ 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.")
+                try:
+                    import gradio as gr
+                    raise gr.Error("Detected ROI is empty.")
+                except Exception:
+                    raise RuntimeError("Detected ROI is empty.")
 
-            # --- Segmentation (your logic + fallback) ---
+            # --- Segmentation (model-first + KMeans fallback) ---
             mask_u8_255 = segment_wound(roi)               # 0..255
-            # Clean up & keep largest component (in 0/1)
             mask01 = (mask_u8_255 > 127).astype(np.uint8)
-            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)
+            if mask01.any():
+                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)
 
             # --- Measurement ---
             if mask01.any():
@@ -482,6 +525,7 @@ class AIProcessor:
                 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()
+                box_pts = None
 
             # --- Save visualizations ---
             out_dir = self._ensure_analysis_dir()
@@ -498,8 +542,12 @@ class AIProcessor:
             segmentation_path = None
             annotated_seg_path = None
             if mask01.any():
+                # Raw mask (ROI size)
+                mask_path = os.path.join(out_dir, f"segmentation_mask_{ts}.png")
+                cv2.imwrite(mask_path, (mask01 * 255).astype(np.uint8))
+
+                # Segmentation overlay (paste back to full image)
                 seg_full = image_cv.copy()
-                # safe masked blend (no mask kwarg)
                 red = np.zeros_like(roi); red[:] = (0, 0, 255)
                 blended = cv2.addWeighted(roi, 1.0, red, 0.3, 0)
                 m3 = np.dstack([mask01 * 255] * 3).astype("uint8")
@@ -509,6 +557,7 @@ class AIProcessor:
                 segmentation_path = os.path.join(out_dir, f"segmentation_{ts}.png")
                 cv2.imwrite(segmentation_path, seg_full)
 
+                # Annotated (arrows + labels)
                 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")
@@ -722,4 +771,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }