Update src/ai_processor.py

src/ai_processor.py

@@ -17,7 +17,6 @@ import spaces
 from .config import Config
 import os
 
-# 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, "
@@ -40,21 +39,31 @@ class AIProcessor:
 
     def _initialize_models(self):
         """Initialize CPU-only AI models; MedGemma is loaded on demand within GPU context."""
-        # Set HuggingFace token early
+        # Set HF token
         if self.config.HF_TOKEN:
             HfFolder.save_token(self.config.HF_TOKEN)
             logging.info("HuggingFace token set successfully")
 
-        # YOLO detection on CPU (force CPU to avoid CUDA init)
+        # YOLO detection on CPU (force CPU)
         try:
-            self.models_cache['det'] = YOLO(self.config.YOLO_MODEL_PATH, device='cpu')
+            self.models_cache['det'] = YOLO(
+                self.config.YOLO_MODEL_PATH,
+                device='cpu'
+            )
             logging.info("✅ YOLO detection model loaded on CPU")
         except Exception as e:
-            logging.warning(f"YOLO model not available: {e}")
+            logging.error(
+                f"Failed to load YOLO model at '{self.config.YOLO_MODEL_PATH}': {e}"
+            )
+            # fail fast so you’ll immediately see why 'det' was never set
+            raise
 
         # Segmentation model on CPU
         try:
-            self.models_cache['seg'] = load_model(self.config.SEG_MODEL_PATH, compile=False)
+            self.models_cache['seg'] = load_model(
+                self.config.SEG_MODEL_PATH,
+                compile=False
+            )
             logging.info("✅ Segmentation model loaded on CPU")
         except Exception as e:
             logging.warning(f"Segmentation model not available: {e}")
@@ -71,7 +80,7 @@ class AIProcessor:
         except Exception as e:
             logging.warning(f"Wound classification model not available: {e}")
 
-        # Embedding for knowledge base on CPU
+        # Embeddings for KB on CPU
         try:
             self.models_cache['embedding_model'] = HuggingFaceEmbeddings(
                 model_name='sentence-transformers/all-MiniLM-L6-v2',
@@ -81,11 +90,10 @@ class AIProcessor:
         except Exception as e:
             logging.warning(f"Embedding model not available: {e}")
 
-        # Load knowledge base
+        # Load PDF guidelines
         self._load_knowledge_base()
 
     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):
@@ -94,92 +102,132 @@ class AIProcessor:
                 logging.info(f"Loaded PDF: {pdf}")
 
         if docs and 'embedding_model' in self.models_cache:
-            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+            splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1000, chunk_overlap=100
+            )
             chunks = splitter.split_documents(docs)
-            vs = FAISS.from_documents(chunks, self.models_cache['embedding_model'])
+            vs = FAISS.from_documents(
+                chunks, self.models_cache['embedding_model']
+            )
             self.knowledge_base_cache['vectorstore'] = vs
             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):
         """Detect & segment on CPU; return only paths + metrics."""
+        if 'det' not in self.models_cache:
+            raise RuntimeError(
+                "YOLO detection model ('det') not loaded; cannot perform visual analysis"
+            )
+
         try:
-            img_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
-            # YOLO detect
-            res = self.models_cache['det'].predict(img_cv, verbose=False)[0]
+            img_cv = cv2.cvtColor(
+                np.array(image_pil), cv2.COLOR_RGB2BGR
+            )
+            res = self.models_cache['det'].predict(
+                img_cv, verbose=False
+            )[0]
             if not res.boxes:
                 raise ValueError("No wound detected")
+
             # Bounding box
             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)
+            cv2.rectangle(det_vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
             os.makedirs(f"{self.config.UPLOADS_DIR}/analysis", exist_ok=True)
             ts = datetime.now().strftime('%Y%m%d_%H%M%S')
-            det_path = f"{self.config.UPLOADS_DIR}/analysis/detection_{ts}.png"
+            det_path = (
+                f"{self.config.UPLOADS_DIR}/analysis/detection_{ts}.png"
+            )
             cv2.imwrite(det_path, det_vis)
+
             # Initialize metrics & seg
             length = breadth = area = 0
             seg_path = None
-            # Segmentation
+
+            # Segmentation (if available)
             if 'seg' in self.models_cache:
                 h, w = self.models_cache['seg'].input_shape[1:3]
-                inp = cv2.resize(region, (w,h)) / 255.0
-                mask = (self.models_cache['seg'].predict(np.expand_dims(inp,0))[0,:,:,0] > 0.5).astype(np.uint8)
-                mask_rs = cv2.resize(mask, (region.shape[1], region.shape[0]), interpolation=cv2.INTER_NEAREST)
+                inp = cv2.resize(region, (w, h)) / 255.0
+                mask = (
+                    self.models_cache['seg']
+                    .predict(np.expand_dims(inp, 0))[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_path = f"{self.config.UPLOADS_DIR}/analysis/segmentation_{ts}.png"
+                ov[mask_rs == 1] = [0, 0, 255]
+                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)
-                # measure
-                cnts, _ = cv2.findContours(mask_rs, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+                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:
-                    label = self.models_cache['cls'](Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB)))
-                    wound_type = max(label, key=lambda x: x['score'])['label']
+                    labels = self.models_cache['cls'](
+                        Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB))
+                    )
+                    wound_type = max(labels, key=lambda x: x['score'])['label']
                 except Exception:
                     pass
+
             return {
                 'wound_type': wound_type,
                 'length_cm': length,
                 'breadth_cm': breadth,
                 'surface_area_cm2': area,
-                'detection_confidence': float(res.boxes.conf[0].cpu().item()),
+                'detection_confidence': float(
+                    res.boxes.conf[0].cpu().item()
+                ),
                 'detection_image_path': det_path,
                 'segmentation_image_path': seg_path
             }
+
         except Exception as e:
             logging.error(f"Visual analysis error: {e}")
             raise
 
     def query_guidelines(self, query: str):
-        """Retrieve clinical guidelines from vectorstore."""
         vs = self.knowledge_base_cache.get('vectorstore')
         if not vs:
             return "Clinical guidelines unavailable"
         docs = vs.as_retriever(search_kwargs={'k':10}).invoke(query)
         return '\n\n'.join(
-            f"Source: {d.metadata.get('source','?')}, Page: {d.metadata.get('page','?')}\n{d.page_content}"
+            f"Source: {d.metadata.get('source','?')}, Page: {d.metadata.get('page','?')}\n"
+            f"{d.page_content}"
             for d in docs
         )
 
     @spaces.GPU(enable_queue=True, duration=120)
-    def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
+    def generate_final_report(
+        self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None
+    ):
         """Run MedGemma on GPU; return markdown report."""
-        # lazy-load MedGemma pipeline here to avoid CUDA init in main process
         if 'medgemma_pipe' not in self.models_cache:
             try:
                 self.models_cache['medgemma_pipe'] = pipeline(
@@ -193,22 +241,26 @@ class AIProcessor:
                 logging.info("✅ MedGemma pipeline loaded on GPU")
             except Exception as e:
                 logging.warning(f"MedGemma pipeline not available: {e}")
-                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+                return self._generate_fallback_report(
+                    patient_info, visual_results, guideline_context
+                )
 
-        # build messages
         msgs = [
             {'role':'system','content':[{'type':'text','text':default_system_prompt}]},
             {'role':'user','content':[]}
         ]
-        # images
         if image_pil:
             msgs[1]['content'].append({'type':'image','image':image_pil})
         for key in ('detection_image_path','segmentation_image_path'):
             p = visual_results.get(key)
             if p and os.path.exists(p):
-                msgs[1]['content'].append({'type':'image','image':Image.open(p)})
-        # text prompt
-        prompt = f"## Patient\n{patient_info}\n## Wound Type: {visual_results['wound_type']}"
+                msgs[1]['content'].append(
+                    {'type':'image','image':Image.open(p)}
+                )
+        prompt = (
+            f"## Patient\n{patient_info}\n"
+            f"## Wound Type: {visual_results['wound_type']}"
+        )
         msgs[1]['content'].append({'type':'text','text':prompt})
 
         out = self.models_cache['medgemma_pipe'](
@@ -222,7 +274,6 @@ class AIProcessor:
         )
 
     def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
-        """Produce text-only fallback."""
         dp = visual_results.get('detection_image_path','N/A')
         sp = visual_results.get('segmentation_image_path','N/A')
         return (
@@ -234,12 +285,11 @@ class AIProcessor:
         )
 
     def save_and_commit_image(self, image_pil):
-        """Save locally and optionally to HuggingFace."""
         os.makedirs(self.config.UPLOADS_DIR, exist_ok=True)
         fn = f"{datetime.now():%Y%m%d_%H%M%S}.png"
         path = os.path.join(self.config.UPLOADS_DIR, fn)
         image_pil.convert('RGB').save(path)
-        if self.config.HF_TOKEN and hasattr(self.config,'DATASET_ID') and self.config.DATASET_ID:
+        if self.config.HF_TOKEN and getattr(self.config,'DATASET_ID', None):
             try:
                 api = HfApi()
                 api.upload_file(
@@ -253,77 +303,51 @@ class AIProcessor:
         return path
 
     def full_analysis_pipeline(self, image_pil, questionnaire_data):
-        """Orchestrate CPU steps + GPU report."""
         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)
+            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}
+            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)}
+            return {'success': False, 'error': str(e)}
 
     def analyze_wound(self, image, questionnaire_data):
-        """Legacy wrapper."""
-        if isinstance(image,str):
+        if isinstance(image, str):
             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
-        
+        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
+            duration = questionnaire_data.get('wound_duration','').lower()
+            if any(term in duration for term 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
+
+            if   risk_score>=7: level="High"
+            elif risk_score>=4: level="Moderate"
+            else: level="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':[]}