Update src/ai_processor.py

src/ai_processor.py

@@ -1,7 +1,11 @@
 import os
-# Ensure all CPU-only models never touch CUDA
+# Force CPU-only until we enter the GPU context
 os.environ['CUDA_VISIBLE_DEVICES'] = ''
 
+import torch
+# Prevent any CUDA initialization in the main process
+torch.cuda.is_available = lambda: False
+
 import io
 import base64
 import logging
@@ -20,7 +24,6 @@ from huggingface_hub import HfApi, HfFolder
 import spaces
 from .config import Config
 
-# System prompt for MedGemma
 default_system_prompt = (
     "You are a world-class medical AI assistant specializing in wound care "
     "with expertise in wound assessment and treatment. Provide concise, "
@@ -42,11 +45,7 @@ def generate_medgemma_report(
     segmentation_image_path: str,
     max_new_tokens: int = None
 ) -> str:
-    """
-    Runs on GPU. Lazy-loads the MedGemma pipeline and returns the markdown report.
-    Accepts only primitive types and file-paths, so pickling works.
-    """
-    # Lazy-load pipeline
+    # Lazy-load HF pipeline inside GPU context
     if not hasattr(generate_medgemma_report, "_pipe"):
         try:
             cfg = Config()
@@ -65,16 +64,13 @@ def generate_medgemma_report(
 
     pipe = generate_medgemma_report._pipe
 
-    # Assemble messages
     msgs = [
         {'role':'system','content':[{'type':'text','text':default_system_prompt}]},
         {'role':'user','content':[]}
     ]
-    # Attach images
     for path in (detection_image_path, segmentation_image_path):
         if path and os.path.exists(path):
             msgs[1]['content'].append({'type':'image','image': Image.open(path)})
-    # Attach text
     prompt = f"## Patient\n{patient_info}\n## Wound Type: {visual_results.get('wound_type','Unknown')}"
     msgs[1]['content'].append({'type':'text','text': prompt})
 
@@ -96,28 +92,24 @@ class AIProcessor:
         self._load_knowledge_base()
 
     def _initialize_models(self):
-        """Load all CPU-only models here."""
-        # Set HuggingFace token
         if self.config.HF_TOKEN:
             HfFolder.save_token(self.config.HF_TOKEN)
             logging.info("✅ HuggingFace token set")
 
-        # YOLO detection (CPU)
         try:
+            # YOLO on CPU only (no CUDA)
             self.models_cache['det'] = YOLO(self.config.YOLO_MODEL_PATH)
             logging.info("✅ YOLO model loaded (CPU only)")
         except Exception as e:
             logging.error(f"YOLO load failed: {e}")
             raise
 
-        # Segmentation (CPU)
         try:
             self.models_cache['seg'] = load_model(self.config.SEG_MODEL_PATH, compile=False)
             logging.info("✅ Segmentation model loaded (CPU)")
         except Exception as e:
             logging.warning(f"Segmentation model not available: {e}")
 
-        # Classification (CPU)
         try:
             self.models_cache['cls'] = pipeline(
                 'image-classification',
@@ -129,7 +121,6 @@ class AIProcessor:
         except Exception as e:
             logging.warning(f"Classification pipeline not available: {e}")
 
-        # Embedding model (CPU)
         try:
             self.models_cache['embedding_model'] = HuggingFaceEmbeddings(
                 model_name='sentence-transformers/all-MiniLM-L6-v2',
@@ -140,26 +131,22 @@ class AIProcessor:
             logging.warning(f"Embedding model not available: {e}")
 
     def _load_knowledge_base(self):
-        """Load PDF guidelines into a FAISS vector store."""
         docs = []
         for pdf in self.config.GUIDELINE_PDFS:
             if os.path.exists(pdf):
-                loader = PyPDFLoader(pdf)
-                docs.extend(loader.load())
+                docs.extend(PyPDFLoader(pdf).load())
                 logging.info(f"Loaded PDF: {pdf}")
 
         if docs and 'embedding_model' in self.models_cache:
             splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
             chunks = splitter.split_documents(docs)
-            vs = FAISS.from_documents(chunks, self.models_cache['embedding_model'])
-            self.knowledge_base_cache['vectorstore'] = vs
+            self.knowledge_base_cache['vectorstore'] = FAISS.from_documents(chunks, self.models_cache['embedding_model'])
             logging.info(f"✅ Knowledge base loaded ({len(chunks)} chunks)")
         else:
             self.knowledge_base_cache['vectorstore'] = None
             logging.warning("Knowledge base unavailable")
 
     def perform_visual_analysis(self, image_pil: Image.Image) -> dict:
-        """Detect & segment on CPU; return metrics + file paths."""
         if 'det' not in self.models_cache:
             raise RuntimeError("YOLO model ('det') not loaded")
 
@@ -171,7 +158,6 @@ class AIProcessor:
         x1, y1, x2, y2 = res.boxes.xyxy[0].cpu().numpy().astype(int)
         region = img_cv[y1:y2, x1:x2]
 
-        # Save detection overlay
         det_vis = img_cv.copy()
         cv2.rectangle(det_vis, (x1, y1), (x2, y2), (0,255,0), 2)
         os.makedirs(f"{self.config.UPLOADS_DIR}/analysis", exist_ok=True)
@@ -179,7 +165,6 @@ class AIProcessor:
         det_path = f"{self.config.UPLOADS_DIR}/analysis/detection_{ts}.png"
         cv2.imwrite(det_path, det_vis)
 
-        # Segmentation metrics
         length = breadth = area = 0
         seg_path = None
         if 'seg' in self.models_cache:
@@ -199,7 +184,6 @@ class AIProcessor:
                 breadth= round(w0/self.px_per_cm,2)
                 area   = round(cv2.contourArea(cnt)/(self.px_per_cm**2),2)
 
-        # Classification
         wound_type = 'Unknown'
         if 'cls' in self.models_cache:
             try:
@@ -238,9 +222,6 @@ class AIProcessor:
         image_pil: Image.Image,
         max_new_tokens: int = None
     ) -> str:
-        """
-        Signature unchanged. Gathers arguments, calls GPU function, and falls back if needed.
-        """
         det = visual_results.get('detection_image_path', '')
         seg = visual_results.get('segmentation_image_path', '')
         report = generate_medgemma_report(
@@ -314,59 +295,34 @@ class AIProcessor:
             image = Image.open(image)
         return self.full_analysis_pipeline(image, questionnaire_data)
 
-    def _assess_risk_legacy(self, questionnaire_data):
-        """Legacy risk assessment for backward compatibility"""
-        risk_factors = []
-        risk_score = 0
-        
+    def _assess_risk_legacy(self, questionnaire_data: dict) -> dict:
+        risk_factors, risk_score = [], 0
         try:
-            # Age factor
             age = questionnaire_data.get('patient_age', 0)
             if age > 65:
-                risk_factors.append("Advanced age (>65)")
-                risk_score += 2
+                risk_factors.append("Advanced age (>65)"); risk_score += 2
             elif age > 50:
-                risk_factors.append("Older adult (50-65)")
-                risk_score += 1
-            
-            # Duration factor
-            duration = questionnaire_data.get('wound_duration', '').lower()
-            if any(term in duration for term in ['month', 'months', 'year']):
-                risk_factors.append("Chronic wound (>4 weeks)")
-                risk_score += 3
-            
-            # Pain level
-            pain_level = questionnaire_data.get('pain_level', 0)
-            if pain_level >= 7:
-                risk_factors.append("High pain level")
-                risk_score += 2
-            
-            # Medical history risk factors
-            medical_history = 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
-            
-            # Determine risk level
-            if risk_score >= 7:
-                risk_level = "High"
-            elif risk_score >= 4:
-                risk_level = "Moderate"
-            else:
-                risk_level = "Low"
-            
-            return {
-                'risk_score': risk_score,
-                'risk_level': risk_level,
-                'risk_factors': risk_factors
-            }
-            
+                risk_factors.append("Older adult (50-65)"); risk_score += 1
+            dur = questionnaire_data.get('wound_duration','').lower()
+            if any(t in dur for t in ['month','year']):
+                risk_factors.append("Chronic wound (>4 weeks)"); risk_score += 3
+            pain = questionnaire_data.get('pain_level', 0)
+            if pain >= 7:
+                risk_factors.append("High pain level"); risk_score += 2
+            hist = questionnaire_data.get('medical_history','').lower()
+            if 'diabetes' in hist:
+                risk_factors.append("Diabetes mellitus"); risk_score += 3
+            if 'vascular' in hist:
+                risk_factors.append("Vascular issues"); risk_score += 2
+            if 'immune' in hist:
+                risk_factors.append("Immune compromise"); risk_score += 2
+
+            level = (
+                "High" if risk_score >= 7 else
+                "Moderate" if risk_score >= 4 else
+                "Low"
+            )
+            return {'risk_score': risk_score, 'risk_level': 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': []}