Update src/ai_processor.py

src/ai_processor.py

@@ -1,435 +1,58 @@
-# smartheal_ai_processor.py
-# Verbose, instrumented version — preserves public class/function names
-# Turn on deep logging: export LOGLEVEL=DEBUG SMARTHEAL_DEBUG=1
-
+# src/ai_processor.py
 import os
 import logging
-from datetime import datetime
-from typing import Optional, Dict, List, Tuple
-
-# ---- Environment defaults (mask CUDA in main process) ----
-os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # ensure main never touches CUDA
-LOGLEVEL = os.getenv("LOGLEVEL", "INFO").upper()
-SMARTHEAL_DEBUG = os.getenv("SMARTHEAL_DEBUG", "0") == "1"
-
 import cv2
 import numpy as np
 from PIL import Image
-from PIL.ExifTags import TAGS
-
-# --- Logging config ---
-logging.basicConfig(
-    level=getattr(logging, LOGLEVEL, logging.INFO),
-    format="%(asctime)s - %(levelname)s - %(message)s",
-)
-
-def _log_kv(prefix: str, kv: Dict):
-    logging.debug(prefix + " | " + " | ".join(f"{k}={v}" for k, v in kv.items()))
-
-# --- Spaces GPU (non-optional) ---
-import spaces  # required; do not stub/optionalize
-
-UPLOADS_DIR = "uploads"
-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
-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
-PX_PER_CM_MIN, PX_PER_CM_MAX = 5.0, 1200.0
-
-# Segmentation preprocessing knobs
-SEG_EXPECTS_RGB = os.getenv("SEG_EXPECTS_RGB", "1") == "1"  # most TF models trained on RGB
-SEG_NORM = os.getenv("SEG_NORM", "0to1")                    # "0to1" | "imagenet"
-SEG_THRESH = float(os.getenv("SEG_THRESH", "0.5"))
-
-models_cache: Dict[str, object] = {}
-knowledge_base_cache: Dict[str, object] = {}
-
-# ---------- Lazy imports ----------
-def _import_ultralytics():
-    from ultralytics import YOLO
-    return YOLO
-
-def _import_tf_loader():
-    import tensorflow as tf
-    try:
-        tf.config.set_visible_devices([], "GPU")  # keep TF on CPU
-    except Exception:
-        pass
-    from tensorflow.keras.models import load_model
-    return load_model
-
-def _import_hf_cls():
-    from transformers import pipeline
-    return pipeline
-
-def _import_embeddings():
-    # updated per LangChain deprecations
-    from langchain_huggingface import HuggingFaceEmbeddings
-    return HuggingFaceEmbeddings
-
-def _import_langchain_pdf():
-    from langchain_community.document_loaders import PyPDFLoader
-    return PyPDFLoader
-
-def _import_langchain_faiss():
-    from langchain_community.vectorstores import FAISS
-    return FAISS
-
-def _import_hf_hub():
-    from huggingface_hub import HfApi, HfFolder
-    return HfApi, HfFolder
-
-# ---------- VLM (MedGemma replacement under the same public function name) ----------
-SMARTHEAL_VLM_ID = os.getenv("SMARTHEAL_VLM_ID", "Qwen/Qwen2-VL-2B-Instruct")
-SMARTHEAL_VLM_MAX_NEW_TOKENS = int(os.getenv("SMARTHEAL_VLM_MAX_NEW_TOKENS", "600"))
-SMARTHEAL_VLM_TEMPERATURE = float(os.getenv("SMARTHEAL_VLM_TEMPERATURE", "0.2"))
-
-SMARTHEAL_SYSTEM_PROMPT = """You are SmartHeal, a medical decision-support assistant specialized in wound assessment.
-You are given: (1) a wound photograph, (2) basic patient context, and (3) visual measurements (length, width, area)
-estimated from computer vision. You must:
-
-- Summarize clinically-relevant visual cues (tissue type, exudate amount, slough/necrosis, peri-wound condition).
-- Interpret in context of diabetes/infection/moisture/bleeding risks.
-- Provide clear next-step care: cleansing, debridement criteria, dressing selection, offloading, escalation triggers.
-- Include risk flags (ischemia, cellulitis, osteomyelitis suspicion) and monitoring frequency.
-- Be concise, structured, and avoid speculation beyond the image and given data.
-- Always add a short disclaimer: “Decision-support only; verify clinically.”"""
-
-def _build_vlm_messages(patient_info: str, visual_results: Dict, guideline_context: str) -> list:
-    wt = visual_results.get("wound_type", "Unknown")
-    L = visual_results.get("length_cm", 0)
-    W = visual_results.get("breadth_cm", 0)
-    A = visual_results.get("surface_area_cm2", 0)
-    ppcm = visual_results.get("px_per_cm", "?")
-
-    ctx = (guideline_context or "")
-    if ctx:
-        ctx = f"\n\nRelevant guideline snippets:\n{ctx[:1200]}{'...' if len(ctx)>1200 else ''}"
-
-    text = (
-        f"{SMARTHEAL_SYSTEM_PROMPT}\n\n"
-        f"Patient: {patient_info}\n"
-        f"Wound visual summary (from CV): type={wt}, length={L} cm, width={W} cm, "
-        f"area={A} cm² (calibration {ppcm} px/cm)."
-        f"{ctx}\n\n"
-        "Analyze the image and provide:\n"
-        "1) Clinical Summary\n2) Dressing & Treatment Plan\n"
-        "3) Risk/Red Flags\n4) Monitoring Plan\n"
-        "Format with short headings and bullets.\n"
-    )
-    return [{"role": "user", "content": [{"type": "text", "text": text}]}]
-
-@spaces.GPU  # non-optional: ensure CUDA work happens only inside the ZeroGPU worker
-def _vlm_infer_gpu(
-    image_pil: Image.Image,
-    messages: list,
-    max_new_tokens: int,
-    temperature: float,
-    model_id: str,
-    token: Optional[str],
-) -> str:
-    import torch
-    from transformers import AutoProcessor, AutoModelForCausalLM
-
-    torch.backends.cuda.matmul.allow_tf32 = True
-    device = "cuda"
-
-    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, use_fast=True, token=token)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-        token=token,
-    ).to(device)
-
-    inputs = processor(messages=messages, images=[image_pil], return_tensors="pt").to(device)
-    gen_ids = model.generate(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=False,
-        temperature=temperature,
-    )
-    out = processor.batch_decode(gen_ids, skip_special_tokens=True)[0]
-    return out.strip()
-
-def _vlm_infer_cpu(
-    image_pil: Image.Image,
-    messages: list,
-    max_new_tokens: int,
-    temperature: float,
-    model_id: str,
-    token: Optional[str],
-) -> str:
-    from transformers import pipeline
-    pipe = pipeline(
-        task="image-text-to-text",
-        model=model_id,
-        device="cpu",
-        trust_remote_code=True,
-        token=token,
-    )
-    out = pipe(
-        text=[{"role": "user", "content": [{"type": "image", "image": image_pil}, *messages[0]["content"]]}],
-        max_new_tokens=max_new_tokens,
-        do_sample=False,
-        temperature=temperature,
-    )
-    try:
-        return (out[0]["generated_text"][-1].get("content", "") or "").strip()
-    except Exception:
-        return (out[0].get("generated_text", "") or "").strip()
-
-def generate_medgemma_report(  # <-- keep the original PUBLIC NAME
-    patient_info: str,
-    visual_results: Dict,
-    guideline_context: str,
-    image_pil: Image.Image,
-    max_new_tokens: Optional[int] = None,
-) -> str:
-    """
-    Re-implemented to use Qwen/Qwen2-VL-* via ZeroGPU (@spaces.GPU) with CPU fallback.
-    Name preserved for compatibility with existing callers.
-    """
-    msgs = _build_vlm_messages(patient_info, visual_results, guideline_context)
-    try:
-        return _vlm_infer_gpu(
-            image_pil=image_pil,
-            messages=msgs,
-            max_new_tokens=max_new_tokens or SMARTHEAL_VLM_MAX_NEW_TOKENS,
-            temperature=SMARTHEAL_VLM_TEMPERATURE,
-            model_id=SMARTHEAL_VLM_ID,
-            token=HF_TOKEN,
-        )
-    except Exception as e:
-        logging.warning(f"GPU VLM failed; falling back to CPU: {e!r}")
-        return _vlm_infer_cpu(
-            image_pil=image_pil,
-            messages=msgs,
-            max_new_tokens=max_new_tokens or SMARTHEAL_VLM_MAX_NEW_TOKENS,
-            temperature=SMARTHEAL_VLM_TEMPERATURE,
-            model_id=SMARTHEAL_VLM_ID,
-            token=HF_TOKEN,
-        ) or "⚠️ VLM returned empty output"
-
-# ---------- Initialize CPU models ----------
-def load_yolo_model():
-    YOLO = _import_ultralytics()
-    return YOLO(YOLO_MODEL_PATH)
-
-def load_segmentation_model():
-    load_model = _import_tf_loader()
-    return load_model(SEG_MODEL_PATH, compile=False)
-
-def load_classification_pipeline():
-    pipe = _import_hf_cls()
-    return pipe("image-classification", model="Hemg/Wound-classification", token=HF_TOKEN, device="cpu")
-
-def load_embedding_model():
-    Emb = _import_embeddings()
-    return Emb(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": "cpu"})
-
-def initialize_cpu_models() -> None:
-    if HF_TOKEN:
-        try:
-            HfApi, HfFolder = _import_hf_hub()
-            HfFolder.save_token(HF_TOKEN)
-            logging.info("✅ HF token set")
-        except Exception as e:
-            logging.warning(f"HF token save failed: {e}")
+from datetime import datetime
 
-    if "det" not in models_cache:
-        try:
-            models_cache["det"] = load_yolo_model()
-            logging.info("✅ YOLO loaded (CPU; CUDA masked in main)")
-        except Exception as e:
-            logging.error(f"YOLO load failed: {e}")
+# ---- Safe env defaults (do NOT init CUDA in main) ----
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # mask GPU in main
+HF_HUB_DISABLE_TELEMETRY = os.environ.get("HF_HUB_DISABLE_TELEMETRY", "1")
 
-    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)
-                oshape = getattr(m, "output_shape", None)
-                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.")
-        except Exception as e:
-            models_cache["seg"] = None
-            logging.warning(f"Segmentation unavailable: {e}")
+import json
 
-    if "cls" not in models_cache:
-        try:
-            models_cache["cls"] = load_classification_pipeline()
-            logging.info("✅ Classifier loaded (CPU)")
-        except Exception as e:
-            models_cache["cls"] = None
-            logging.warning(f"Classifier unavailable: {e}")
+# Light imports that won't trigger CUDA
+from transformers import pipeline
+from ultralytics import YOLO
 
-    if "embedding_model" not in models_cache:
-        try:
-            models_cache["embedding_model"] = load_embedding_model()
-            logging.info("✅ Embeddings loaded (CPU)")
-        except Exception as e:
-            models_cache["embedding_model"] = None
-            logging.warning(f"Embeddings unavailable: {e}")
+# TensorFlow: keep on CPU in main process
+import tensorflow as tf
+try:
+    tf.config.set_visible_devices([], "GPU")
+except Exception:
+    pass
 
-def setup_knowledge_base() -> None:
-    if "vector_store" in knowledge_base_cache:
-        return
-    docs: List = []
-    try:
-        PyPDFLoader = _import_langchain_pdf()
-        for pdf in GUIDELINE_PDFS:
-            if os.path.exists(pdf):
-                try:
-                    docs.extend(PyPDFLoader(pdf).load())
-                    logging.info(f"Loaded PDF: {pdf}")
-                except Exception as e:
-                    logging.warning(f"PDF load failed ({pdf}): {e}")
-    except Exception as e:
-        logging.warning(f"LangChain PDF loader unavailable: {e}")
+# LangChain bits (match your old code; no function name change)
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
 
-    if docs and models_cache.get("embedding_model"):
-        try:
-            from langchain.text_splitter import RecursiveCharacterTextSplitter
-            FAISS = _import_langchain_faiss()
-            chunks = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100).split_documents(docs)
-            knowledge_base_cache["vector_store"] = FAISS.from_documents(chunks, models_cache["embedding_model"])
-            logging.info(f"✅ Knowledge base ready ({len(chunks)} chunks)")
-        except Exception as e:
-            knowledge_base_cache["vector_store"] = None
-            logging.warning(f"KB build failed: {e}")
-    else:
-        knowledge_base_cache["vector_store"] = None
-        logging.warning("KB disabled (no docs or embeddings).")
-
-initialize_cpu_models()
-setup_knowledge_base()
-
-# ---------- Calibration helpers ----------
-def _exif_to_dict(pil_img: Image.Image) -> Dict[str, object]:
-    out = {}
-    try:
-        exif = pil_img.getexif()
-        if not exif:
-            return out
-        for k, v in exif.items():
-            tag = TAGS.get(k, k)
-            out[tag] = v
-    except Exception:
-        pass
-    return out
+from huggingface_hub import HfApi, HfFolder
 
-def _to_float(val) -> Optional[float]:
-    try:
-        if val is None:
-            return None
-        if isinstance(val, tuple) and len(val) == 2:
-            num, den = float(val[0]), float(val[1]) if float(val[1]) != 0 else 1.0
-            return num / den
-        return float(val)
-    except Exception:
-        return None
+# Spaces (ZeroGPU)
+try:
+    import spaces
+    SPACES_AVAILABLE = True
+except Exception:
+    spaces = None
+    SPACES_AVAILABLE = False
 
-def _estimate_sensor_width_mm(f_mm: Optional[float], f35: Optional[float]) -> Optional[float]:
-    if f_mm and f35 and f35 > 0:
-        return 36.0 * f_mm / f35
-    return None
+from .config import Config
 
-def estimate_px_per_cm_from_exif(pil_img: Image.Image, default_px_per_cm: float = DEFAULT_PX_PER_CM) -> Tuple[float, Dict]:
-    meta = {"used": "default", "f_mm": None, "f35": None, "sensor_w_mm": None, "distance_m": None}
-    try:
-        exif = _exif_to_dict(pil_img)
-        f_mm = _to_float(exif.get("FocalLength"))
-        f35 = _to_float(exif.get("FocalLengthIn35mmFilm") or exif.get("FocalLengthIn35mm"))
-        subj_dist_m = _to_float(exif.get("SubjectDistance"))
-        sensor_w_mm = _estimate_sensor_width_mm(f_mm, f35)
-        meta.update({"f_mm": f_mm, "f35": f35, "sensor_w_mm": sensor_w_mm, "distance_m": subj_dist_m})
-
-        if f_mm and sensor_w_mm and subj_dist_m and subj_dist_m > 0:
-            w_px = pil_img.width
-            field_w_mm = sensor_w_mm * (subj_dist_m * 1000.0) / f_mm
-            field_w_cm = field_w_mm / 10.0
-            px_per_cm = w_px / max(field_w_cm, 1e-6)
-            px_per_cm = float(np.clip(px_per_cm, PX_PER_CM_MIN, PX_PER_CM_MAX))
-            meta["used"] = "exif"
-            return px_per_cm, meta
-        return float(default_px_per_cm), meta
-    except Exception:
-        return float(default_px_per_cm), meta
-
-# ---------- Segmentation helpers ----------
-def _imagenet_norm(arr: np.ndarray) -> np.ndarray:
-    mean = np.array([123.675, 116.28, 103.53], dtype=np.float32)
-    std  = np.array([58.395, 57.12, 57.375], dtype=np.float32)
-    return (arr.astype(np.float32) - mean) / std
-
-def _preprocess_for_seg(bgr_roi: np.ndarray, target_hw: Tuple[int, int]) -> np.ndarray:
-    H, W = target_hw
-    resized = cv2.resize(bgr_roi, (W, H), interpolation=cv2.INTER_LINEAR)
-    if SEG_EXPECTS_RGB:
-        resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
-    if SEG_NORM.lower() == "imagenet":
-        x = _imagenet_norm(resized)
-    else:
-        x = resized.astype(np.float32) / 255.0
-    x = np.expand_dims(x, axis=0)  # (1,H,W,3)
-    return x
-
-def _to_prob(pred: np.ndarray) -> np.ndarray:
-    p = np.squeeze(pred)
-    pmin, pmax = float(p.min()), float(p.max())
-    if pmax > 1.0 or pmin < 0.0:
-        p = 1.0 / (1.0 + np.exp(-p))
-    return p.astype(np.float32)
-
-# ---- Adaptive threshold + GrabCut grow ----
-def _adaptive_prob_threshold(p: np.ndarray) -> float:
-    """
-    Choose a threshold that avoids tiny blobs while not swallowing skin.
-    Try Otsu and the 90th percentile, clamp to [0.25, 0.65], pick by area heuristic.
-    """
-    p01 = np.clip(p.astype(np.float32), 0, 1)
-    p255 = (p01 * 255).astype(np.uint8)
-
-    ret_otsu, _ = cv2.threshold(p255, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-    thr_otsu = float(np.clip(ret_otsu / 255.0, 0.25, 0.65))
-    thr_pctl = float(np.clip(np.percentile(p01, 90), 0.25, 0.65))
-
-    def area_frac(thr: float) -> float:
-        return float((p01 >= thr).sum()) / float(p01.size)
-
-    af_otsu = area_frac(thr_otsu)
-    af_pctl = area_frac(thr_pctl)
-
-    def score(af: float) -> float:
-        target_low, target_high = 0.03, 0.10
-        if af < target_low: return abs(af - target_low) * 3.0
-        if af > target_high: return abs(af - target_high) * 1.5
-        return 0.0
-
-    return thr_otsu if score(af_otsu) <= score(af_pctl) else thr_pctl
 
-def _grabcut_refine(bgr: np.ndarray, seed01: np.ndarray, iters: int = 3) -> np.ndarray:
-    """Grow from a confident core into low-contrast margins."""
-    h, w = bgr.shape[:2]
-    gc = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
-    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-    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
-    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)
-    return np.where((gc == cv2.GC_FGD) | (gc == cv2.GC_PR_FGD), 1, 0).astype(np.uint8)
+# ----------------------------- utils -----------------------------
+def _largest_component(mask01: np.ndarray, min_area_px: int = 50) -> np.ndarray:
+    num, labels, stats, _ = cv2.connectedComponentsWithStats(mask01.astype(np.uint8), 8)
+    if num <= 1:
+        return (mask01 > 0).astype(np.uint8)
+    areas = stats[1:, cv2.CC_STAT_AREA]
+    if areas.size == 0 or areas.max() < min_area_px:
+        return (mask01 > 0).astype(np.uint8)
+    idx = 1 + int(np.argmax(areas))
+    return (labels == idx).astype(np.uint8)
 
 def _fill_holes(mask01: np.ndarray) -> np.ndarray:
     h, w = mask01.shape[:2]
@@ -441,558 +64,610 @@ def _fill_holes(mask01: np.ndarray) -> np.ndarray:
     return out.astype(np.uint8)
 
 def _clean_mask(mask01: np.ndarray) -> np.ndarray:
-    """Open → Close → Fill holes → Largest component (no dilation)."""
     mask01 = (mask01 > 0).astype(np.uint8)
     k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
     k5 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
     mask01 = cv2.morphologyEx(mask01, cv2.MORPH_OPEN, k3, iterations=1)
     mask01 = cv2.morphologyEx(mask01, cv2.MORPH_CLOSE, k5, iterations=1)
     mask01 = _fill_holes(mask01)
-    # Keep largest component only
-    num, labels, stats, _ = cv2.connectedComponentsWithStats(mask01, 8)
-    if num > 1:
-        areas = stats[1:, cv2.CC_STAT_AREA]
-        if areas.size:
-            largest_idx = 1 + int(np.argmax(areas))
-            mask01 = (labels == largest_idx).astype(np.uint8)
+    mask01 = _largest_component(mask01)
     return (mask01 > 0).astype(np.uint8)
 
-# ---------- Segmentation pipeline ----------
-def segment_wound(image_bgr: np.ndarray, ts: str, out_dir: str) -> Tuple[np.ndarray, Dict[str, object]]:
-    """
-    TF model → adaptive threshold on prob → GrabCut grow → cleanup.
-    Fallback: KMeans-Lab.
-    Returns (mask_uint8_0_255, debug_dict)
-    """
-    debug = {"used": None, "reason": None, "positive_fraction": 0.0,
-             "thr": None, "heatmap_path": None, "roi_seen_by_model": None}
-
-    seg_model = models_cache.get("seg", None)
-
-    # --- 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])
-
-            x = _preprocess_for_seg(image_bgr, (th, tw))
-            roi_seen_path = None
-            if SMARTHEAL_DEBUG:
-                roi_seen_path = os.path.join(out_dir, f"roi_for_seg_{ts}.png")
-                cv2.imwrite(roi_seen_path, image_bgr)
-
-            pred = seg_model.predict(x, verbose=0)
-            if isinstance(pred, (list, tuple)): pred = pred[0]
-            p = _to_prob(pred)
-            p = cv2.resize(p, (image_bgr.shape[1], image_bgr.shape[0]), interpolation=cv2.INTER_LINEAR)
-
-            heatmap_path = None
-            if SMARTHEAL_DEBUG:
-                hm = (np.clip(p, 0, 1) * 255).astype(np.uint8)
-                heat = cv2.applyColorMap(hm, cv2.COLORMAP_JET)
-                heatmap_path = os.path.join(out_dir, f"seg_pred_heatmap_{ts}.png")
-                cv2.imwrite(heatmap_path, heat)
-
-            thr = _adaptive_prob_threshold(p)
-            core01 = (p >= thr).astype(np.uint8)
-            core_frac = float(core01.sum()) / float(core01.size)
-
-            if core_frac < 0.005:
-                thr2 = max(thr - 0.10, 0.15)
-                core01 = (p >= thr2).astype(np.uint8)
-                thr = thr2
-                core_frac = float(core01.sum()) / float(core01.size)
-
-            if core01.any():
-                gc01 = _grabcut_refine(image_bgr, core01, iters=3)
-                mask01 = _clean_mask(gc01)
-            else:
-                mask01 = np.zeros(core01.shape, np.uint8)
-
-            pos_frac = float(mask01.sum()) / float(mask01.size)
-            logging.info(f"SegModel USED | thr={float(thr):.2f} core_frac={core_frac:.4f} final_frac={pos_frac:.4f}")
-
-            debug.update({
-                "used": "tf_model",
-                "reason": "ok",
-                "positive_fraction": pos_frac,
-                "thr": float(thr),
-                "heatmap_path": heatmap_path,
-                "roi_seen_by_model": roi_seen_path
-            })
-            return (mask01 * 255).astype(np.uint8), debug
-
-        except Exception as e:
-            logging.warning(f"⚠️ Segmentation model failed → fallback. Reason: {e}")
-            debug.update({"used": "fallback_kmeans", "reason": f"model_failed: {e}"})
+def _grabcut_refine(bgr: np.ndarray, seed01: np.ndarray, iters: int = 3) -> np.ndarray:
+    h, w = bgr.shape[:2]
+    gc = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
+    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    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
+    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)
+    return np.where((gc == cv2.GC_FGD) | (gc == cv2.GC_PR_FGD), 1, 0).astype(np.uint8)
 
-    # --- Fallback: KMeans in Lab (reddest cluster as wound) ---
-    Z = image_bgr.reshape((-1, 3)).astype(np.float32)
+def _fallback_segment(roi_bgr: np.ndarray) -> np.ndarray:
+    """Robust OpenCV fallback: Lab 2-cluster (maximize a*), then GrabCut grow + cleanup."""
+    Z = roi_bgr.reshape((-1, 3)).astype(np.float32)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
     _, 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]))  # maximize a* (red)
-    mask01 = (labels.reshape(image_bgr.shape[:2]) == wound_idx).astype(np.uint8)
-    mask01 = _clean_mask(mask01)
-
-    pos_frac = float(mask01.sum()) / float(mask01.size)
-    logging.info(f"KMeans USED | final_frac={pos_frac:.4f}")
-
-    debug.update({
-        "used": "fallback_kmeans",
-        "reason": debug.get("reason") or "no_model",
-        "positive_fraction": pos_frac,
-        "thr": None
-    })
-    return (mask01 * 255).astype(np.uint8), debug
-
-# ---------- Measurement + overlay helpers ----------
-def largest_component_mask(binary01: np.ndarray, min_area_px: int = 50) -> np.ndarray:
-    num, labels, stats, _ = cv2.connectedComponentsWithStats(binary01.astype(np.uint8), connectivity=8)
-    if num <= 1:
-        return binary01.astype(np.uint8)
-    areas = stats[1:, cv2.CC_STAT_AREA]
-    if areas.size == 0 or areas.max() < min_area_px:
-        return binary01.astype(np.uint8)
-    largest_idx = 1 + int(np.argmax(areas))
-    return (labels == largest_idx).astype(np.uint8)
-
-def measure_min_area_rect(mask01: np.ndarray, px_per_cm: float) -> Tuple[float, float, Tuple]:
-    contours, _ = cv2.findContours(mask01.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if not contours:
-        return 0.0, 0.0, (None, None)
-    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_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)
-    return length_cm, breadth_cm, (box, rect[0])
-
-def area_cm2_from_contour(mask01: np.ndarray, px_per_cm: float) -> Tuple[float, Optional[np.ndarray]]:
-    """Area from largest polygon (sub-pixel); returns (area_cm2, contour)."""
-    m = (mask01 > 0).astype(np.uint8)
-    contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if not contours:
-        return 0.0, None
-    cnt = max(contours, key=cv2.contourArea)
-    poly_area_px2 = float(cv2.contourArea(cnt))
-    area_cm2 = round(poly_area_px2 / (max(px_per_cm, 1e-6) ** 2), 2)
-    return area_cm2, cnt
-
-def clamp_area_with_minrect(cnt: np.ndarray, px_per_cm: float, area_cm2_poly: float) -> float:
-    rect = cv2.minAreaRect(cnt)
-    (w_px, h_px) = rect[1]
-    rect_area_px2 = float(max(w_px, 0.0) * max(h_px, 0.0))
-    rect_area_cm2 = rect_area_px2 / (max(px_per_cm, 1e-6) ** 2)
-    return round(min(area_cm2_poly, rect_area_cm2 * 1.05), 2)
-
-def draw_measurement_overlay(
-    base_bgr: np.ndarray,
-    mask01: np.ndarray,
-    rect_box: np.ndarray,
-    length_cm: float,
-    breadth_cm: float,
-    thickness: int = 2
-) -> np.ndarray:
-    """
-    1) Strong red mask overlay + white contour
-    2) Min-area rectangle
-    3) Double-headed arrows labeled Length/Width
-    """
-    overlay = base_bgr.copy()
-
-    # Mask tint
-    mask255 = (mask01 * 255).astype(np.uint8)
-    mask3 = cv2.merge([mask255, mask255, mask255])
-    red = np.zeros_like(overlay); red[:] = (0, 0, 255)
-    alpha = 0.55
-    tinted = cv2.addWeighted(overlay, 1 - alpha, red, alpha, 0)
-    overlay = np.where(mask3 > 0, tinted, overlay)
-
-    # Contour
-    cnts, _ = cv2.findContours(mask255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if cnts:
-        cv2.drawContours(overlay, cnts, -1, (255, 255, 255), 2)
-
-    if rect_box is not None:
-        cv2.polylines(overlay, [rect_box], True, (255, 255, 255), thickness)
-        pts = rect_box.reshape(-1, 2)
-
-        def midpoint(a, b): return (int((a[0] + b[0]) / 2), int((a[1] + b[1]) / 2))
-        e = [np.linalg.norm(pts[i] - pts[(i + 1) % 4]) for i in range(4)]
-        long_edge_idx = int(np.argmax(e))
-        mids = [midpoint(pts[i], pts[(i + 1) % 4]) for i in range(4)]
-        long_pair = (long_edge_idx, (long_edge_idx + 2) % 4)
-        short_pair = ((long_edge_idx + 1) % 4, (long_edge_idx + 3) % 4)
-
-        def draw_double_arrow(img, p1, p2):
-            cv2.arrowedLine(img, p1, p2, (0, 0, 0), thickness + 2, tipLength=0.05)
-            cv2.arrowedLine(img, p2, p1, (0, 0, 0), thickness + 2, tipLength=0.05)
-            cv2.arrowedLine(img, p1, p2, (255, 255, 255), thickness, tipLength=0.05)
-            cv2.arrowedLine(img, p2, p1, (255, 255, 255), thickness, tipLength=0.05)
-
-        def put_label(text, anchor):
-            org = (anchor[0] + 6, anchor[1] - 6)
-            cv2.putText(overlay, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 4, cv2.LINE_AA)
-            cv2.putText(overlay, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
-
-        draw_double_arrow(overlay, mids[long_pair[0]], mids[long_pair[1]])
-        draw_double_arrow(overlay, mids[short_pair[0]], mids[short_pair[1]])
-        put_label(f"Length: {length_cm:.2f} cm", mids[long_pair[0]])
-        put_label(f"Width:  {breadth_cm:.2f} cm", mids[short_pair[0]])
-
-    return overlay
-
-# ---------- AI PROCESSOR ----------
+    wound_idx = int(np.argmax(centers_lab[:, 1]))  # reddest cluster
+    seed01 = (labels.reshape(roi_bgr.shape[:2]) == wound_idx).astype(np.uint8)
+    gc01 = _grabcut_refine(roi_bgr, seed01, iters=3)
+    return _clean_mask(gc01)
+
+def _safe_load_seg_model(path: str):
+    """Try multiple loaders to survive Keras 3 / TF 2.15 / h5 mismatches."""
+    if not os.path.exists(path):
+        return None
+    try:
+        # Keras legacy API (present in TF 2.13+ with legacy shim)
+        from tensorflow import keras as tfk
+        if hasattr(tfk, "saving") and hasattr(tfk.saving, "legacy"):
+            return tfk.saving.legacy.load_model(path, compile=False)
+    except Exception:
+        pass
+    try:
+        # tf.keras standard loader
+        from tensorflow.keras.models import load_model as tf_load_model
+        return tf_load_model(path, compile=False)
+    except Exception as e:
+        logging.warning(f"Segmentation model failed to load with legacy + tf.keras: {e}")
+        return None
+
+
+# ----------------------------- main class -----------------------------
 class AIProcessor:
     def __init__(self):
-        self.models_cache = models_cache
-        self.knowledge_base_cache = knowledge_base_cache
-        self.uploads_dir = UPLOADS_DIR
-        self.dataset_id = DATASET_ID
-        self.hf_token = HF_TOKEN
-
-    def _ensure_analysis_dir(self) -> str:
-        out_dir = os.path.join(self.uploads_dir, "analysis")
-        os.makedirs(out_dir, exist_ok=True)
-        return out_dir
-
-    def perform_visual_analysis(self, image_pil: Image.Image) -> Dict:
-        """
-        YOLO detect → crop ROI → segment_wound(ROI) → clean mask →
-        minAreaRect measurement (cm) using EXIF px/cm → save outputs.
-        """
+        self.models_cache = {}
+        self.knowledge_base_cache = {}
+        self.config = Config()
+        self.px_per_cm = self.config.PIXELS_PER_CM
+        self._initialize_models()
+
+    def _initialize_models(self):
+        """Initialize all AI models except GPU VLM (that one loads inside the GPU worker)."""
         try:
-            px_per_cm, exif_meta = estimate_px_per_cm_from_exif(image_pil, DEFAULT_PX_PER_CM)
-            # Guardrails for calibration to avoid huge area blow-ups
-            px_per_cm = float(np.clip(px_per_cm, 20.0, 350.0))
-            if (exif_meta or {}).get("used") != "exif":
-                logging.warning(f"Calibration fallback used: px_per_cm={px_per_cm:.2f} (default). Prefer ruler/Aruco for accuracy.")
+            # HF token
+            if self.config.HF_TOKEN:
+                HfFolder.save_token(self.config.HF_TOKEN)
+                logging.info("✅ HF token set")
+
+            # YOLO (CPU-only in main)
+            try:
+                self.models_cache["det"] = YOLO(self.config.YOLO_MODEL_PATH)
+                logging.info("✅ YOLO loaded (CPU; CUDA masked in main)")
+            except Exception as e:
+                logging.warning(f"YOLO not available: {e}")
+
+            # Segmentation (safe loader, stays on CPU)
+            try:
+                seg = _safe_load_seg_model(self.config.SEG_MODEL_PATH)
+                if seg is None:
+                    raise RuntimeError("segmentation file missing or incompatible")
+                self.models_cache["seg"] = seg
+                logging.info("✅ Segmentation model loaded (CPU)")
+            except Exception as e:
+                self.models_cache["seg"] = None
+                logging.warning(f"Segmentation unavailable: {e}")
+
+            # Wound classifier (CPU)
+            try:
+                self.models_cache["cls"] = pipeline(
+                    "image-classification",
+                    model="Hemg/Wound-classification",
+                    token=self.config.HF_TOKEN,
+                    device="cpu",
+                )
+                logging.info("✅ Classifier loaded (CPU)")
+            except Exception as e:
+                self.models_cache["cls"] = None
+                logging.warning(f"Classifier unavailable: {e}")
+
+            # Embeddings for KB (CPU)
+            try:
+                self.models_cache["embedding_model"] = HuggingFaceEmbeddings(
+                    model_name="sentence-transformers/all-MiniLM-L6-v2",
+                    model_kwargs={"device": "cpu"},
+                )
+                logging.info("✅ Embeddings loaded (CPU)")
+            except Exception as e:
+                self.models_cache["embedding_model"] = None
+                logging.warning(f"Embeddings unavailable: {e}")
+
+            self._load_knowledge_base()
+        except Exception as e:
+            logging.error(f"Error initializing models: {e}")
 
+    def _load_knowledge_base(self):
+        """Load guideline PDFs into a FAISS vector store."""
+        try:
+            documents = []
+            for pdf_path in self.config.GUIDELINE_PDFS:
+                if os.path.exists(pdf_path):
+                    try:
+                        loader = PyPDFLoader(pdf_path)
+                        docs = loader.load()
+                        documents.extend(docs)
+                        logging.info(f"Loaded PDF: {pdf_path}")
+                    except Exception as e:
+                        logging.warning(f"PDF load failed ({pdf_path}): {e}")
+
+            if documents and self.models_cache.get("embedding_model"):
+                splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+                chunks = splitter.split_documents(documents)
+                vectorstore = FAISS.from_documents(chunks, self.models_cache["embedding_model"])
+                self.knowledge_base_cache["vectorstore"] = vectorstore
+                logging.info(f"✅ Knowledge base ready ({len(chunks)} chunks)")
+            else:
+                self.knowledge_base_cache["vectorstore"] = None
+                logging.warning("Knowledge base not available (no PDFs or embeddings).")
+        except Exception as e:
+            logging.warning(f"Knowledge base loading error: {e}")
+            self.knowledge_base_cache["vectorstore"] = None
+
+    # ------------------------ vision core ------------------------
+    def perform_visual_analysis(self, image_pil):
+        """Perform comprehensive visual analysis of wound image."""
+        try:
             image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
 
-            # --- Detection ---
-            det_model = self.models_cache.get("det")
-            if det_model is None:
-                raise RuntimeError("YOLO model not loaded")
-            results = det_model.predict(image_cv, verbose=False, device="cpu")
+            # YOLO detection
+            if "det" not in self.models_cache or self.models_cache["det"] is None:
+                raise ValueError("YOLO detection model not available.")
+            results = self.models_cache["det"].predict(image_cv, verbose=False, device="cpu")
             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.")
+                raise ValueError("No wound detected in the image.")
 
-            box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
+            box = results[0].boxes.xyxy[0].cpu().numpy().astype(int)
             x1, y1, x2, y2 = [int(v) for v in box]
             x1, y1 = max(0, x1), max(0, y1)
             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:
-                try:
-                    import gradio as gr
-                    raise gr.Error("Detected ROI is empty.")
-                except Exception:
-                    raise RuntimeError("Detected ROI is empty.")
+            region_cv = image_cv[y1:y2, x1:x2].copy()
 
-            out_dir = self._ensure_analysis_dir()
+            # Save detection vis
+            os.makedirs(os.path.join(self.config.UPLOADS_DIR, "analysis"), exist_ok=True)
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-
-            # --- Segmentation (model-first + KMeans fallback) ---
-            mask_u8_255, seg_debug = segment_wound(roi, ts, out_dir)
-            mask01 = (mask_u8_255 > 127).astype(np.uint8)
-
-            if mask01.any():
-                mask01 = _clean_mask(mask01)
-                logging.debug(f"Mask postproc: px_after={int(mask01.sum())}")
-
-            # --- Measurement (accurate & conservative) ---
-            if mask01.any():
-                length_cm, breadth_cm, (box_pts, _) = measure_min_area_rect(mask01, px_per_cm)
-                area_poly_cm2, largest_cnt = area_cm2_from_contour(mask01, px_per_cm)
-                if largest_cnt is not None:
-                    surface_area_cm2 = clamp_area_with_minrect(largest_cnt, px_per_cm, area_poly_cm2)
-                else:
-                    surface_area_cm2 = area_poly_cm2
-
-                anno_roi = draw_measurement_overlay(roi, mask01, box_pts, length_cm, breadth_cm)
-                segmentation_empty = False
-            else:
-                # Fallback if seg failed: use ROI dimensions
-                h_px = max(0, y2 - y1); w_px = max(0, x2 - x1)
-                length_cm = round(max(h_px, w_px) / px_per_cm, 2)
-                breadth_cm = round(min(h_px, w_px) / px_per_cm, 2)
-                surface_area_cm2 = round((h_px * w_px) / (px_per_cm ** 2), 2)
-                anno_roi = roi.copy()
-                cv2.rectangle(anno_roi, (2, 2), (anno_roi.shape[1]-3, anno_roi.shape[0]-3), (0, 0, 255), 3)
-                cv2.line(anno_roi, (0, 0), (anno_roi.shape[1]-1, anno_roi.shape[0]-1), (0, 0, 255), 2)
-                cv2.line(anno_roi, (anno_roi.shape[1]-1, 0), (0, anno_roi.shape[0]-1), (0, 0, 255), 2)
-                box_pts = None
-                segmentation_empty = True
-
-            # --- Save visualizations ---
-            original_path = os.path.join(out_dir, f"original_{ts}.png")
-            cv2.imwrite(original_path, image_cv)
-
             det_vis = image_cv.copy()
             cv2.rectangle(det_vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            detection_path = os.path.join(out_dir, f"detection_{ts}.png")
-            cv2.imwrite(detection_path, det_vis)
-
-            roi_mask_path = os.path.join(out_dir, f"roi_mask_{ts}.png")
-            cv2.imwrite(roi_mask_path, (mask01 * 255).astype(np.uint8))
-
-            # ROI overlay (mask tint + contour, without arrows)
-            mask255 = (mask01 * 255).astype(np.uint8)
-            mask3   = cv2.merge([mask255, mask255, mask255])
-            red     = np.zeros_like(roi); red[:] = (0, 0, 255)
-            alpha   = 0.55
-            tinted  = cv2.addWeighted(roi, 1 - alpha, red, alpha, 0)
-            if mask255.any():
-                roi_overlay = np.where(mask3 > 0, tinted, roi)
-                cnts, _ = cv2.findContours(mask255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-                cv2.drawContours(roi_overlay, cnts, -1, (255, 255, 255), 2)
+            detection_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"detection_{ts}.png")
+            cv2.imwrite(detection_image_path, det_vis)
+            detection_image_pil = Image.fromarray(cv2.cvtColor(det_vis, cv2.COLOR_BGR2RGB))
+
+            # --- segmentation ---
+            length = breadth = area = 0.0
+            segmentation_image_pil = None
+            segmentation_image_path = None
+
+            mask01 = None
+            seg_model = self.models_cache.get("seg", None)
+            if seg_model is not None:
+                try:
+                    ishape = getattr(seg_model, "input_shape", None)
+                    th, tw = int(ishape[1]), int(ishape[2]) if ishape and len(ishape) >= 3 else (224, 224)
+                    resized = cv2.resize(region_cv, (tw, th), interpolation=cv2.INTER_LINEAR)
+                    x = np.expand_dims(resized.astype(np.float32) / 255.0, 0)
+                    pred = seg_model.predict(x, verbose=0)
+                    if isinstance(pred, (list, tuple)):
+                        pred = pred[0]
+                    p = np.squeeze(pred)
+                    # sigmoid if raw
+                    if p.max() > 1.0 or p.min() < 0.0:
+                        p = 1.0 / (1.0 + np.exp(-p))
+                    p = cv2.resize(p.astype(np.float32), (region_cv.shape[1], region_cv.shape[0]), interpolation=cv2.INTER_LINEAR)
+                    # adaptive threshold
+                    p255 = (np.clip(p, 0, 1) * 255).astype(np.uint8)
+                    thr_val, _ = cv2.threshold(p255, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+                    thr = float(np.clip(thr_val / 255.0, 0.25, 0.65))
+                    seed01 = (p >= thr).astype(np.uint8)
+                    if seed01.sum() == 0:
+                        seed01 = (p >= max(thr - 0.1, 0.15)).astype(np.uint8)
+                    gc01 = _grabcut_refine(region_cv, seed01, iters=3)
+                    mask01 = _clean_mask(gc01)
+                except Exception as e:
+                    logging.warning(f"Segmentation model failed; using OpenCV fallback: {e}")
+                    mask01 = _fallback_segment(region_cv)
             else:
-                roi_overlay = anno_roi
-
+                mask01 = _fallback_segment(region_cv)
+
+            # overlay + measurements
+            overlay = region_cv.copy()
+            red = overlay.copy(); red[:] = (0, 0, 255)
+            if mask01 is not None and mask01.any():
+                mask255 = (mask01 * 255).astype(np.uint8)
+                mask3 = cv2.merge([mask255, mask255, mask255])
+                tinted = cv2.addWeighted(region_cv, 0.45, red, 0.55, 0)
+                overlay = np.where(mask3 > 0, tinted, region_cv)
+                cnts, _ = cv2.findContours(mask255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                if cnts:
+                    cnt = max(cnts, key=cv2.contourArea)
+                    x, y, w, h = cv2.boundingRect(cnt)
+                    length = round(h / self.px_per_cm, 2)
+                    breadth = round(w / self.px_per_cm, 2)
+                    area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
+                    cv2.drawContours(overlay, [cnt], -1, (255, 255, 255), 2)
+
+            segmentation_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"segmentation_{ts}.png")
             seg_full = image_cv.copy()
-            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)
+            seg_full[y1:y2, x1:x2] = overlay
+            cv2.imwrite(segmentation_image_path, seg_full)
+            segmentation_image_pil = Image.fromarray(cv2.cvtColor(seg_full, cv2.COLOR_BGR2RGB))
 
-            segmentation_roi_path = os.path.join(out_dir, f"segmentation_roi_{ts}.png")
-            cv2.imwrite(segmentation_roi_path, roi_overlay)
-
-            # Annotated (mask + arrows + labels) in full-frame
-            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 ---
+            # classification
             wound_type = "Unknown"
-            cls_pipe = self.models_cache.get("cls")
-            if cls_pipe is not None:
+            if self.models_cache.get("cls") is not None:
                 try:
-                    preds = cls_pipe(Image.fromarray(cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)))
+                    region_pil = Image.fromarray(cv2.cvtColor(region_cv, cv2.COLOR_BGR2RGB))
+                    preds = self.models_cache["cls"](region_pil)
                     if preds:
                         wound_type = max(preds, key=lambda x: x.get("score", 0)).get("label", "Unknown")
                 except Exception as e:
-                    logging.warning(f"Classification failed: {e}")
-
-            # Log end-of-seg summary
-            seg_summary = {
-                "seg_used": seg_debug.get("used"),
-                "seg_reason": seg_debug.get("reason"),
-                "positive_fraction": round(float(seg_debug.get("positive_fraction", 0.0)), 6),
-                "threshold": seg_debug.get("thr"),
-                "segmentation_empty": segmentation_empty,
-                "exif_px_per_cm": round(px_per_cm, 3),
-            }
-            _log_kv("SEG_SUMMARY", seg_summary)
+                    logging.warning(f"Wound classification error: {e}")
+
+            conf = 0.0
+            try:
+                conf = float(results[0].boxes.conf[0].cpu().item())
+            except Exception:
+                pass
 
             return {
                 "wound_type": wound_type,
-                "length_cm": length_cm,
-                "breadth_cm": breadth_cm,
-                "surface_area_cm2": surface_area_cm2,
-                "px_per_cm": round(px_per_cm, 2),
-                "calibration_meta": exif_meta,
-                "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": annotated_seg_path,
-                "segmentation_annotated_path": annotated_seg_path,
-                "segmentation_roi_path": segmentation_roi_path,
-                "roi_mask_path": roi_mask_path,
-                "segmentation_empty": segmentation_empty,
-                "segmentation_debug": seg_debug,
-                "original_image_path": original_path,
+                "length_cm": float(length),
+                "breadth_cm": float(breadth),
+                "surface_area_cm2": float(area),
+                "detection_confidence": conf,
+                "bounding_box": [int(x1), int(y1), int(x2), int(y2)],
+                "detection_image_path": detection_image_path,
+                "detection_image_pil": detection_image_pil,
+                "segmentation_image_path": segmentation_image_path,
+                "segmentation_image_pil": segmentation_image_pil,
             }
+
         except Exception as e:
-            logging.error(f"Visual analysis failed: {e}", exc_info=True)
-            raise
+            logging.error(f"Visual analysis error: {e}", exc_info=True)
+            raise ValueError(f"Visual analysis failed: {str(e)}")
 
-    # ---------- Knowledge base + reporting ----------
-    def query_guidelines(self, query: str) -> str:
+    # ------------------------ KB / RAG ------------------------
+    def query_guidelines(self, query: str):
+        """Query the knowledge base for relevant guidelines"""
         try:
-            vs = self.knowledge_base_cache.get("vector_store")
-            if not vs:
-                return "Knowledge base is not available."
-            retriever = vs.as_retriever(search_kwargs={"k": 5})
-            # LangChain deprecation fix: use invoke()
-            docs = retriever.invoke(query)
-            lines: List[str] = []
+            vector_store = self.knowledge_base_cache.get("vectorstore")
+            if not vector_store:
+                return "Knowledge base unavailable - clinical guidelines not loaded"
+
+            retriever = vector_store.as_retriever(search_kwargs={"k": 10})
+            try:
+                docs = retriever.invoke(query)
+            except Exception:
+                # old API fallback
+                docs = retriever.get_relevant_documents(query)
+
+            if not docs:
+                return "No relevant guidelines found for the query"
+
+            out = []
             for d in docs:
-                src = (d.metadata or {}).get("source", "N/A")
-                txt = (d.page_content or "")[:300]
-                lines.append(f"Source: {src}\nContent: {txt}...")
-            return "\n\n".join(lines) if lines else "No relevant guideline snippets found."
+                meta = d.metadata or {}
+                src = meta.get("source", "Unknown")
+                page = meta.get("page", "N/A")
+                content = (d.page_content or "").strip()
+                out.append(f"Source: {src}, Page: {page}\nContent: {content}")
+            return "\n\n".join(out)
         except Exception as e:
-            logging.warning(f"Guidelines query failed: {e}")
-            return f"Guidelines query failed: {str(e)}"
+            logging.error(f"Guidelines query error: {e}")
+            return f"Error querying guidelines: {str(e)}"
+
+    # ------------------------ Reporting (VLM + fallback) ------------------------
+    def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
+        """Generate comprehensive medical report using a VLM if available (loaded inside GPU worker), else fallback."""
+        try:
+            # If medgemma/qwen pipeline wasn't cached by GPU worker, fallback right away.
+            if "medgemma_pipe" not in self.models_cache or self.models_cache["medgemma_pipe"] is None:
+                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+
+            max_tokens = max_new_tokens or self.config.MAX_NEW_TOKENS
+            detection_path = visual_results.get("detection_image_path", "")
+            segmentation_path = visual_results.get("segmentation_image_path", "")
 
-    def _generate_fallback_report(self, patient_info: str, visual_results: Dict, guideline_context: str) -> str:
-        return f"""# 🩺 SmartHeal AI - Comprehensive Wound Analysis Report
+            prompt = f"""
+# Wound Care Report
 
-## 📋 Patient Information
+## 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.
-"""
-
-    def generate_final_report(
-        self,
-        patient_info: str,
-        visual_results: Dict,
-        guideline_context: str,
-        image_pil: Image.Image,
-        max_new_tokens: Optional[int] = None,
-    ) -> str:
-        try:
-            # call the preserved function name (now backed by Qwen2-VL)
-            report = generate_medgemma_report(
-                patient_info, visual_results, guideline_context, image_pil, max_new_tokens
+## Visual Analysis Summary
+- Wound Type: {visual_results.get('wound_type', 'Unknown')}
+- Length: {visual_results.get('length_cm', 0)} cm
+- Breadth: {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):.2f}
+
+## Clinical Reference
+{guideline_context}
+
+You are SmartHeal-AI Agent, a world-class wound care AI specialist trained in clinical wound assessment and guideline-based treatment planning.
+Generate a concise, actionable, evidence-based report with: Clinical Summary, Dressing/Medication Recommendations, Key Risk Factors, and Prognosis & Monitoring. Avoid generic advice; tailor to the data above.
+""".strip()
+
+            content_list = [{"type": "text", "text": prompt}]
+            if image_pil:
+                content_list.insert(0, {"type": "image", "image": image_pil})
+            if visual_results.get("detection_image_pil"):
+                content_list.append({"type": "image", "image": visual_results["detection_image_pil"]})
+            if visual_results.get("segmentation_image_pil"):
+                content_list.append({"type": "image", "image": visual_results["segmentation_image_pil"]})
+
+            messages = [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": "You are a medical AI assistant specializing in wound care. Be precise, objective, and recommendation-focused."}],
+                },
+                {
+                    "role": "user",
+                    "content": content_list,
+                },
+            ]
+
+            out = self.models_cache["medgemma_pipe"](
+                text=messages,
+                max_new_tokens=int(max_tokens),
+                do_sample=False,
             )
-            if report and report.strip() and not report.startswith(("⚠️", "❌")):
-                return report
-            logging.warning("VLM unavailable/invalid; using fallback.")
+            generated = ""
+            try:
+                generated = (out[0]["generated_text"][-1].get("content", "") or "").strip()
+            except Exception:
+                generated = (out[0].get("generated_text", "") or "").strip()
+
+            if generated:
+                images_sec = f"\n\n## Analysis Images\n- Detection: {detection_path}\n- Segmentation: {segmentation_path}\n"
+                return images_sec + generated
+
             return self._generate_fallback_report(patient_info, visual_results, guideline_context)
         except Exception as e:
-            logging.error(f"Report generation failed: {e}")
+            logging.error(f"VLM report generation error: {e}")
             return self._generate_fallback_report(patient_info, visual_results, guideline_context)
 
-    def save_and_commit_image(self, image_pil: Image.Image) -> str:
+    def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
+        detection_path = visual_results.get("detection_image_path", "Not available")
+        segmentation_path = visual_results.get("segmentation_image_path", "Not available")
+        return f"""
+# 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):.2f}
+
+## Analysis Images
+- **Detection Image**: {detection_path}
+- **Segmentation Image**: {segmentation_path}
+
+## Assessment
+Automated measurements provided. Verify via clinical exam.
+
+## Recommendations
+- Cleanse wound; choose dressing per moisture/infection risk
+- Consider debridement if indicated
+- Document with serial photos & measurements
+
+## Clinical Guidelines
+{(guideline_context or '')[:500]}...
+
+*Note: Decision support only; not a diagnosis.*
+""".strip()
+
+    # ------------------------ I/O ------------------------
+    def save_and_commit_image(self, image_pil):
+        """Save image locally and optionally upload to HuggingFace dataset"""
         try:
-            os.makedirs(self.uploads_dir, exist_ok=True)
-            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-            filename = f"{ts}.png"
-            path = os.path.join(self.uploads_dir, filename)
-            image_pil.convert("RGB").save(path)
-            logging.info(f"✅ Image saved locally: {path}")
+            os.makedirs(self.config.UPLOADS_DIR, exist_ok=True)
+            filename = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
+            local_path = os.path.join(self.config.UPLOADS_DIR, filename)
+            image_pil.convert("RGB").save(local_path)
+            logging.info(f"✅ Image saved locally: {local_path}")
 
-            if HF_TOKEN and DATASET_ID:
+            if self.config.HF_TOKEN and self.config.DATASET_ID:
                 try:
-                    HfApi, HfFolder = _import_hf_hub()
-                    HfFolder.save_token(HF_TOKEN)
+                    HfFolder.save_token(self.config.HF_TOKEN)
                     api = HfApi()
                     api.upload_file(
-                        path_or_fileobj=path,
+                        path_or_fileobj=local_path,
                         path_in_repo=f"images/{filename}",
-                        repo_id=DATASET_ID,
+                        repo_id=self.config.DATASET_ID,
                         repo_type="dataset",
-                        token=HF_TOKEN,
                         commit_message=f"Upload wound image: {filename}",
                     )
                     logging.info("✅ Image committed to HF dataset")
                 except Exception as e:
                     logging.warning(f"HF upload failed: {e}")
 
-            return path
+            return local_path
         except Exception as e:
-            logging.error(f"Failed to save/commit image: {e}")
-            return ""
+            logging.error(f"Image saving error: {e}")
+            return None
 
-    def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: Dict) -> Dict:
+    # ------------------------ Pipeline (GPU-safe) ------------------------
+    # decoration is evaluated at import; only the selected branch executes
+    @spaces.GPU(enable_queue=True, duration=120) if SPACES_AVAILABLE else (lambda f: f)
+    def full_analysis_pipeline(self, image, questionnaire_data):
+        """Complete analysis pipeline. VLM loads here (inside GPU worker) to avoid main-process CUDA init."""
         try:
-            saved_path = self.save_and_commit_image(image_pil)
-            visual_results = self.perform_visual_analysis(image_pil)
-
+            # Try to build a VLM inside the worker. If ZeroGPU fails, we fallback to CPU.
+            if "medgemma_pipe" not in self.models_cache or self.models_cache["medgemma_pipe"] is None:
+                vlm_loaded = False
+                # Prefer a small VLM (Qwen2-VL 2B) for Spaces; if it fails on GPU, retry on CPU.
+                try:
+                    self.models_cache["medgemma_pipe"] = pipeline(
+                        "image-text-to-text",
+                        model=os.environ.get("SMARTHEAL_VLM", "Qwen/Qwen2-VL-2B-Instruct"),
+                        token=self.config.HF_TOKEN,
+                        device_map="cuda",  # we're inside ZeroGPU worker now
+                        torch_dtype="auto",
+                        max_new_tokens=self.config.MAX_NEW_TOKENS,
+                        trust_remote_code=True,
+                    )
+                    vlm_loaded = True
+                    logging.info("✅ VLM loaded on GPU worker")
+                except Exception as e:
+                    logging.warning(f"GPU VLM failed; falling back to CPU: {e}")
+                    try:
+                        self.models_cache["medgemma_pipe"] = pipeline(
+                            "image-text-to-text",
+                            model=os.environ.get("SMARTHEAL_VLM", "Qwen/Qwen2-VL-2B-Instruct"),
+                            token=self.config.HF_TOKEN,
+                            device_map="cpu",
+                            torch_dtype="auto",
+                            max_new_tokens=self.config.MAX_NEW_TOKENS,
+                            trust_remote_code=True,
+                        )
+                        vlm_loaded = True
+                        logging.info("✅ VLM loaded on CPU")
+                    except Exception as e2:
+                        self.models_cache["medgemma_pipe"] = None
+                        logging.error(f"❌ Could not load any VLM: {e2}")
+
+            # Save image
+            saved_path = self.save_and_commit_image(image)
+
+            # Visual analysis
+            visual_results = self.perform_visual_analysis(image)
+
+            # Patient info string
             pi = questionnaire_data or {}
-            patient_info = (
-                f"Age: {pi.get('age','N/A')}, "
-                f"Diabetic: {pi.get('diabetic','N/A')}, "
-                f"Allergies: {pi.get('allergies','N/A')}, "
-                f"Date of Wound: {pi.get('date_of_injury','N/A')}, "
-                f"Professional Care: {pi.get('professional_care','N/A')}, "
-                f"Oozing/Bleeding: {pi.get('oozing_bleeding','N/A')}, "
-                f"Infection: {pi.get('infection','N/A')}, "
-                f"Moisture: {pi.get('moisture','N/A')}"
-            )
-
-            query = (
-                f"best practices for managing a {visual_results.get('wound_type','Unknown')} "
-                f"with moisture '{pi.get('moisture','unknown')}' and infection '{pi.get('infection','unknown')}' "
-                f"in a diabetic status '{pi.get('diabetic','unknown')}'"
-            )
+            patient_info = ", ".join([f"{k}: {v}" for k, v in pi.items() if str(v).strip() != ""])
+
+            # KB query
+            wound_type = visual_results.get("wound_type", "wound")
+            moisture = pi.get("moisture", "unknown")
+            infection = pi.get("infection", "unknown")
+            diabetic = pi.get("diabetic", "unknown")
+            query = f"best practices for managing a {wound_type} with moisture level '{moisture}' and signs of infection '{infection}' in a patient who is diabetic '{diabetic}'"
             guideline_context = self.query_guidelines(query)
 
-            report = self.generate_final_report(patient_info, visual_results, guideline_context, image_pil)
+            # Report
+            final_report = self.generate_final_report(patient_info, visual_results, guideline_context, image)
 
             return {
                 "success": True,
                 "visual_analysis": visual_results,
-                "report": report,
+                "report": final_report,
                 "saved_image_path": saved_path,
-                "guideline_context": (guideline_context or "")[:500] + (
-                    "..." if guideline_context and len(guideline_context) > 500 else ""
-                ),
+                "timestamp": datetime.now().isoformat(),
             }
+
         except Exception as e:
-            logging.error(f"Pipeline error: {e}")
+            logging.error(f"Full analysis pipeline error: {e}", exc_info=True)
             return {
                 "success": False,
                 "error": str(e),
-                "visual_analysis": {},
-                "report": f"Analysis failed: {str(e)}",
-                "saved_image_path": None,
-                "guideline_context": "",
+                "timestamp": datetime.now().isoformat(),
             }
 
-    def analyze_wound(self, image, questionnaire_data: Dict) -> Dict:
+    # ------------------------ Legacy API ------------------------
+    def analyze_wound(self, image, questionnaire_data):
+        """Legacy method for backward compatibility"""
         try:
             if isinstance(image, str):
-                if not os.path.exists(image):
-                    raise ValueError(f"Image file not found: {image}")
-                image_pil = Image.open(image)
-            elif isinstance(image, Image.Image):
-                image_pil = image
-            elif isinstance(image, np.ndarray):
-                image_pil = Image.fromarray(image)
+                try:
+                    image = Image.open(image)
+                    logging.info("Converted path to PIL Image")
+                except Exception as e:
+                    logging.error(f"Error opening image: {e}")
+            if not isinstance(image, Image.Image):
+                # file-like?
+                if hasattr(image, "read"):
+                    try:
+                        if hasattr(image, "seek"):
+                            image.seek(0)
+                        image = Image.open(image)
+                    except Exception as e:
+                        logging.error(f"Error reading file-like image: {e}")
+                        raise ValueError(f"Invalid image format: {type(image)}")
+
+            result = self.full_analysis_pipeline(image, questionnaire_data)
+
+            if result.get("success"):
+                return {
+                    "timestamp": result["timestamp"],
+                    "summary": f"Analysis completed for {questionnaire_data.get('patient_name', 'patient')}",
+                    "recommendations": result["report"],
+                    "wound_detection": {
+                        "status": "success",
+                        "detections": [result["visual_analysis"]],
+                        "total_wounds": 1,
+                    },
+                    "segmentation_result": {
+                        "status": "success",
+                        "wound_area_percentage": result["visual_analysis"].get("surface_area_cm2", 0),
+                    },
+                    "risk_assessment": self._assess_risk_legacy(questionnaire_data),
+                    "guideline_recommendations": [result["report"][:200] + "..."],
+                }
             else:
-                raise ValueError(f"Unsupported image type: {type(image)}")
+                return {
+                    "timestamp": result["timestamp"],
+                    "summary": f"Analysis failed: {result.get('error','unknown')}",
+                    "recommendations": "Please consult with a healthcare professional.",
+                    "wound_detection": {"status": "error", "message": result.get("error", "")},
+                    "segmentation_result": {"status": "error", "message": result.get("error", "")},
+                    "risk_assessment": {"risk_score": 0, "risk_level": "Unknown", "risk_factors": []},
+                    "guideline_recommendations": ["Analysis unavailable due to error"],
+                }
 
-            return self.full_analysis_pipeline(image_pil, questionnaire_data or {})
         except Exception as e:
-            logging.error(f"Wound analysis error: {e}")
+            logging.error(f"Legacy analyze_wound error: {e}")
             return {
-                "success": False,
-                "error": str(e),
-                "visual_analysis": {},
-                "report": f"Analysis initialization failed: {str(e)}",
-                "saved_image_path": None,
-                "guideline_context": "",
-            }
+                "timestamp": datetime.now().isoformat(),
+                "summary": f"Analysis error: {str(e)}",
+                "recommendations": "Please consult with a healthcare professional.",
+                "wound_detection": {"status": "error", "message": str(e)},
+                "segmentation_result": {"status": "error", "message": str(e)},
+                "risk_assessment": {"risk_score": 0, "risk_level": "Unknown", "risk_factors": []},
+                "guideline_recommendations": ["Analysis unavailable due to error"],
+            }
+
+    def _assess_risk_legacy(self, questionnaire_data):
+        """Legacy risk assessment for backward compatibility"""
+        risk_factors = []
+        risk_score = 0
+        try:
+            age = int(questionnaire_data.get("patient_age", 0) or 0)
+            if age > 65:
+                risk_factors.append("Advanced age (>65)")
+                risk_score += 2
+            elif age > 50:
+                risk_factors.append("Older adult (50-65)")
+                risk_score += 1
+
+            duration = str(questionnaire_data.get("wound_duration", "")).lower()
+            if any(t in duration for t in ["month", "months", "year"]):
+                risk_factors.append("Chronic wound (>4 weeks)")
+                risk_score += 3
+
+            pain_level = int(questionnaire_data.get("pain_level", 0) or 0)
+            if pain_level >= 7:
+                risk_factors.append("High pain level")
+                risk_score += 2
+
+            medical_history = str(questionnaire_data.get("medical_history", "")).lower()
+            if "diabetes" in medical_history:
+                risk_factors.append("Diabetes mellitus")
+                risk_score += 3
+            if "circulation" in medical_history or "vascular" in medical_history:
+                risk_factors.append("Vascular/circulation issues")
+                risk_score += 2
+            if "immune" in medical_history:
+                risk_factors.append("Immune system compromise")
+                risk_score += 2
+
+            risk_level = "High" if risk_score >= 7 else ("Moderate" if risk_score >= 4 else "Low")
+            return {"risk_score": risk_score, "risk_level": risk_level, "risk_factors": risk_factors}
+        except Exception as e:
+            logging.error(f"Risk assessment error: {e}")
+            return {"risk_score": 0, "risk_level": "Unknown", "risk_factors": []}