Update src/ai_processor.py

src/ai_processor.py

@@ -24,6 +24,7 @@ from huggingface_hub import HfApi, HfFolder
 import spaces
 from .config import Config
 
+# Inline system prompt for MedGemma GPU pipeline
 default_system_prompt = (
     "You are a world-class medical AI assistant specializing in wound care "
     "with expertise in wound assessment and treatment. Provide concise, "
@@ -45,7 +46,7 @@ def generate_medgemma_report(
     segmentation_image_path: str,
     max_new_tokens: int = None
 ) -> str:
-    # Lazy-load HF pipeline inside GPU context
+    """Runs on GPU. Lazy-loads the MedGemma pipeline and returns the markdown report."""
     if not hasattr(generate_medgemma_report, "_pipe"):
         try:
             cfg = Config()
@@ -64,22 +65,27 @@ def generate_medgemma_report(
 
     pipe = generate_medgemma_report._pipe
 
+    # Assemble messages
     msgs = [
-        {'role':'system','content':[{'type':'text','text':default_system_prompt}]},
-        {'role':'user','content':[]}
+        {'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)})
+            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})
+    msgs[1]['content'].append({'type': 'text', 'text': prompt})
 
     out = pipe(
         text=msgs,
         max_new_tokens=max_new_tokens or Config().MAX_NEW_TOKENS,
         do_sample=False
     )
-    return out[0]['generated_text'][-1].get('content','')
+    return out[0]['generated_text'][-1].get('content', '')
 
 
 class AIProcessor:
@@ -92,24 +98,28 @@ 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-only)
         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 model (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 pipeline (CPU)
         try:
             self.models_cache['cls'] = pipeline(
                 'image-classification',
@@ -121,32 +131,38 @@ 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',
-                model_kwargs={'device':'cpu'}
+                model_kwargs={'device': 'cpu'}
             )
             logging.info("✅ Embedding model loaded (CPU)")
         except Exception as e:
             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):
-                docs.extend(PyPDFLoader(pdf).load())
+                loader = PyPDFLoader(pdf)
+                docs.extend(loader.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)
-            self.knowledge_base_cache['vectorstore'] = FAISS.from_documents(chunks, self.models_cache['embedding_model'])
+            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")
 
@@ -158,6 +174,7 @@ 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)
@@ -165,31 +182,31 @@ class AIProcessor:
         det_path = f"{self.config.UPLOADS_DIR}/analysis/detection_{ts}.png"
         cv2.imwrite(det_path, det_vis)
 
+        # Segmentation
         length = breadth = area = 0
         seg_path = None
         if 'seg' in self.models_cache:
             h, w = self.models_cache['seg'].input_shape[1:3]
-            inp = cv2.resize(region, (w,h)) / 255.0
+            inp = cv2.resize(region, (w, h)) / 255.0
             mask = (self.models_cache['seg'].predict(inp[None])[0,:,:,0] > 0.5).astype(np.uint8)
             mask_rs = cv2.resize(mask, (region.shape[1], region.shape[0]), interpolation=cv2.INTER_NEAREST)
             ov = region.copy(); ov[mask_rs==1] = [0,0,255]
-            seg_vis = cv2.addWeighted(region,0.7,ov,0.3,0)
+            seg_vis = cv2.addWeighted(region, 0.7, ov, 0.3, 0)
             seg_path = f"{self.config.UPLOADS_DIR}/analysis/segmentation_{ts}.png"
             cv2.imwrite(seg_path, seg_vis)
             cnts, _ = cv2.findContours(mask_rs, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
             if cnts:
                 cnt = max(cnts, key=cv2.contourArea)
-                _,_,w0,h0 = cv2.boundingRect(cnt)
-                length = round(h0/self.px_per_cm,2)
-                breadth= round(w0/self.px_per_cm,2)
-                area   = round(cv2.contourArea(cnt)/(self.px_per_cm**2),2)
+                _, _, w0, h0 = cv2.boundingRect(cnt)
+                length = round(h0 / self.px_per_cm, 2)
+                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:
-                preds = self.models_cache['cls'](
-                    Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB))
-                )
+                preds = self.models_cache['cls'](Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB)))
                 wound_type = max(preds, key=lambda x: x['score'])['label']
             except Exception:
                 pass
@@ -225,12 +242,8 @@ class AIProcessor:
         det = visual_results.get('detection_image_path', '')
         seg = visual_results.get('segmentation_image_path', '')
         report = generate_medgemma_report(
-            patient_info,
-            visual_results,
-            guideline_context,
-            det,
-            seg,
-            max_new_tokens
+            patient_info, visual_results, guideline_context,
+            det, seg, max_new_tokens
         )
         if report:
             return report
@@ -269,23 +282,14 @@ class AIProcessor:
                 logging.warning(f"HF upload failed: {e}")
         return path
 
-    def full_analysis_pipeline(
-        self,
-        image_pil: Image.Image,
-        questionnaire_data: dict
-    ) -> dict:
+    def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: dict) -> dict:
         try:
             saved = self.save_and_commit_image(image_pil)
-            vis = self.perform_visual_analysis(image_pil)
-            info = ", ".join(f"{k}:{v}" for k,v in questionnaire_data.items() if v)
-            gc = self.query_guidelines(info)
-            report = self.generate_final_report(info, vis, gc, image_pil)
-            return {
-                'success': True,
-                'visual_analysis': vis,
-                'report': report,
-                'saved_image_path': saved
-            }
+            vis   = self.perform_visual_analysis(image_pil)
+            info  = ", ".join(f"{k}:{v}" for k,v in questionnaire_data.items() if v)
+            gc    = self.query_guidelines(info)
+            report= self.generate_final_report(info, vis, gc, image_pil)
+            return {'success': True, 'visual_analysis': vis, 'report': report, 'saved_image_path': saved}
         except Exception as e:
             logging.error(f"Pipeline error: {e}")
             return {'success': False, 'error': str(e)}
@@ -303,7 +307,7 @@ class AIProcessor:
                 risk_factors.append("Advanced age (>65)"); risk_score += 2
             elif age > 50:
                 risk_factors.append("Older adult (50-65)"); risk_score += 1
-            dur = questionnaire_data.get('wound_duration','').lower()
+            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)
@@ -316,13 +320,8 @@ class AIProcessor:
                 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"
-            )
+            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': []}
+            return {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []}