Update src/ai_processor.py

src/ai_processor.py

@@ -1,11 +1,11 @@
 # smartheal_ai_processor.py
 # Preserves ALL original class/function names.
-# Same logic you confirmed on Colab:
-#  - Uses segmentation_model.h5 first (fallback to KMeans)
-#  - Safe overlay (no 'mask' kwarg in addWeighted)
-#  - CPU-only by default to avoid ZeroGPU cuda probe
-#  - Registers a harmless @spaces.GPU stub (enable_queue=False) to silence
-#    "No @spaces.GPU function detected during startup" without starting a GPU worker.
+# 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
 
 import os
 import time
@@ -13,28 +13,28 @@ import logging
 from datetime import datetime
 from typing import Optional, Dict, List, Tuple
 
-# Quieter tokenizer + default CPU
+# Quiet tokenizers; default to CPU for safety on ZeroGPU/Spaces
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # keep torch/TF on CPU
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
 
 import cv2
 import numpy as np
 from PIL import Image
 from PIL.ExifTags import TAGS
 
-# --- Register a non-queue GPU stub so Spaces detects @spaces.GPU but doesn't start a worker ---
+# --- Optional: register a harmless @spaces.GPU-decorated stub to silence startup warning ---
 try:
     import spaces as _spaces
 
-    @_spaces.GPU(enable_queue=False)  # NOTE: no queue, so ZeroGPU worker is not launched
-    def _spaces_gpu_stub(ping: int = 0) -> str:
-        """Harmless stub to satisfy Spaces startup scan without touching CUDA."""
+    @_spaces.GPU(enable_queue=False)  # not queued -> won't start a ZeroGPU worker
+    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:
-    _spaces = None
-    logging.info("No 'spaces' module or stub registration failed: %s", _e)
+    # It's fine if 'spaces' isn't available locally.
+    pass
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
@@ -60,7 +60,7 @@ def _import_ultralytics():
 def _import_tf_loader():
     import tensorflow as tf
     try:
-        tf.config.set_visible_devices([], "GPU")  # force TF CPU
+        tf.config.set_visible_devices([], "GPU")  # keep TF on CPU
     except Exception:
         pass
     from tensorflow.keras.models import load_model
@@ -86,7 +86,7 @@ def _import_hf_hub():
     from huggingface_hub import HfApi, HfFolder
     return HfApi, HfFolder
 
-# ---------- LLM report: CPU-only path (safe on ZeroGPU) ----------
+# ---------- LLM report (OFF by default; enable with SMARTHEAL_ENABLE_VLM=1) ----------
 def generate_medgemma_report(
     patient_info: str,
     visual_results: Dict,
@@ -95,16 +95,21 @@ def generate_medgemma_report(
     max_new_tokens: Optional[int] = None,
 ) -> str:
     """
-    CPU-only MedGemma call (safe on Spaces/ZeroGPU). If it fails, fallback text is provided by caller.
+    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(
-            "image-text-to-text",
+            task="image-text-to-text",
             model="google/medgemma-4b-it",
-            device_map=None,   # CPU
+            device_map=None,               # CPU
             token=HF_TOKEN,
-            model_kwargs={"low_cpu_mem_usage": True, "use_cache": True},
+            trust_remote_code=True,
+            model_kwargs={"low_cpu_mem_usage": True},  # avoid 'use_cache' arg mismatch
         )
 
         prompt = (
@@ -121,15 +126,7 @@ def generate_medgemma_report(
             {"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")
-
+        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:
                 return out[0]["generated_text"][-1].get("content", "").strip() or "⚠️ Empty response"
@@ -138,7 +135,7 @@ def generate_medgemma_report(
         return "⚠️ No output generated"
     except Exception as e:
         logging.error(f"❌ MedGemma generation error: {e}")
-        return "⚠️ GPU/LLM worker unavailable"
+        return "⚠️ VLM error"
 
 # ---------- Initialize CPU models ----------
 def load_yolo_model():
@@ -461,14 +458,15 @@ class AIProcessor:
                 surface_area_cm2 = count_area_cm2(mask01, px_per_cm)
                 anno_roi = draw_measurement_overlay(roi, mask01, box_pts, length_cm, breadth_cm)
             else:
-                # fallback to detection box
+                # 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()
+                box_pts = None
 
-            # --- Save visualizations ---
+            # --- Save visualizations (ALWAYS create a segmentation image) ---
             out_dir = self._ensure_analysis_dir()
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
 
@@ -480,29 +478,31 @@ class AIProcessor:
             detection_path = os.path.join(out_dir, f"detection_{ts}.png")
             cv2.imwrite(detection_path, det_vis)
 
-            segmentation_path = None
-            annotated_seg_path = None
+            # 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():
-                # 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()
-                red = np.zeros_like(roi); red[:] = (0, 0, 255)
-                blended = cv2.addWeighted(roi, 1.0, red, 0.3, 0)
+                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))
-                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)
+            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)
+
+            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)
-                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)
+            # 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")
+            cv2.imwrite(annotated_seg_path, anno_full)
 
             # --- Optional classification ---
             wound_type = "Unknown"
@@ -525,8 +525,9 @@ 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,
+                "segmentation_image_path": segmentation_path,            # always present
                 "segmentation_annotated_path": annotated_seg_path,
+                "roi_mask_path": roi_mask_path,                          # helpful for debugging
                 "original_image_path": original_path,
             }
         except Exception as e: