Update src/ai_processor.py

src/ai_processor.py

@@ -1,5 +1,6 @@
 # smartheal_ai_processor.py
-# Full, functional module with conditional Spaces GPU support and CPU fallbacks.
+# Full, functional module with an always-present @spaces.GPU function (if `spaces` is importable)
+# and robust CPU fallbacks to avoid crashes when GPU isn't actually available yet.
 
 import os
 import time
@@ -11,32 +12,32 @@ import cv2
 import numpy as np
 from PIL import Image
 
-# =============== LOGGING SETUP ===============
+# =============== LOGGING ===============
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
-# =============== CONFIGURATION ===============
+# =============== CONFIG ===============
 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
+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"  # optional (set HF_TOKEN too)
-PIXELS_PER_CM = 38  # heuristic
+DATASET_ID = "SmartHeal/wound-image-uploads"   # optional (requires HF_TOKEN)
+PIXELS_PER_CM = 38
 
-# =============== GLOBAL CACHES ===============
+# =============== CACHES ===============
 models_cache: Dict[str, object] = {}
 knowledge_base_cache: Dict[str, object] = {}
 
-# ---------- Optional imports guarded ----------
+# =============== Optional imports (lazy) ===============
 def _import_ultralytics():
     from ultralytics import YOLO
     return YOLO
 
 def _import_tf_loader():
     import tensorflow as tf
-    tf.config.set_visible_devices([], "GPU")  # force CPU
+    tf.config.set_visible_devices([], "GPU")  # force CPU for TF
     from tensorflow.keras.models import load_model
     return load_model
 
@@ -60,39 +61,33 @@ 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 function (always defined if `spaces` import works) ===============
+try:
+    import spaces
 
     @spaces.GPU(enable_queue=True, duration=90)
-    def generate_medgemma_report_with_timeout(
+    def generate_medgemma_report(
         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
+        """
+        This function MUST exist at import time so Spaces Zero detects it.
+        It is guarded internally so if anything fails (no GPU yet, model load error),
+        it returns a warning and your pipeline will use the fallback report.
+        """
         try:
-            torch.cuda.empty_cache()
+            import torch
+            from transformers import pipeline
+
+            # Try to free cache; if no CUDA, this will raise and we return a warning.
+            try:
+                if hasattr(torch, "cuda") and torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+            except Exception:
+                pass
 
             prompt = f"""
 You are a medical AI assistant. Analyze this wound image and patient data.
@@ -110,21 +105,16 @@ Provide a structured report with:
             pipe = pipeline(
                 "image-text-to-text",
                 model="google/medgemma-4b-it",
-                torch_dtype=torch.bfloat16,
+                torch_dtype=getattr(torch, "bfloat16", None),
                 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},
-                    ],
-                }
-            ]
+            messages = [{"role": "user", "content": [
+                {"type": "image", "image": image_pil},
+                {"type": "text",  "text": prompt},
+            ]}]
 
             t0 = time.time()
             out = pipe(
@@ -134,10 +124,10 @@ Provide a structured report with:
                 temperature=0.7,
                 pad_token_id=pipe.tokenizer.eos_token_id,
             )
-            logging.info(f"✅ MedGemma completed in {time.time() - t0:.2f}s")
+            logging.info(f"✅ MedGemma finished in {time.time()-t0:.2f}s")
 
             if out and len(out) > 0:
-                # Defensive extraction
+                # Defensive extraction (different transformers versions)
                 try:
                     return out[0]["generated_text"][-1].get("content", "").strip() or "⚠️ Empty response"
                 except Exception:
@@ -145,13 +135,9 @@ Provide a structured report with:
             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:
+            return "⚠️ GPU worker unavailable"
+except Exception:
+    # If `spaces` cannot be imported locally, expose a CPU-safe stub with same signature.
     def generate_medgemma_report_with_timeout(
         patient_info: str,
         visual_results: Dict,
@@ -159,10 +145,9 @@ else:
         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) ===============
+# =============== Model init (CPU-safe) ===============
 def load_yolo_model():
     YOLO = _import_ultralytics()
     return YOLO(YOLO_MODEL_PATH)
@@ -173,32 +158,25 @@ def load_segmentation_model():
 
 def load_classification_pipeline():
     pipe = _import_hf_cls()
-    return pipe(
-        "image-classification",
-        model="Hemg/Wound-classification",
-        token=HF_TOKEN,
-        device="cpu",
-    )
+    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:
-    """Initialize all CPU-only models once with robust fallbacks."""
-    # Hugging Face auth (optional)
     if HF_TOKEN:
         try:
             HfApi, HfFolder = _import_hf_hub()
             HfFolder.save_token(HF_TOKEN)
-            logging.info("✅ HuggingFace token set")
+            logging.info("✅ HF 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()
-            logging.info("✅ YOLO model loaded (CPU)")
+            logging.info("✅ YOLO loaded (CPU)")
         except Exception as e:
             logging.error(f"YOLO load failed: {e}")
 
@@ -209,43 +187,41 @@ def initialize_cpu_models() -> None:
                 logging.info("✅ Segmentation model loaded (CPU)")
             else:
                 models_cache["seg"] = None
-                logging.warning("Segmentation model file not found; skipping seg.")
+                logging.warning("Segmentation model file missing; skipping.")
         except Exception as e:
             models_cache["seg"] = None
-            logging.warning(f"Segmentation model not available: {e}")
+            logging.warning(f"Segmentation unavailable: {e}")
 
     if "cls" not in models_cache:
         try:
             models_cache["cls"] = load_classification_pipeline()
-            logging.info("✅ Classification pipeline loaded (CPU)")
+            logging.info("✅ Classifier loaded (CPU)")
         except Exception as e:
             models_cache["cls"] = None
-            logging.warning(f"Classification pipeline not available: {e}")
+            logging.warning(f"Classifier unavailable: {e}")
 
     if "embedding_model" not in models_cache:
         try:
             models_cache["embedding_model"] = load_embedding_model()
-            logging.info("✅ Embedding model loaded (CPU)")
+            logging.info("✅ Embeddings loaded (CPU)")
         except Exception as e:
             models_cache["embedding_model"] = None
-            logging.warning(f"Embedding model not available: {e}")
+            logging.warning(f"Embeddings unavailable: {e}")
 
 def setup_knowledge_base() -> None:
-    """Load PDFs and create FAISS vector store (optional)."""
     if "vector_store" in knowledge_base_cache:
         return
 
-    docs = []
+    docs: List = []
     try:
         PyPDFLoader = _import_langchain_pdf()
         for pdf in GUIDELINE_PDFS:
             if os.path.exists(pdf):
                 try:
-                    loader = PyPDFLoader(pdf)
-                    docs.extend(loader.load())
+                    docs.extend(PyPDFLoader(pdf).load())
                     logging.info(f"Loaded PDF: {pdf}")
                 except Exception as e:
-                    logging.warning(f"Failed to load PDF {pdf}: {e}")
+                    logging.warning(f"PDF load failed ({pdf}): {e}")
     except Exception as e:
         logging.warning(f"LangChain PDF loader unavailable: {e}")
 
@@ -253,18 +229,17 @@ def setup_knowledge_base() -> None:
         try:
             from langchain.text_splitter import RecursiveCharacterTextSplitter
             FAISS = _import_langchain_faiss()
-            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-            chunks = splitter.split_documents(docs)
+            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 with {len(chunks)} chunks")
+            logging.info(f"✅ Knowledge base ready ({len(chunks)} chunks)")
         except Exception as e:
             knowledge_base_cache["vector_store"] = None
-            logging.warning(f"Knowledge base unavailable: {e}")
+            logging.warning(f"KB build failed: {e}")
     else:
         knowledge_base_cache["vector_store"] = None
-        logging.warning("Knowledge base disabled (no docs or embeddings).")
+        logging.warning("KB disabled (no docs or embeddings).")
 
-# Initialize on import
+# Initialize on import so app is ready
 initialize_cpu_models()
 setup_knowledge_base()
 
@@ -278,14 +253,13 @@ class AIProcessor:
         self.dataset_id = DATASET_ID
         self.hf_token = HF_TOKEN
 
-    # ---------- 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."""
+        """YOLO detect → (optional) Keras seg → (optional) HF classify → save visuals."""
         try:
             image_cv = cv2.cvtColor(np.array(image_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
 
@@ -338,7 +312,7 @@ class AIProcessor:
                         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}")
+                    logging.warning(f"Segmentation skipped: {e}")
 
             # Optional classification
             wound_type = "Unknown"
@@ -350,7 +324,7 @@ class AIProcessor:
                     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}")
+                    logging.warning(f"Classification failed: {e}")
 
             # Save detection & original
             out_dir = self._ensure_analysis_dir()
@@ -380,19 +354,17 @@ class AIProcessor:
             raise
 
     def query_guidelines(self, query: str) -> str:
-        """Query the knowledge base (optional)."""
+        """Query the (optional) guideline knowledge base."""
         try:
             vs = self.knowledge_base_cache.get("vector_store")
             if not vs:
                 return "Knowledge base is not available."
-            # 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)
+                docs = retriever.invoke(query)                 # older LC
             lines: List[str] = []
             for d in docs:
                 src = (d.metadata or {}).get("source", "N/A")
@@ -403,9 +375,7 @@ class AIProcessor:
             logging.warning(f"Guidelines query failed: {e}")
             return f"Guidelines query failed: {str(e)}"
 
-    # ---------- 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
@@ -431,11 +401,11 @@ Automated analysis provides quantitative measurements; verify via clinical exami
 - Document with serial photos and measurements
 
 ## 📅 Monitoring
-- Daily in week 1, then every 2-3 days (or as indicated)
+- 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 ''}
+{(guideline_context or '')[:800]}{"..." if guideline_context and len(guideline_context) > 800 else ''}
 
 **Disclaimer:** Automated, for decision support only. Verify clinically.
 """
@@ -448,7 +418,7 @@ Automated analysis provides quantitative measurements; verify via clinical exami
         image_pil: Image.Image,
         max_new_tokens: Optional[int] = None,
     ) -> str:
-        """Try MedGemma (GPU) → fallback report."""
+        """Use GPU path when available, fallback otherwise."""
         try:
             report = generate_medgemma_report_with_timeout(
                 patient_info, visual_results, guideline_context, image_pil, max_new_tokens
@@ -461,9 +431,8 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             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 upload to HF dataset."""
+        """Save locally and (optionally) upload to HF dataset."""
         try:
             os.makedirs(self.uploads_dir, exist_ok=True)
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -472,17 +441,17 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             image_pil.convert("RGB").save(path)
             logging.info(f"✅ Image saved locally: {path}")
 
-            if self.hf_token and self.dataset_id:
+            if HF_TOKEN and DATASET_ID:
                 try:
                     HfApi, HfFolder = _import_hf_hub()
-                    HfFolder.save_token(self.hf_token)
+                    HfFolder.save_token(HF_TOKEN)
                     api = HfApi()
                     api.upload_file(
                         path_or_fileobj=path,
                         path_in_repo=f"images/{filename}",
-                        repo_id=self.dataset_id,
+                        repo_id=DATASET_ID,
                         repo_type="dataset",
-                        token=self.hf_token,
+                        token=HF_TOKEN,
                         commit_message=f"Upload wound image: {filename}",
                     )
                     logging.info("✅ Image committed to HF dataset")
@@ -494,28 +463,24 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             logging.error(f"Failed to save/commit image: {e}")
             return ""
 
-    # ---------- Orchestrator ----------
     def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: Dict) -> Dict:
-        """End-to-end analysis with robust fallbacks."""
+        """End-to-end analysis."""
         try:
             saved_path = self.save_and_commit_image(image_pil)
-
             visual_results = self.perform_visual_analysis(image_pil)
 
-            # 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')}"
+                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.get('wound_type','Unknown')} "
                 f"with moisture '{pi.get('moisture','unknown')}' and infection '{pi.get('infection','unknown')}' "
@@ -523,18 +488,16 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             )
             guideline_context = self.query_guidelines(query)
 
-            # Generate final report
-            report = self.generate_final_report(patient_info=patient_info,
-                                                visual_results=visual_results,
-                                                guideline_context=guideline_context,
-                                                image_pil=image_pil)
+            report = self.generate_final_report(patient_info, visual_results, guideline_context, image_pil)
 
             return {
                 "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 ""),
+                "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}")
@@ -548,7 +511,7 @@ Automated analysis provides quantitative measurements; verify via clinical exami
             }
 
     def analyze_wound(self, image, questionnaire_data: Dict) -> Dict:
-        """Public entrypoint used by your UI."""
+        """Public entrypoint used by UI."""
         try:
             if isinstance(image, str):
                 if not os.path.exists(image):
@@ -571,4 +534,4 @@ Automated analysis provides quantitative measurements; verify via clinical exami
                 "report": f"Analysis initialization failed: {str(e)}",
                 "saved_image_path": None,
                 "guideline_context": "",
-            }
+            }