Update src/ai_processor.py

src/ai_processor.py

@@ -1,99 +1,225 @@
+# smartheal_ai_processor.py
+# Full, functional module with conditional Spaces GPU support and CPU fallbacks.
+
 import os
+import time
 import logging
+from datetime import datetime
+from typing import Optional, Dict, List
+
 import cv2
 import numpy as np
 from PIL import Image
-from datetime import datetime
-import gradio as gr
-import spaces
-import torch
-import time
-
-from huggingface_hub import HfApi, HfFolder
-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
 
 # =============== LOGGING SETUP ===============
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
 # =============== CONFIGURATION ===============
 UPLOADS_DIR = "uploads"
-if not os.path.exists(UPLOADS_DIR):
-    os.makedirs(UPLOADS_DIR)
-    logging.info(f"Created uploads directory: {UPLOADS_DIR}")
+os.makedirs(UPLOADS_DIR, exist_ok=True)
 
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN", None)
 YOLO_MODEL_PATH = "src/best.pt"
-SEG_MODEL_PATH = "src/segmentation_model.h5"
+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"
-MAX_NEW_TOKENS = 1024  # Reduced for stability
-PIXELS_PER_CM = 38
+DATASET_ID = "SmartHeal/wound-image-uploads"  # optional (set HF_TOKEN too)
+PIXELS_PER_CM = 38  # heuristic
 
 # =============== GLOBAL CACHES ===============
-models_cache = {}
-knowledge_base_cache = {}
+models_cache: Dict[str, object] = {}
+knowledge_base_cache: Dict[str, object] = {}
 
-# =============== LAZY LOADING FUNCTIONS (CPU-SAFE) ===============
-def load_yolo_model(yolo_model_path):
-    """Lazy import and load YOLO model to avoid CUDA initialization."""
+# ---------- Optional imports guarded ----------
+def _import_ultralytics():
     from ultralytics import YOLO
-    return YOLO(yolo_model_path)
+    return YOLO
 
-def load_segmentation_model(seg_model_path):
-    """Lazy import and load segmentation model."""
+def _import_tf_loader():
     import tensorflow as tf
-    tf.config.set_visible_devices([], 'GPU')  # Force CPU for TensorFlow
+    tf.config.set_visible_devices([], "GPU")  # force CPU
     from tensorflow.keras.models import load_model
-    return load_model(seg_model_path, compile=False)
+    return load_model
 
-def load_classification_pipeline(hf_token):
-    """Lazy import and load classification pipeline (CPU only)."""
+def _import_hf_cls():
     from transformers import pipeline
-    return pipeline(
+    return pipeline
+
+def _import_embeddings():
+    from langchain_community.embeddings 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
+
+# =============== SPACES GPU CONDITIONAL ===============
+def _spaces_gpu_available() -> bool:
+    try:
+        import torch
+        return bool(torch.cuda.is_available())
+    except Exception:
+        return False
+
+def _spaces_lib_available() -> bool:
+    try:
+        import spaces  # noqa
+        return True
+    except Exception:
+        return False
+
+HAVE_SPACES_GPU = _spaces_gpu_available() and _spaces_lib_available()
+
+if HAVE_SPACES_GPU:
+    import spaces  # define only if available & GPU present
+
+    @spaces.GPU(enable_queue=True, duration=90)
+    def generate_medgemma_report_with_timeout(
+        patient_info: str,
+        visual_results: Dict,
+        guideline_context: str,
+        image_pil: Image.Image,
+        max_new_tokens: Optional[int] = None,
+    ) -> str:
+        """Runs on Spaces GPU only; callers keep one signature on both paths."""
+        import torch
+        from transformers import pipeline
+        try:
+            torch.cuda.empty_cache()
+
+            prompt = f"""
+You are a medical AI assistant. Analyze this wound image and patient data.
+
+Patient: {patient_info}
+Wound: {visual_results.get('wound_type', 'Unknown')} - {visual_results.get('length_cm', 0)}×{visual_results.get('breadth_cm', 0)} cm
+
+Provide a structured report with:
+1. Clinical Summary
+2. Treatment Recommendations
+3. Risk Assessment
+4. Monitoring Plan
+""".strip()
+
+            pipe = pipeline(
+                "image-text-to-text",
+                model="google/medgemma-4b-it",
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                token=HF_TOKEN,
+                model_kwargs={"low_cpu_mem_usage": True, "use_cache": True},
+            )
+
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image_pil},
+                        {"type": "text", "text": prompt},
+                    ],
+                }
+            ]
+
+            t0 = time.time()
+            out = pipe(
+                text=messages,
+                max_new_tokens=max_new_tokens or 800,
+                do_sample=False,
+                temperature=0.7,
+                pad_token_id=pipe.tokenizer.eos_token_id,
+            )
+            logging.info(f"✅ MedGemma completed in {time.time() - t0:.2f}s")
+
+            if out and len(out) > 0:
+                # Defensive extraction
+                try:
+                    return out[0]["generated_text"][-1].get("content", "").strip() or "⚠️ Empty response"
+                except Exception:
+                    return (out[0].get("generated_text", "") or "").strip() or "⚠️ Empty response"
+            return "⚠️ No output generated"
+        except Exception as e:
+            logging.error(f"❌ MedGemma generation error: {e}")
+            return f"❌ Report generation failed: {str(e)}"
+        finally:
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
+else:
+    def generate_medgemma_report_with_timeout(
+        patient_info: str,
+        visual_results: Dict,
+        guideline_context: str,
+        image_pil: Image.Image,
+        max_new_tokens: Optional[int] = None,
+    ) -> str:
+        """CPU-only path: return a warning so caller uses fallback."""
+        return "⚠️ GPU not available"
+
+# =============== MODEL INITIALIZATION (CPU-SAFE) ===============
+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"
+        token=HF_TOKEN,
+        device="cpu",
     )
 
 def load_embedding_model():
-    """Load embedding model for knowledge base."""
-    return HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2",
-        model_kwargs={"device": "cpu"}
-    )
+    Emb = _import_embeddings()
+    return Emb(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": "cpu"})
 
-# =============== MODEL INITIALIZATION ===============
-def initialize_cpu_models():
-    """Initialize all CPU-only models once."""
-    global models_cache
-    
+def initialize_cpu_models() -> None:
+    """Initialize all CPU-only models once with robust fallbacks."""
+    # Hugging Face auth (optional)
     if HF_TOKEN:
-        HfFolder.save_token(HF_TOKEN)
-        logging.info("✅ HuggingFace token set")
-    
+        try:
+            HfApi, HfFolder = _import_hf_hub()
+            HfFolder.save_token(HF_TOKEN)
+            logging.info("✅ HuggingFace token set")
+        except Exception as e:
+            logging.warning(f"HF token save failed: {e}")
+
     if "det" not in models_cache:
         try:
-            models_cache["det"] = load_yolo_model(YOLO_MODEL_PATH)
-            logging.info("✅ YOLO model loaded (CPU only)")
+            models_cache["det"] = load_yolo_model()
+            logging.info("✅ YOLO model loaded (CPU)")
         except Exception as e:
             logging.error(f"YOLO load failed: {e}")
 
     if "seg" not in models_cache:
         try:
-            models_cache["seg"] = load_segmentation_model(SEG_MODEL_PATH)
-            logging.info("✅ Segmentation model loaded (CPU)")
+            if os.path.exists(SEG_MODEL_PATH):
+                models_cache["seg"] = load_segmentation_model()
+                logging.info("✅ Segmentation model loaded (CPU)")
+            else:
+                models_cache["seg"] = None
+                logging.warning("Segmentation model file not found; skipping seg.")
         except Exception as e:
+            models_cache["seg"] = None
             logging.warning(f"Segmentation model not available: {e}")
 
     if "cls" not in models_cache:
         try:
-            models_cache["cls"] = load_classification_pipeline(HF_TOKEN)
+            models_cache["cls"] = load_classification_pipeline()
             logging.info("✅ Classification pipeline loaded (CPU)")
         except Exception as e:
+            models_cache["cls"] = None
             logging.warning(f"Classification pipeline not available: {e}")
 
     if "embedding_model" not in models_cache:
@@ -101,119 +227,48 @@ def initialize_cpu_models():
             models_cache["embedding_model"] = load_embedding_model()
             logging.info("✅ Embedding model loaded (CPU)")
         except Exception as e:
+            models_cache["embedding_model"] = None
             logging.warning(f"Embedding model not available: {e}")
 
-def setup_knowledge_base():
-    """Load PDF documents and create FAISS vector store."""
-    global knowledge_base_cache
+def setup_knowledge_base() -> None:
+    """Load PDFs and create FAISS vector store (optional)."""
     if "vector_store" in knowledge_base_cache:
         return
 
     docs = []
-    for pdf_path in GUIDELINE_PDFS:
-        if os.path.exists(pdf_path):
-            try:
-                loader = PyPDFLoader(pdf_path)
-                docs.extend(loader.load())
-                logging.info(f"Loaded PDF: {pdf_path}")
-            except Exception as e:
-                logging.warning(f"Failed to load PDF {pdf_path}: {e}")
-
-    if docs and "embedding_model" in models_cache:
-        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-        chunks = splitter.split_documents(docs)
-        knowledge_base_cache["vector_store"] = FAISS.from_documents(chunks, models_cache["embedding_model"])
-        logging.info(f"✅ Knowledge base ready with {len(chunks)} chunks")
+    try:
+        PyPDFLoader = _import_langchain_pdf()
+        for pdf in GUIDELINE_PDFS:
+            if os.path.exists(pdf):
+                try:
+                    loader = PyPDFLoader(pdf)
+                    docs.extend(loader.load())
+                    logging.info(f"Loaded PDF: {pdf}")
+                except Exception as e:
+                    logging.warning(f"Failed to load PDF {pdf}: {e}")
+    except Exception as e:
+        logging.warning(f"LangChain PDF loader unavailable: {e}")
+
+    if docs and models_cache.get("embedding_model"):
+        try:
+            from langchain.text_splitter import RecursiveCharacterTextSplitter
+            FAISS = _import_langchain_faiss()
+            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+            chunks = splitter.split_documents(docs)
+            knowledge_base_cache["vector_store"] = FAISS.from_documents(chunks, models_cache["embedding_model"])
+            logging.info(f"✅ Knowledge base ready with {len(chunks)} chunks")
+        except Exception as e:
+            knowledge_base_cache["vector_store"] = None
+            logging.warning(f"Knowledge base unavailable: {e}")
     else:
         knowledge_base_cache["vector_store"] = None
-        logging.warning("Knowledge base unavailable")
+        logging.warning("Knowledge base disabled (no docs or embeddings).")
 
-# Initialize models on app startup
+# Initialize on import
 initialize_cpu_models()
 setup_knowledge_base()
 
-# =============== GPU-DECORATED MEDGEMMA FUNCTION WITH TIMEOUT HANDLING ===============
-@spaces.GPU(enable_queue=True, duration=90)  # Reduced duration for stability
-def generate_medgemma_report_with_timeout(
-    patient_info,
-    visual_results,
-    guideline_context,
-    image_pil,
-    max_new_tokens=None,
-):
-    """GPU-only function for MedGemma report generation with improved timeout handling."""
-    import torch
-    from transformers import pipeline
-    
-    try:
-        # Clear GPU cache first
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        # Use a shorter, more focused prompt to reduce processing time
-        prompt = f"""
-You are a medical AI assistant. Analyze this wound image and patient data to provide a clinical assessment.
-
-Patient: {patient_info}
-Wound: {visual_results.get('wound_type', 'Unknown')} - {visual_results.get('length_cm', 0)}×{visual_results.get('breadth_cm', 0)}cm
-
-Provide a structured report with:
-1. Clinical Summary (wound appearance, size, location)
-2. Treatment Recommendations (dressings, care protocols)
-3. Risk Assessment (healing factors)
-4. Monitoring Plan (follow-up schedule)
-
-Keep response concise but medically comprehensive.
-"""
-
-        # Initialize pipeline with optimized settings
-        pipe = pipeline(
-            "image-text-to-text",
-            model="google/medgemma-4b-it",
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            token=HF_TOKEN,
-            model_kwargs={"low_cpu_mem_usage": True, "use_cache": True}
-        )
-        
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": image_pil},
-                    {"type": "text", "text": prompt},
-                ]
-            }
-        ]
-        
-        # Generate with conservative settings
-        start_time = time.time()
-        output = pipe(
-            text=messages,
-            max_new_tokens=max_new_tokens or 800,  # Reduced for stability
-            do_sample=False,
-            temperature=0.7,
-            pad_token_id=pipe.tokenizer.eos_token_id
-        )
-        
-        processing_time = time.time() - start_time
-        logging.info(f"✅ MedGemma processing completed in {processing_time:.2f} seconds")
-        
-        if output and len(output) > 0:
-            result = output[0]["generated_text"][-1].get("content", "").strip()
-            return result if result else "⚠️ Empty response generated"
-        else:
-            return "⚠️ No output generated"
-            
-    except Exception as e:
-        logging.error(f"❌ MedGemma generation error: {e}")
-        return f"❌ Report generation failed: {str(e)}"
-    finally:
-        # Clear GPU memory
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-# =============== AI PROCESSOR CLASS ===============
+# =============== AI PROCESSOR ===============
 class AIProcessor:
     def __init__(self):
         self.models_cache = models_cache
@@ -223,237 +278,204 @@ class AIProcessor:
         self.dataset_id = DATASET_ID
         self.hf_token = HF_TOKEN
 
-    def perform_visual_analysis(self, image_pil: Image.Image) -> dict:
-        """Performs the full visual analysis pipeline."""
+    # ---------- Image utilities ----------
+    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 → (optional) Keras seg → (optional) HF classifier → save visuals."""
         try:
-            # Convert PIL to OpenCV format
-            image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
-            
-            # YOLO Detection
-            results = self.models_cache["det"].predict(image_cv, verbose=False, device="cpu")
-            if not results or not results[0].boxes: 
+            image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
+
+            det = self.models_cache.get("det")
+            if det is None:
+                raise RuntimeError("YOLO model not loaded")
+
+            # YOLO on CPU
+            results = det.predict(image_cv, verbose=False, device="cpu")
+            if not results or not getattr(results[0], "boxes", None) or len(results[0].boxes) == 0:
                 raise ValueError("No wound could be detected.")
-                
+
             box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
-            detected_region_cv = image_cv[box[1]:box[3], box[0]:box[2]]
-            
-            # Segmentation
-            input_size = self.models_cache["seg"].input_shape[1:3]
-            resized = cv2.resize(detected_region_cv, (input_size[1], input_size[0]))
-            mask_pred = self.models_cache["seg"].predict(np.expand_dims(resized / 255.0, 0), verbose=0)[0]
-            mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
-            
-            # Calculate measurements
-            contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            length, breadth, area = (0, 0, 0)
-            if contours:
-                cnt = max(contours, key=cv2.contourArea)
-                x, y, w, h = cv2.boundingRect(cnt)
-                length, breadth, area = round(h / self.px_per_cm, 2), round(w / self.px_per_cm, 2), round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
-
-            # Classification
-            detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
-            wound_type = max(self.models_cache["cls"](detected_image_pil), key=lambda x: x["score"])["label"]
-
-            # Save visualization images
-            os.makedirs(f"{self.uploads_dir}/analysis", exist_ok=True)
+            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)
+            detected_region_cv = image_cv[y1:y2, x1:x2]
+
+            # Optional segmentation
+            seg_model = self.models_cache.get("seg")
+            length = breadth = area = 0.0
+            seg_path = None
+            if seg_model is not None and detected_region_cv.size > 0:
+                try:
+                    input_size = seg_model.input_shape[1:3]
+                    resized = cv2.resize(detected_region_cv, (input_size[1], input_size[0]))
+                    mask_pred = seg_model.predict(np.expand_dims(resized / 255.0, 0), verbose=0)[0]
+                    mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
+
+                    contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                    if contours:
+                        cnt = max(contours, 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)
+
+                        # overlay visualization
+                        mask_resized = cv2.resize(
+                            mask_np * 255,
+                            (detected_region_cv.shape[1], detected_region_cv.shape[0]),
+                            interpolation=cv2.INTER_NEAREST,
+                        )
+                        overlay = detected_region_cv.copy()
+                        overlay[mask_resized > 127] = [0, 0, 255]
+                        seg_vis = cv2.addWeighted(detected_region_cv, 0.7, overlay, 0.3, 0)
+
+                        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+                        out_dir = self._ensure_analysis_dir()
+                        seg_path = os.path.join(out_dir, f"segmentation_{ts}.png")
+                        cv2.imwrite(seg_path, seg_vis)
+                except Exception as e:
+                    logging.warning(f"Segmentation step skipped: {e}")
+
+            # Optional classification
+            wound_type = "Unknown"
+            cls_pipe = self.models_cache.get("cls")
+            if cls_pipe is not None:
+                try:
+                    detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
+                    preds = cls_pipe(detected_image_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 step failed: {e}")
+
+            # Save detection & original
+            out_dir = self._ensure_analysis_dir()
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-            
-            # Detection visualization
             det_vis = image_cv.copy()
-            cv2.rectangle(det_vis, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
-            det_path = f"{self.uploads_dir}/analysis/detection_{ts}.png"
+            cv2.rectangle(det_vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            det_path = os.path.join(out_dir, f"detection_{ts}.png")
             cv2.imwrite(det_path, det_vis)
-            
-            # Original image
-            original_path = f"{self.uploads_dir}/analysis/original_{ts}.png"
+
+            original_path = os.path.join(out_dir, f"original_{ts}.png")
             cv2.imwrite(original_path, image_cv)
 
-            # Segmentation visualization
-            seg_path = None
-            if contours:
-                mask_resized = cv2.resize(mask_np * 255, (detected_region_cv.shape[1], detected_region_cv.shape[0]), interpolation=cv2.INTER_NEAREST)
-                overlay = detected_region_cv.copy()
-                overlay[mask_resized > 127] = [0, 0, 255]  # Red overlay for wound area
-                seg_vis = cv2.addWeighted(detected_region_cv, 0.7, overlay, 0.3, 0)
-                seg_path = f"{self.uploads_dir}/analysis/segmentation_{ts}.png"
-                cv2.imwrite(seg_path, seg_vis)
-
-            visual_results = {
-                "wound_type": wound_type, 
-                "length_cm": length, 
-                "breadth_cm": breadth, 
+            return {
+                "wound_type": wound_type,
+                "length_cm": length,
+                "breadth_cm": breadth,
                 "surface_area_cm2": area,
-                "detection_confidence": float(results[0].boxes.conf[0].cpu().item()) if results[0].boxes.conf is not None else 0.0,
+                "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": det_path,
                 "segmentation_image_path": seg_path,
-                "original_image_path": original_path
+                "original_image_path": original_path,
             }
-            return visual_results
-
         except Exception as e:
             logging.error(f"Visual analysis failed: {e}")
-            raise e
+            raise
 
     def query_guidelines(self, query: str) -> str:
-        """Query the knowledge base for relevant information."""
+        """Query the knowledge base (optional)."""
         try:
-            vector_store = self.knowledge_base_cache.get("vector_store")
-            if not vector_store:
+            vs = self.knowledge_base_cache.get("vector_store")
+            if not vs:
                 return "Knowledge base is not available."
-            
-            retriever = vector_store.as_retriever(search_kwargs={"k": 5})  # Reduced for efficiency
-            docs = retriever.invoke(query)
-            return "\n\n".join([f"Source: {doc.metadata.get('source', 'N/A')}\nContent: {doc.page_content[:300]}..." for doc in docs])
-            
+            # support both old and new retriever APIs
+            try:
+                retriever = vs.as_retriever(search_kwargs={"k": 5})
+                docs = retriever.get_relevant_documents(query)  # LC >= 0.2
+            except Exception:
+                retriever = vs.as_retriever(search_kwargs={"k": 5})
+                # older invoke API
+                docs = retriever.invoke(query)
+            lines: List[str] = []
+            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."
         except Exception as e:
-            logging.error(f"Guidelines query failed: {e}")
+            logging.warning(f"Guidelines query failed: {e}")
             return f"Guidelines query failed: {str(e)}"
 
-    def generate_final_report(
-        self, patient_info: str, visual_results: dict, guideline_context: str, 
-        image_pil: Image.Image, max_new_tokens: int = None
-    ) -> str:
-        """Generate final report using MedGemma with timeout handling."""
-        try:
-            # Try MedGemma with timeout handling
-            report = generate_medgemma_report_with_timeout(
-                patient_info, visual_results, guideline_context, image_pil, max_new_tokens
-            )
-            
-            # Check if report is valid
-            if report and report.strip() and not report.startswith("❌") and not report.startswith("⚠️"):
-                return report
-            else:
-                logging.warning("MedGemma returned invalid response, using fallback")
-                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
-                
-        except Exception as e:
-            logging.error(f"MedGemma report generation failed: {e}")
-            return self._generate_fallback_report(patient_info, visual_results, guideline_context)
-
-    def _generate_fallback_report(
-        self, patient_info: str, visual_results: dict, guideline_context: str
-    ) -> str:
-        """Generate comprehensive fallback report if MedGemma fails."""
-        
-        report = f"""# 🩺 SmartHeal AI - Comprehensive Wound Analysis Report
+    # ---------- Report builders ----------
+    def _generate_fallback_report(self, patient_info: str, visual_results: Dict, guideline_context: str) -> str:
+        """Plaintext/markdown fallback when MedGemma is unavailable."""
+        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  
+- **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%}
 
-## 📊 Analysis Images Available
-- **Original Image**: {visual_results.get('original_image_path', 'Available')}
-- **Detection Visualization**: {visual_results.get('detection_image_path', 'Available')}
-- **Segmentation Overlay**: {visual_results.get('segmentation_image_path', 'Available')}
-
-## 🎯 Clinical Assessment Summary
-
-### Wound Classification
-Based on automated analysis, this wound has been classified as **{visual_results.get('wound_type', 'Unspecified')}** with the following characteristics:
-- Size: {visual_results.get('length_cm', 0)} × {visual_results.get('breadth_cm', 0)} cm
-- Total area: {visual_results.get('surface_area_cm2', 0)} cm²
-- Detection confidence: {visual_results.get('detection_confidence', 0):.1%}
-
-### Clinical Observations
-The automated visual analysis provides quantitative measurements that should be verified through clinical examination. The wound type classification helps guide initial treatment considerations.
-
-## 💊 Treatment Recommendations
-
-### Wound Care Protocol
-1. **Assessment**: Comprehensive clinical evaluation by qualified healthcare professional
-2. **Cleaning**: Gentle wound cleansing with appropriate solution
-3. **Debridement**: Remove necrotic tissue if present (professional assessment required)
-4. **Dressing Selection**: Choose appropriate dressing based on wound characteristics:
-   - Moisture level assessment
-   - Infection risk evaluation
-   - Patient comfort and mobility
-
-### Monitoring Plan
-- **Initial Phase**: Daily assessment for first week
-- **Ongoing Care**: Reassessment every 2-3 days or as clinically indicated
-- **Documentation**: Regular photo documentation and measurement tracking
-- **Progress Evaluation**: Weekly review of healing progression
-
-## ⚠️ Risk Factors & Considerations
-
-### Patient-Specific Factors
-Review patient history for factors that may impact healing:
-- Age and general health status
-- Diabetes or metabolic conditions
-- Circulation and vascular health
-- Nutritional status
-- Mobility and pressure relief
-
-### Warning Signs
-Monitor for signs requiring immediate attention:
-- Increased pain, redness, or swelling
-- Purulent drainage or odor
-- Fever or systemic signs of infection
-- Wound expansion or deterioration
-- Delayed healing beyond expected timeframe
-
-## 📚 Clinical Guidelines Context
-{guideline_context[:800]}{'...' if len(guideline_context) > 800 else ''}
-
-## 🏥 Next Steps
-
-### Immediate Actions
-1. **Professional Consultation**: Schedule appointment with wound care specialist
-2. **Baseline Documentation**: Establish comprehensive baseline assessment
-3. **Treatment Plan**: Develop individualized care protocol
-4. **Patient Education**: Provide wound care instructions and warning signs
-
-### Follow-up Schedule
-- **Week 1**: Daily monitoring and assessment
-- **Week 2-4**: Every 2-3 days or as indicated
-- **Monthly**: Comprehensive reassessment and plan review
-- **As Needed**: Immediate evaluation for any concerning changes
-
-## ⚖️ Important Medical Disclaimer
-
-**This automated analysis is provided for informational and educational purposes only.**
-
-- This report does not constitute medical diagnosis or treatment advice
-- All measurements are computer-generated estimates requiring clinical verification
-- Professional medical evaluation is essential for proper diagnosis and treatment
-- This AI tool should supplement, not replace, clinical judgment
-- Always consult qualified healthcare professionals for medical decisions
-
-### Clinical Correlation Required
-- Verify all measurements with standard clinical tools
-- Correlate findings with patient symptoms and history
-- Consider factors not captured in automated analysis
-- Follow institutional protocols and guidelines
-
----
-*Generated by SmartHeal AI - Advanced Wound Care Analysis System*  
-*Report Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*  
-*Version: AI-Processor v1.2 with Enhanced Fallback Reporting*
+## 📊 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')}
+
+## 🎯 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.
 """
-        return report
 
+    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 MedGemma (GPU) → fallback report."""
+        try:
+            report = generate_medgemma_report_with_timeout(
+                patient_info, visual_results, guideline_context, image_pil, max_new_tokens
+            )
+            if report and report.strip() and not report.startswith(("⚠️", "❌")):
+                return report
+            logging.warning("MedGemma 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}")
+            return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+
+    # ---------- HF dataset commit ----------
     def save_and_commit_image(self, image_pil: Image.Image) -> str:
-        """Save image locally and optionally commit to HF dataset."""
+        """Save image locally and optionally upload to HF dataset."""
         try:
             os.makedirs(self.uploads_dir, exist_ok=True)
-            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-            filename = f"{timestamp}.png"
+            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filename = f"{ts}.png"
             path = os.path.join(self.uploads_dir, filename)
-            
-            # Save image
             image_pil.convert("RGB").save(path)
             logging.info(f"✅ Image saved locally: {path}")
-            
-            # Upload to HuggingFace dataset if configured
+
             if self.hf_token and self.dataset_id:
                 try:
+                    HfApi, HfFolder = _import_hf_hub()
+                    HfFolder.save_token(self.hf_token)
                     api = HfApi()
                     api.upload_file(
                         path_or_fileobj=path,
@@ -461,85 +483,92 @@ Monitor for signs requiring immediate attention:
                         repo_id=self.dataset_id,
                         repo_type="dataset",
                         token=self.hf_token,
-                        commit_message=f"Upload wound image: {filename}"
+                        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
-            
         except Exception as e:
-            logging.error(f"Failed to save image: {e}")
+            logging.error(f"Failed to save/commit image: {e}")
             return ""
 
-    def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: dict) -> dict:
-        """Run full analysis pipeline."""
+    # ---------- Orchestrator ----------
+    def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: Dict) -> Dict:
+        """End-to-end analysis with robust fallbacks."""
         try:
-            # Save image first
             saved_path = self.save_and_commit_image(image_pil)
-            logging.info(f"Image saved: {saved_path}")
-            
-            # Perform visual analysis
+
             visual_results = self.perform_visual_analysis(image_pil)
-            logging.info(f"Visual analysis completed: {visual_results}")
-            
-            # Process questionnaire data
-            patient_info = f"Age: {questionnaire_data.get('age', 'N/A')}, Diabetic: {questionnaire_data.get('diabetic', 'N/A')}, Allergies: {questionnaire_data.get('allergies', 'N/A')}, Date of Wound Sustained: {questionnaire_data.get('date_of_injury', 'N/A')}, Professional Care: {questionnaire_data.get('professional_care', 'N/A')}, Oozing/Bleeding: {questionnaire_data.get('oozing_bleeding', 'N/A')}, Infection: {questionnaire_data.get('infection', 'N/A')}, Moisture: {questionnaire_data.get('moisture', 'N/A')}"
-            
+
+            # Patient info summary text
+            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 guidelines
-            query = f"best practices for managing a {visual_results['wound_type']} with moisture level '{questionnaire_data.get('moisture', 'unknown')}' and signs of infection '{questionnaire_data.get('infection', 'unknown')}' in a patient who is diabetic '{questionnaire_data.get('diabetic', 'unknown')}'"
+            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')}'"
+            )
             guideline_context = self.query_guidelines(query)
-            logging.info("Guidelines queried successfully")
-            
+
             # Generate final report
-            report = self.generate_final_report(patient_info, visual_results, guideline_context, image_pil)
-            logging.info("Report generated successfully")
-            
+            report = self.generate_final_report(patient_info=patient_info,
+                                                visual_results=visual_results,
+                                                guideline_context=guideline_context,
+                                                image_pil=image_pil)
+
             return {
-                'success': True, 
-                'visual_analysis': visual_results, 
-                'report': report, 
-                'saved_image_path': saved_path,
-                'guideline_context': guideline_context[:500] + "..." if len(guideline_context) > 500 else guideline_context
+                "success": True,
+                "visual_analysis": visual_results,
+                "report": report,
+                "saved_image_path": saved_path,
+                "guideline_context": (guideline_context or "")[:500] + ("..." if guideline_context and len(guideline_context) > 500 else ""),
             }
-            
         except Exception as e:
             logging.error(f"Pipeline error: {e}")
             return {
-                'success': False, 
-                'error': str(e),
-                'visual_analysis': {},
-                'report': f"Analysis failed: {str(e)}",
-                'saved_image_path': None,
-                'guideline_context': ""
+                "success": False,
+                "error": str(e),
+                "visual_analysis": {},
+                "report": f"Analysis failed: {str(e)}",
+                "saved_image_path": None,
+                "guideline_context": "",
             }
 
-    def analyze_wound(self, image, questionnaire_data: dict) -> dict:
-        """Main analysis entry point - maintains original function name."""
+    def analyze_wound(self, image, questionnaire_data: Dict) -> Dict:
+        """Public entrypoint used by your UI."""
         try:
-            # Handle different image input formats
             if isinstance(image, str):
-                if os.path.exists(image):
-                    image_pil = Image.open(image)
-                else:
+                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)
             else:
                 raise ValueError(f"Unsupported image type: {type(image)}")
-            
-            return self.full_analysis_pipeline(image_pil, questionnaire_data)
-            
+
+            return self.full_analysis_pipeline(image_pil, questionnaire_data or {})
         except Exception as e:
             logging.error(f"Wound analysis error: {e}")
             return {
-                'success': False, 
-                'error': str(e),
-                'visual_analysis': {},
-                'report': f"Analysis initialization failed: {str(e)}",
-                'saved_image_path': None,
-                'guideline_context': ""
-            }
+                "success": False,
+                "error": str(e),
+                "visual_analysis": {},
+                "report": f"Analysis initialization failed: {str(e)}",
+                "saved_image_path": None,
+                "guideline_context": "",
+            }