Update src/ai_processor.py

src/ai_processor.py

@@ -28,15 +28,12 @@ logging.basicConfig(
 def _log_kv(prefix: str, kv: Dict):
     logging.debug(prefix + " | " + " | ".join(f"{k}={v}" for k, v in kv.items()))
 
-# --- Optional Spaces GPU stub (harmless) ---
-try:
-    import spaces as _spaces
-    @_spaces.GPU(enable_queue=False)
-    def smartheal_gpu_stub(ping: int = 0) -> str:
-        return "ready"
-    logging.info("Registered @spaces.GPU stub (enable_queue=False).")
-except Exception:
-    pass
+# --- Spaces GPU decorator (REQUIRED) ---
+from spaces import GPU as _SPACES_GPU
+
+@_SPACES_GPU(enable_queue=False)
+def smartheal_gpu_stub(ping: int = 0) -> str:
+    return "ready"
 
 UPLOADS_DIR = "uploads"
 os.makedirs(UPLOADS_DIR, exist_ok=True)
@@ -91,49 +88,105 @@ def _import_hf_hub():
     from huggingface_hub import HfApi, HfFolder
     return HfApi, HfFolder
 
-# ---------- VLM (disabled by default) ----------
-def generate_medgemma_report(
+# ---------- SmartHeal prompts (system + user prefix) ----------
+SMARTHEAL_SYSTEM_PROMPT = """\
+You are SmartHeal Clinical Assistant, a wound-care decision-support system.
+You analyze wound photographs and brief patient context to produce careful,
+specific, guideline-informed recommendations WITHOUT diagnosing. You always:
+- Use the measurements calculated by the vision pipeline as ground truth.
+- Prefer concise, actionable steps tailored to exudate level, infection risk, and pain.
+- Flag uncertainties and red flags that need escalation to a clinician.
+- Avoid contraindicated advice; do not infer unseen comorbidities.
+- Keep under 300 words and use the requested headings exactly.
+- Tone: professional, clear, and conservative; no definitive medical claims.
+- Safety: remind the user to seek clinician review for changes or red flags.
+"""
+
+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')
+3. Treatment Plan (specific dressing choices and frequency based on exudate/infection risk)
+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.
+"""
+
+# ---------- VLM (MedGemma replaced with Qwen2-VL) ----------
+def generate_medgemma_report(  # kept name so callers don't change
     patient_info: str,
     visual_results: Dict,
     guideline_context: str,
     image_pil: Image.Image,
     max_new_tokens: Optional[int] = None,
 ) -> str:
-    if os.getenv("SMARTHEAL_ENABLE_VLM", "0") != "1":
+    """
+    MedGemma replacement using Qwen/Qwen2-VL-2B-Instruct via image-text-to-text.
+    """
+    if os.getenv("SMARTHEAL_ENABLE_VLM", "1") != "1":
         return "⚠️ VLM disabled"
+
+    model_id = os.getenv("SMARTHEAL_VLM_MODEL", "Qwen/Qwen2-VL-2B-Instruct")
+    max_new_tokens = max_new_tokens or int(os.getenv("SMARTHEAL_VLM_MAX_TOKENS", "600"))
+
     try:
         from transformers import pipeline
         pipe = pipeline(
             task="image-text-to-text",
-            model="google/medgemma-4b-it",
-            device_map=None,
+            model=model_id,
+            device_map=None,              # keep CPU by default for Spaces stability
             token=HF_TOKEN,
             trust_remote_code=True,
             model_kwargs={"low_cpu_mem_usage": 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 600, do_sample=False, temperature=0.7)
+    except Exception as e:
+        logging.error(f"❌ Could not load VLM ({model_id}): {e}")
+        return "⚠️ VLM error"
+
+    uprompt = SMARTHEAL_USER_PREFIX.format(
+        patient_info=patient_info,
+        wound_type=visual_results.get("wound_type", "Unknown"),
+        length_cm=visual_results.get("length_cm", 0),
+        breadth_cm=visual_results.get("breadth_cm", 0),
+        area_cm2=visual_results.get("surface_area_cm2", 0),
+        det_conf=float(visual_results.get("detection_confidence", 0.0)),
+        px_per_cm=visual_results.get("px_per_cm", "?"),
+        guideline_context=(guideline_context or "")[:900],
+    )
+
+    try:
+        messages = [
+            {"role": "system", "content": [{"type": "text", "text": SMARTHEAL_SYSTEM_PROMPT}]},
+            {"role": "user", "content": [
+                {"type": "image", "image": image_pil},
+                {"type": "text", "text": uprompt},
+            ]},
+        ]
+        out = pipe(text=messages, max_new_tokens=max_new_tokens, do_sample=False, temperature=0.2)
         if out and len(out) > 0:
             try:
-                return out[0]["generated_text"][-1].get("content", "").strip() or "⚠️ Empty response"
+                text = out[0]["generated_text"][-1].get("content", "")
             except Exception:
-                return (out[0].get("generated_text", "") or "").strip() or "⚠️ Empty response"
+                text = out[0].get("generated_text", "")
+            text = (text or "").strip()
+            return text if text else "⚠️ Empty response"
         return "⚠️ No output generated"
     except Exception as e:
-        logging.error(f"❌ MedGemma generation error: {e}")
+        logging.error(f"❌ VLM generation error: {e}")
         return "⚠️ VLM error"
 
+UPLOADS_DIR = "uploads"
+os.makedirs(UPLOADS_DIR, exist_ok=True)
+
 # ---------- Initialize CPU models ----------
 def load_yolo_model():
     YOLO = _import_ultralytics()
@@ -343,7 +396,7 @@ def _grabcut_refine(bgr: np.ndarray, seed01: np.ndarray, iters: int = 3) -> np.n
     seed_dil = cv2.dilate(seed01, k, iterations=1)
     gc[seed01.astype(bool)] = cv2.GC_PR_FGD
     gc[seed_dil.astype(bool)] = cv2.GC_FGD
-    gc[0, :], gc[-1, :], gc[:, 0], gc[:, -1] = cv2.GC_BGD, cv2.GC_BGD, cv2.GC_BGD, cv2.GC_BGD
+    gc[0, :], gc[-1, :], gc[:, 0], gc[:, 1] = cv2.GC_BGD, cv2.GC_BGD, cv2.GC_BGD, cv2.GC_BGD
     bgdModel = np.zeros((1, 65), np.float64)
     fgdModel = np.zeros((1, 65), np.float64)
     cv2.grabCut(bgr, gc, None, bgdModel, fgdModel, iters, cv2.GC_INIT_WITH_MASK)
@@ -811,7 +864,7 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             )
             if report and report.strip() and not report.startswith(("⚠️", "❌")):
                 return report
-            logging.warning("MedGemma unavailable/invalid; using fallback.")
+            logging.warning("VLM unavailable/invalid; using fallback.")
             return self._generate_fallback_report(patient_info, visual_results, guideline_context)
         except Exception as e:
             logging.error(f"Report generation failed: {e}")
@@ -917,4 +970,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }