Update src/ai_processor.py

src/ai_processor.py

@@ -16,7 +16,7 @@ import cv2
 import numpy as np
 from PIL import Image
 from PIL.ExifTags import TAGS
-import spaces
+
 # --- Logging config ---
 logging.basicConfig(
     level=getattr(logging, LOGLEVEL, logging.INFO),
@@ -26,6 +26,12 @@ logging.basicConfig(
 def _log_kv(prefix: str, kv: Dict):
     logging.debug(prefix + " | " + " | ".join(f"{k}={v}" for k, v in kv.items()))
 
+# --- Spaces GPU decorator (REQUIRED) ---
+from spaces import GPU as _SPACES_GPU
+
+@_SPACES_GPU(enable_queue=True)
+def smartheal_gpu_stub(ping: int = 0) -> str:
+    return "ready"
 
 # ---- Paths / constants ----
 UPLOADS_DIR = "uploads"
@@ -33,7 +39,7 @@ os.makedirs(UPLOADS_DIR, exist_ok=True)
 
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 YOLO_MODEL_PATH = "src/best.pt"
-SEG_MODEL_PATH = "src/segmentation_model_fixed.h5"   # optional
+SEG_MODEL_PATH = "src/segmentation_model.h5"   # optional; legacy .h5 supported
 GUIDELINE_PDFS = ["src/eHealth in Wound Care.pdf", "src/IWGDF Guideline.pdf", "src/evaluation.pdf"]
 DATASET_ID = "SmartHeal/wound-image-uploads"
 DEFAULT_PX_PER_CM = 38.0
@@ -117,10 +123,8 @@ SMARTHEAL_USER_PREFIX = """\
 Patient: {patient_info}
 Visual findings: type={wound_type}, size={length_cm}x{breadth_cm} cm, area={area_cm2} cm^2,
 detection_conf={det_conf:.2f}, calibration={px_per_cm} px/cm.
-
 Guideline context (snippets you can draw principles from; do not quote at length):
 {guideline_context}
-
 Write a structured answer with these headings exactly:
 1. Clinical Summary (max 4 bullet points)
 2. Likely Stage/Type (if uncertain, say 'uncertain')
@@ -128,238 +132,53 @@ Write a structured answer with these headings exactly:
 4. Red Flags (what to escalate and when)
 5. Follow-up Cadence (days)
 6. Notes (assumptions/uncertainties)
-
 Keep to 220–300 words. Do NOT provide diagnosis. Avoid contraindicated advice.
 """
 
-
-def _vlm_infer_gpu(messages, model_id: str, max_new_tokens: int, token: Optional[str]):
+# ---------- MedGemma-only text generator ----------
+@_SPACES_GPU(enable_queue=True)
+def _medgemma_generate_gpu(prompt: str, model_id: str, max_new_tokens: int, token: Optional[str]):
     """
-    Runs entirely inside a Spaces GPU worker. It's the ONLY place we allow CUDA init.
-    Safe for:
-      - CUDA device selection (no 'Invalid device id')
-      - BF16/FP16 choice via compute capability
-      - LLaVA processors with patch_size=None
-      - Processors WITHOUT a chat template (fallback to plain/LLaVA-style prompt)
+    Runs entirely inside a Spaces GPU worker. Uses Med-Gemma (text-only) to draft the report.
     """
-    import logging
     import torch
-    from typing import Optional, List
-    from transformers import (
-        AutoProcessor,
-        AutoModelForVision2Seq,
-        StoppingCriteria,
-        StoppingCriteriaList,
-    )
+    from transformers import pipeline
 
-    # -------- Device & dtype (robust) --------
-    def _pick_device_and_dtype():
-        if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
-            logging.warning("CUDA not available; using CPU.")
-            return "cpu", torch.float32
-        idx = 0
-        try:
-            torch.cuda.set_device(idx)
-        except Exception as e:
-            logging.warning(f"torch.cuda.set_device({idx}) failed: {e}; falling back to CPU.")
-            return "cpu", torch.float32
-        device = f"cuda:{idx}"
-        try:
-            props = torch.cuda.get_device_properties(idx)
-            cc = props.major * 10 + props.minor
-            dtype = torch.bfloat16 if cc >= 80 else torch.float16
-        except Exception as e:
-            logging.warning(f"Could not query CUDA props: {e}; defaulting to float16.")
-            dtype = torch.float16
-        return device, dtype
-
-    device, torch_dtype = _pick_device_and_dtype()
-
-    # -------- Load model & processor --------
-    model = AutoModelForVision2Seq.from_pretrained(
-        model_id,
-        torch_dtype=torch_dtype,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True,
-        token=token,
-    ).to(device)
-    model.eval()
-
-    processor = AutoProcessor.from_pretrained(
-        model_id, trust_remote_code=True, token=token
+    pipe = pipeline(
+        "image-text-to-text",
+        model="google/medgemma-4b-it",
+        torch_dtype=torch.bfloat16,
+        device="cuda",
     )
-
-    # -------- Extract image & text --------
-    image_obj = None
-    text_prompt = ""
-    for m in messages:
-        if m.get("role") == "user":
-            for c in m.get("content", []):
-                if c.get("type") == "image":
-                    image_obj = c.get("image")
-                elif c.get("type") == "text":
-                    text_prompt = c.get("text", "")
-            break
-    if image_obj is None:
-        raise ValueError("No image found in messages for VLM inference.")
-
-    # -------- Normalize image to PIL --------
-    from PIL import Image
-    import numpy as np
-    def _to_pil(x):
-        if isinstance(x, Image.Image):
-            return x.convert("RGB")
-        if isinstance(x, str):
-            return Image.open(x).convert("RGB")
-        if isinstance(x, np.ndarray):
-            if x.ndim == 2:
-                x = np.stack([x]*3, axis=-1)
-            if x.dtype != np.uint8:
-                x = x.astype(np.uint8)
-            return Image.fromarray(x, "RGB")
-        if hasattr(x, "read"):
-            return Image.open(x).convert("RGB")
-        raise TypeError(f"Unsupported image type: {type(x)}")
-    image_pil = _to_pil(image_obj)
-
-    # -------- Ensure patch_size for LLaVA processors --------
-    def _ensure_patch_size(proc, mdl):
-        ps = getattr(proc, "patch_size", None)
-        if not ps:
-            candidates = [
-                getattr(getattr(mdl, "vision_tower", None), "config", None),
-                getattr(mdl.config, "vision_config", None),
-                getattr(proc, "image_processor", None),
-                getattr(getattr(proc, "image_processor", None), "config", None),
-            ]
-            for obj in candidates:
-                if obj is None:
-                    continue
-                maybe = getattr(obj, "patch_size", None)
-                if maybe:
-                    ps = int(maybe); break
-        if not ps:
-            ps = 14  # safe default for ViT-L/14-style
-        try:
-            setattr(proc, "patch_size", ps)
-        except Exception:
-            pass
-        return ps
-    _ensure_patch_size(processor, model)
-
-    # -------- Build text (chat-template only if it truly exists) --------
-    # Some processors expose apply_chat_template but tokenizer has no template → ValueError. Guard it.
-    tokenizer = getattr(processor, "tokenizer", None)
-    has_template = bool(getattr(tokenizer, "chat_template", None))
-    used_chat_template = False
-
-    def _looks_like_llava():
-        name = processor.__class__.__name__.lower()
-        mid = (model_id or "").lower()
-        return ("llava" in name) or ("llava" in mid)
-
-    if hasattr(processor, "apply_chat_template") and has_template:
-        try:
-            chat = [{
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": image_pil},
-                    {"type": "text", "text": text_prompt or "Describe the image."},
-                ],
-            }]
-            text_for_model = processor.apply_chat_template(
-                chat, add_generation_prompt=True, tokenize=False
-            )
-            used_chat_template = True
-        except Exception as e:
-            logging.info(f"No usable chat template ({e}); falling back to plain prompt.")
-            text_for_model = (
-                f"USER: <image>\n{text_prompt or 'Describe the image.'}\nASSISTANT:"
-                if _looks_like_llava() else (text_prompt or "Describe the image.")
-            )
-    else:
-        text_for_model = (
-            f"USER: <image>\n{text_prompt or 'Describe the image.'}\nASSISTANT:"
-            if _looks_like_llava() else (text_prompt or "Describe the image.")
-        )
-
-    # -------- Tokenize --------
-    inputs = processor(
-        text=[text_for_model],
-        images=[image_pil],
-        return_tensors="pt",
-        padding=True,
-    ).to(device)
-
-    # -------- Stopping criteria --------
-    class EosTokenCriteria(StoppingCriteria):
-        def __init__(self, eos_token_ids: List[int]):
-            import torch as _t
-            self.eos = _t.tensor(eos_token_ids, dtype=_t.long)
-        def __call__(self, input_ids, scores, **kwargs) -> bool:
-            import torch as _t
-            last_tok = input_ids[:, -1]
-            return _t.isin(last_tok, self.eos.to(last_tok.device)).any().item()
-
-    eos_ids: List[int] = []
-    if tokenizer is not None:
-        for attr in ("eos_token_id", "eot_token_id"):
-            v = getattr(tokenizer, attr, None)
-            if v is None: continue
-            eos_ids.extend([v] if isinstance(v, int) else list(v))
-    if not eos_ids:
-        cfg = getattr(model, "generation_config", None)
-        if cfg and getattr(cfg, "eos_token_id", None) is not None:
-            eos_ids = [cfg.eos_token_id]
-        else:
-            eos_ids = [2]
-    stopping_criteria = StoppingCriteriaList([EosTokenCriteria(eos_ids)])
-
-    if tokenizer is not None and getattr(tokenizer, "pad_token_id", None) is None:
-        try: tokenizer.pad_token_id = eos_ids[0]
-        except Exception: pass
-
-    # -------- Generate --------
-    gen_kwargs = dict(
-        max_new_tokens=int(max_new_tokens or 256),
+    out = pipe(
+        prompt,
+        max_new_tokens=max_new_tokens,
         do_sample=False,
-        stopping_criteria=stopping_criteria,
-        eos_token_id=eos_ids[0] if eos_ids else None,
-        pad_token_id=getattr(tokenizer, "pad_token_id", None) if tokenizer else None,
+        temperature=0.2,
+        return_full_text=True,
     )
-    with torch.inference_mode():
-        out = model.generate(**inputs, **gen_kwargs)
-
-    # -------- Decode --------
-    seq = out[0]
-    if "input_ids" in inputs:
-        cut = inputs["input_ids"].shape[-1]
-        seq = seq[cut:]
-    if tokenizer is not None:
-        text_out = tokenizer.decode(seq, skip_special_tokens=True)
-    elif hasattr(processor, "batch_decode"):
-        text_out = processor.batch_decode(seq.unsqueeze(0), skip_special_tokens=True)[0]
-    else:
-        text_out = str(seq.tolist())
-
-    return text_out.strip()
-
+    text = (out[0].get("generated_text") if isinstance(out, list) else out).strip()
+    # Remove the prompt echo if present
+    if text.startswith(prompt):
+        text = text[len(prompt):].lstrip()
+    return text or "⚠️ Empty response"
 
-def generate_medgemma_report(
+def generate_medgemma_report(  # kept name so callers don't change
     patient_info: str,
     visual_results: Dict,
     guideline_context: str,
-    image_pil: Image.Image,
+    image_pil: Image.Image,  # kept for signature compatibility; not used by MedGemma
     max_new_tokens: Optional[int] = None,
 ) -> str:
     """
-    MedGemma replacement using a vision-language model.
-    Loads & runs ONLY inside a GPU worker to satisfy Stateless GPU constraints.
+    MedGemma (text-only) report generation.
+    The image is analyzed by the vision pipeline; MedGemma formats clinical guidance text.
     """
     if os.getenv("SMARTHEAL_ENABLE_VLM", "1") != "1":
         return "⚠️ VLM disabled"
 
-    model_id = os.getenv("SMARTHEAL_VLM_MODEL", "bczhou/tiny-llava-v1-hf")
+    # Default to a public Med-Gemma instruction-tuned model (update via env if you have access to another).
+    model_id = os.getenv("SMARTHEAL_MEDGEMMA_MODEL", "google/med-gemma-2-2b-it")
     max_new_tokens = max_new_tokens or int(os.getenv("SMARTHEAL_VLM_MAX_TOKENS", "600"))
 
     uprompt = SMARTHEAL_USER_PREFIX.format(
@@ -373,28 +192,69 @@ def generate_medgemma_report(
         guideline_context=(guideline_context or "")[:900],
     )
 
-    # The `messages` structure is passed to the verified `_vlm_infer_gpu` function
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": SMARTHEAL_SYSTEM_PROMPT}]},
-        {"role": "user", "content": [
-            {"type": "image", "image": image_pil},
-            {"type": "text",  "text": uprompt},
-        ]},
-    ]
+    # Compose a single text prompt
+    prompt = f"{SMARTHEAL_SYSTEM_PROMPT}\n\n{uprompt}\n\nAnswer:"
 
     try:
-        return _vlm_infer_gpu(messages, model_id, max_new_tokens, HF_TOKEN)
+        return _medgemma_generate_gpu(prompt, model_id, max_new_tokens, HF_TOKEN)
     except Exception as e:
-        logging.error(f"VLM call failed: {e}", exc_info=True)
-        return f"⚠️ VLM error: {e}"
+        logging.error(f"MedGemma call failed: {e}")
+        return "⚠️ VLM error"
+
+# ---------- Input-shape helpers (avoid `.as_list()` on strings) ----------
+def _shape_to_hw(shape) -> Tuple[Optional[int], Optional[int]]:
+    try:
+        if hasattr(shape, "as_list"):
+            shape = shape.as_list()
+    except Exception:
+        pass
+    if isinstance(shape, (tuple, list)):
+        if len(shape) == 4:   # (None, H, W, C)
+            H, W = shape[1], shape[2]
+        elif len(shape) == 3: # (H, W, C)
+            H, W = shape[0], shape[1]
+        else:
+            return (None, None)
+        try: H = int(H) if (H is not None and str(H).lower() != "none") else None
+        except Exception: H = None
+        try: W = int(W) if (W is not None and str(W).lower() != "none") else None
+        except Exception: W = None
+        return (H, W)
+    return (None, None)
+
+def _get_model_input_hw(model, default_hw: Tuple[int, int] = (224, 224)) -> Tuple[int, int]:
+    H, W = _shape_to_hw(getattr(model, "input_shape", None))
+    if H and W:
+        return H, W
+    try:
+        inputs = getattr(model, "inputs", None)
+        if inputs:
+            H, W = _shape_to_hw(inputs[0].shape)
+            if H and W:
+                return H, W
+    except Exception:
+        pass
+    try:
+        cfg = model.get_config() if hasattr(model, "get_config") else None
+        if isinstance(cfg, dict):
+            for layer in cfg.get("layers", []):
+                conf = (layer or {}).get("config", {})
+                cand = conf.get("batch_input_shape") or conf.get("batch_shape")
+                H, W = _shape_to_hw(cand)
+                if H and W:
+                    return H, W
+    except Exception:
+        pass
+    logging.warning(f"Could not resolve model input shape; using default {default_hw}.")
+    return default_hw
+
 # ---------- Initialize CPU models ----------
 def load_yolo_model():
     YOLO = _import_ultralytics()
-    # Construct model with CUDA masked to avoid auto-selecting cuda:0
     with _no_cuda_env():
         model = YOLO(YOLO_MODEL_PATH)
     return model
-    
+
 def load_segmentation_model():
     import os; os.environ.setdefault("KERAS_BACKEND","tensorflow")
     import tensorflow as tf; tf.config.set_visible_devices([], "GPU")
@@ -428,11 +288,11 @@ def initialize_cpu_models() -> None:
     if "seg" not in models_cache:
         try:
             if os.path.exists(SEG_MODEL_PATH):
-                models_cache["seg"] = load_segmentation_model()
-                m = models_cache["seg"]
-                ishape = getattr(m, "input_shape", None)
+                m = load_segmentation_model()  # uses global path by default
+                models_cache["seg"] = m
+                th, tw = _get_model_input_hw(m, default_hw=(224, 224))
                 oshape = getattr(m, "output_shape", None)
-                logging.info(f"✅ Segmentation model loaded (CPU) | input_shape={ishape} output_shape={oshape}")
+                logging.info(f"✅ Segmentation model loaded (CPU) | input_hw=({th},{tw}) output_shape={oshape}")
             else:
                 models_cache["seg"] = None
                 logging.warning("Segmentation model file missing; skipping.")
@@ -650,11 +510,7 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
     # --- Model path ---
     if seg_model is not None:
         try:
-            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])
-
+            th, tw = _get_model_input_hw(seg_model, default_hw=(224, 224))
             x = _preprocess_for_seg(image_bgr, (th, tw))
             roi_seen_path = None
             if SMARTHEAL_DEBUG:
@@ -745,7 +601,7 @@ def measure_min_area_rect(mask01: np.ndarray, px_per_cm: float) -> Tuple[float,
     cnt = max(contours, key=cv2.contourArea)
     rect = cv2.minAreaRect(cnt)
     (w_px, h_px) = rect[1]
-    length_px, breadth_px = (max(w_px, h_px), min(w_px, h_px))
+    length_px, breadth_px = (max(w_px, h_px), min(h_px, w_px))
     length_cm = round(length_px / max(px_per_cm, 1e-6), 2)
     breadth_cm = round(breadth_px / max(px_per_cm, 1e-6), 2)
     box = cv2.boxPoints(rect).astype(int)
@@ -1004,7 +860,6 @@ class AIProcessor:
             if not vs:
                 return "Knowledge base is not available."
             retriever = vs.as_retriever(search_kwargs={"k": 5})
-            # Modern API (avoid get_relevant_documents deprecation)
             docs = retriever.invoke(query)
             lines: List[str] = []
             for d in docs:
@@ -1018,38 +873,30 @@ class AIProcessor:
 
     def _generate_fallback_report(self, patient_info: str, visual_results: Dict, guideline_context: str) -> str:
         return f"""# 🩺 SmartHeal AI - Comprehensive Wound Analysis Report
-
 ## 📋 Patient Information
 {patient_info}
-
 ## 🔍 Visual Analysis Results
 - **Wound Type**: {visual_results.get('wound_type', 'Unknown')}
 - **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('px_per_cm','?')} px/cm ({(visual_results.get('calibration_meta') or {}).get('used','default')})
-
 ## 📊 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')}
 - **Annotated**: {visual_results.get('segmentation_annotated_path', 'N/A')}
-
 ## 🎯 Clinical Summary
 Automated analysis provides quantitative measurements; verify via clinical examination.
-
 ## 💊 Recommendations
 - Cleanse wound gently; select dressing per exudate/infection risk
 - Debride necrotic tissue if indicated (clinical decision)
 - Document with serial photos and measurements
-
 ## 📅 Monitoring
 - Daily in week 1, then every 2–3 days (or as indicated)
 - Weekly progress review
-
 ## 📚 Guideline Context
 {(guideline_context or '')[:800]}{"..." if guideline_context and len(guideline_context) > 800 else ''}
-
 **Disclaimer:** Automated, for decision support only. Verify clinically.
 """
 
@@ -1103,8 +950,7 @@ Automated analysis provides quantitative measurements; verify via clinical exami
         except Exception as e:
             logging.error(f"Failed to save/commit image: {e}")
             return ""
-            
-    
+
     def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: Dict) -> Dict:
         try:
             saved_path = self.save_and_commit_image(image_pil)
@@ -1150,7 +996,7 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "saved_image_path": None,
                 "guideline_context": "",
             }
-    
+
     def analyze_wound(self, image, questionnaire_data: Dict) -> Dict:
         try:
             if isinstance(image, str):
@@ -1174,4 +1020,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }