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,12 +26,6 @@ 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"
@@ -39,7 +33,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.h5"   # optional; legacy .h5 supported
+SEG_MODEL_PATH = "src/segmentation_model.h5"   # optional
 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
@@ -123,8 +117,10 @@ 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')
@@ -132,53 +128,238 @@ 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.
 """
 
-# ---------- 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]):
+
+def _vlm_infer_gpu(messages, model_id: str, max_new_tokens: int, token: Optional[str]):
     """
-    Runs entirely inside a Spaces GPU worker. Uses Med-Gemma (text-only) to draft the report.
+    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)
     """
+    import logging
     import torch
-    from transformers import pipeline
+    from typing import Optional, List
+    from transformers import (
+        AutoProcessor,
+        AutoModelForVision2Seq,
+        StoppingCriteria,
+        StoppingCriteriaList,
+    )
 
-    pipe = pipeline(
-        "image-text-to-text",
-        model="google/medgemma-4b-it",
-        torch_dtype=torch.bfloat16,
-        device="cuda",
+    # -------- 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
     )
-    out = pipe(
-        prompt,
-        max_new_tokens=max_new_tokens,
+
+    # -------- 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),
         do_sample=False,
-        temperature=0.2,
-        return_full_text=True,
+        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,
     )
-    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"
+    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()
+
 
-def generate_medgemma_report(  # kept name so callers don't change
+def generate_medgemma_report(
     patient_info: str,
     visual_results: Dict,
     guideline_context: str,
-    image_pil: Image.Image,  # kept for signature compatibility; not used by MedGemma
+    image_pil: Image.Image,
     max_new_tokens: Optional[int] = None,
 ) -> str:
     """
-    MedGemma (text-only) report generation.
-    The image is analyzed by the vision pipeline; MedGemma formats clinical guidance text.
+    MedGemma replacement using a vision-language model.
+    Loads & runs ONLY inside a GPU worker to satisfy Stateless GPU constraints.
     """
     if os.getenv("SMARTHEAL_ENABLE_VLM", "1") != "1":
         return "⚠️ VLM disabled"
 
-    # 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")
+    model_id = os.getenv("SMARTHEAL_VLM_MODEL", "bczhou/tiny-llava-v1-hf")
     max_new_tokens = max_new_tokens or int(os.getenv("SMARTHEAL_VLM_MAX_TOKENS", "600"))
 
     uprompt = SMARTHEAL_USER_PREFIX.format(
@@ -192,69 +373,27 @@ def generate_medgemma_report(  # kept name so callers don't change
         guideline_context=(guideline_context or "")[:900],
     )
 
-    # Compose a single text prompt
-    prompt = f"{SMARTHEAL_SYSTEM_PROMPT}\n\n{uprompt}\n\nAnswer:"
+    # 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},
+        ]},
+    ]
 
     try:
-        return _medgemma_generate_gpu(prompt, model_id, max_new_tokens, HF_TOKEN)
+        return _vlm_infer_gpu(messages, model_id, max_new_tokens, HF_TOKEN)
     except Exception as 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
-
+        logging.error(f"VLM call failed: {e}", exc_info=True)
+        return f"⚠️ VLM error: {e}"
 # ---------- 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 tensorflow as tf
     load_model = _import_tf_loader()
@@ -287,11 +426,11 @@ def initialize_cpu_models() -> None:
     if "seg" not in models_cache:
         try:
             if os.path.exists(SEG_MODEL_PATH):
-                m = load_segmentation_model()  # uses global path by default
-                models_cache["seg"] = m
-                th, tw = _get_model_input_hw(m, default_hw=(224, 224))
+                models_cache["seg"] = load_segmentation_model()
+                m = models_cache["seg"]
+                ishape = getattr(m, "input_shape", None)
                 oshape = getattr(m, "output_shape", None)
-                logging.info(f"✅ Segmentation model loaded (CPU) | input_hw=({th},{tw}) output_shape={oshape}")
+                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.")
@@ -509,7 +648,11 @@ def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndar
     # --- Model path ---
     if seg_model is not None:
         try:
-            th, tw = _get_model_input_hw(seg_model, default_hw=(224, 224))
+            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))
             roi_seen_path = None
             if SMARTHEAL_DEBUG:
@@ -600,7 +743,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(h_px, w_px))
+    length_px, breadth_px = (max(w_px, h_px), min(w_px, h_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)
@@ -859,6 +1002,7 @@ 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:
@@ -872,30 +1016,38 @@ 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.
 """
 
@@ -949,7 +1101,8 @@ 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)
@@ -995,7 +1148,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):
@@ -1019,4 +1172,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }