Update src/ai_processor.py

src/ai_processor.py

@@ -36,17 +36,14 @@ knowledge_base_cache = {}
 
 # =============== LAZY LOADING FUNCTIONS (CPU-SAFE) ===============
 def load_yolo_model(yolo_model_path):
-    """Lazy import and load YOLO model to avoid CUDA initialization."""
     from ultralytics import YOLO
     return YOLO(yolo_model_path)
 
 def load_segmentation_model(seg_model_path):
-    """Lazy import and load segmentation model."""
     from tensorflow.keras.models import load_model
     return load_model(seg_model_path, compile=False)
 
 def load_classification_pipeline(hf_token):
-    """Lazy import and load classification pipeline (CPU only)."""
     from transformers import pipeline
     return pipeline(
         "image-classification",
@@ -56,7 +53,6 @@ def load_classification_pipeline(hf_token):
     )
 
 def load_embedding_model():
-    """Load embedding model for knowledge base."""
     return HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2",
         model_kwargs={"device": "cpu"}
@@ -64,34 +60,28 @@ def load_embedding_model():
 
 # =============== MODEL INITIALIZATION ===============
 def initialize_cpu_models():
-    """Initialize all CPU-only models once."""
     global models_cache
-    
     if HF_TOKEN:
         HfFolder.save_token(HF_TOKEN)
         logging.info("✅ HuggingFace token set")
-    
     if "det" not in models_cache:
         try:
             models_cache["det"] = load_yolo_model(YOLO_MODEL_PATH)
             logging.info("✅ YOLO model loaded (CPU only)")
         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)")
         except Exception as e:
             logging.warning(f"Segmentation model not available: {e}")
-
     if "cls" not in models_cache:
         try:
             models_cache["cls"] = load_classification_pipeline(HF_TOKEN)
             logging.info("✅ Classification pipeline loaded (CPU)")
         except Exception as e:
             logging.warning(f"Classification pipeline not available: {e}")
-
     if "embedding_model" not in models_cache:
         try:
             models_cache["embedding_model"] = load_embedding_model()
@@ -100,11 +90,9 @@ def initialize_cpu_models():
             logging.warning(f"Embedding model not available: {e}")
 
 def setup_knowledge_base():
-    """Load PDF documents and create FAISS vector store."""
     global knowledge_base_cache
     if "vector_store" in knowledge_base_cache:
         return
-
     docs = []
     for pdf_path in GUIDELINE_PDFS:
         if os.path.exists(pdf_path):
@@ -114,7 +102,6 @@ def setup_knowledge_base():
                 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)
@@ -124,7 +111,7 @@ def setup_knowledge_base():
         knowledge_base_cache["vector_store"] = None
         logging.warning("Knowledge base unavailable")
 
-# Initialize models on app startup
+# Initialize models at startup
 initialize_cpu_models()
 setup_knowledge_base()
 
@@ -138,8 +125,6 @@ def generate_medgemma_report(
     segmentation_image_path,
     max_new_tokens=None,
 ):
-    """GPU-only function for MedGemma report generation."""
-    # Import GPU libraries ONLY here
     import torch
     from transformers import pipeline
     from PIL import Image
@@ -156,7 +141,6 @@ def generate_medgemma_report(
         "patient context."
     )
 
-    # Lazy-load MedGemma pipeline on GPU
     if not hasattr(generate_medgemma_report, "_pipe"):
         try:
             generate_medgemma_report._pipe = pipeline(
@@ -174,18 +158,15 @@ def generate_medgemma_report(
 
     pipe = generate_medgemma_report._pipe
 
-    # Compose messages
     msgs = [
         {"role": "system", "content": [{"type": "text", "text": default_system_prompt}]},
         {"role": "user", "content": []},
     ]
 
-    # Attach images if available
     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
     prompt = f"## Patient\n{patient_info}\n## Wound Type: {visual_results.get('wound_type','Unknown')}"
     msgs[1]["content"].append({"type": "text", "text": prompt})
 
@@ -207,7 +188,6 @@ class AIProcessor:
         self.hf_token = HF_TOKEN
 
     def perform_visual_analysis(self, image_pil: Image.Image) -> dict:
-        """Detect & segment on CPU; return metrics + file paths."""
         img_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
         yolo = self.models_cache.get("det")
         if yolo is None:
@@ -217,40 +197,57 @@ class AIProcessor:
         if not res.boxes:
             raise ValueError("No wound detected")
 
-        x1, y1, x2, y2 = res.boxes.xyxy.cpu().numpy().astype(int)
-        region = img_cv[y1:y2, x1:x2]
+        # Safely unpack detection boxes
+        try:
+            xyxy = res.boxes.xyxy.cpu().numpy()
+            if xyxy.shape[0] == 0:
+                raise ValueError("No detection boxes found")
+            x1, y1, x2, y2 = xyxy[0]
+        except Exception as e:
+            logging.warning(f"Error unpacking detection boxes: {e}")
+            raise
+
+        region = img_cv[int(y1):int(y2), int(x1):int(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, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
         os.makedirs(f"{self.uploads_dir}/analysis", exist_ok=True)
         ts = datetime.now().strftime("%Y%m%d_%H%M%S")
         det_path = f"{self.uploads_dir}/analysis/detection_{ts}.png"
         cv2.imwrite(det_path, det_vis)
 
         # Segmentation
-        length = breadth = area = 0
+        length, breadth, area = 0, 0, 0
         seg_path = None
         seg_model = self.models_cache.get("seg")
         if seg_model:
-            h, w = seg_model.input_shape[1:3]
-            inp = cv2.resize(region, (w, h)) / 255.0
-            mask = (seg_model.predict(inp[None])[0, :, :, 0] > 0.5).astype(np.uint8)
-            mask_rs = cv2.resize(mask, (region.shape[1], region.shape), 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.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)
+            try:
+                h, w = seg_model.input_shape[1:3]
+                inp = cv2.resize(region, (w, h)) / 255.0
+                mask_pred = seg_model.predict(inp[None])
+                if mask_pred.shape[1:3] != (h, w):
+                    # Resize if needed
+                    mask_pred = np.squeeze(mask_pred)
+                mask = (mask_pred[0, :, :, 0] > 0.5).astype(np.uint8)
+                mask_rs = cv2.resize(mask, (region.shape[1], region.shape[0]), interpolation=cv2.INTER_NEAREST)
+                # Save segmentation visualization
+                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.uploads_dir}/analysis/segmentation_{ts}.png"
+                cv2.imwrite(seg_path, seg_vis)
+
+                # Find contours
+                cnts, _ = cv2.findContours(mask_rs, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                if cnts:
+                    cnt = max(cnts, key=cv2.contourArea)
+                    x, y, w_box, h_box = cv2.boundingRect(cnt)
+                    length = round(h_box / self.px_per_cm, 2)
+                    breadth = round(w_box / self.px_per_cm, 2)
+                    area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
+            except Exception as e:
+                logging.warning(f"Segmentation processing error: {e}")
 
         # Classification
         wound_type = "Unknown"
@@ -258,9 +255,10 @@ class AIProcessor:
         if cls_pipe:
             try:
                 preds = cls_pipe(Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB)))
-                wound_type = max(preds, key=lambda x: x["score"])["label"]
-            except Exception:
-                pass
+                if preds:
+                    wound_type = max(preds, key=lambda x: x["score"])["label"]
+            except Exception as e:
+                logging.warning(f"Classification error: {e}")
 
         return {
             "wound_type": wound_type,
@@ -273,7 +271,6 @@ class AIProcessor:
         }
 
     def query_guidelines(self, query: str) -> str:
-        """Query the knowledge base for relevant information."""
         vs = self.knowledge_base_cache.get("vector_store")
         if not vs:
             return "Clinical guidelines unavailable"
@@ -282,39 +279,28 @@ class AIProcessor:
             f"Source: {d.metadata.get('source','?')}, Page: {d.metadata.get('page','?')}\n{d.page_content}" for d in docs
         )
 
-    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 GPU pipeline."""
-        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)
+    def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
+        det_path = visual_results.get("detection_image_path", "")
+        seg_path = visual_results.get("segmentation_image_path", "")
+        report = generate_medgemma_report(patient_info, visual_results, guideline_context, det_path, seg_path, max_new_tokens)
         if report:
             return report
         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 fallback report if MedGemma fails."""
-        dp = visual_results.get('detection_image_path','N/A')
-        sp = visual_results.get('segmentation_image_path','N/A')
+    def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
         return (
             f"# Fallback Report\n{patient_info}\n"
             f"Type: {visual_results.get('wound_type','Unknown')}\n"
-            f"Detection Image: {dp}\n"
-            f"Segmentation Image: {sp}\n"
+            f"Detection Image: {visual_results.get('detection_image_path','N/A')}\n"
+            f"Segmentation Image: {visual_results.get('segmentation_image_path','N/A')}\n"
             f"Guidelines: {guideline_context[:200]}..."
         )
 
     def save_and_commit_image(self, image_pil: Image.Image) -> str:
-        """Save image locally and optionally commit to HF dataset."""
         os.makedirs(self.uploads_dir, exist_ok=True)
         fn = f"{datetime.now():%Y%m%d_%H%M%S}.png"
         path = os.path.join(self.uploads_dir, fn)
         image_pil.convert("RGB").save(path)
-        
         if self.hf_token and self.dataset_id:
             try:
                 HfApi().upload_file(
@@ -328,27 +314,24 @@ class AIProcessor:
                 logging.warning(f"HF upload failed: {e}")
         return path
 
-    def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: dict) -> dict:
-        """Run full analysis pipeline."""
+    def full_analysis_pipeline(self, image, questionnaire_data):
         try:
-            saved = self.save_and_commit_image(image_pil)
-            vis = self.perform_visual_analysis(image_pil)
+            saved = self.save_and_commit_image(image)
+            vis = self.perform_visual_analysis(image)
             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)
+            report = self.generate_final_report(info, vis, gc, image)
             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)}
 
-    def analyze_wound(self, image, questionnaire_data: dict) -> dict:
-        """Main analysis entry point."""
+    def analyze_wound(self, image, questionnaire_data):
         if isinstance(image, str):
             image = Image.open(image)
         return self.full_analysis_pipeline(image, questionnaire_data)
 
-    def _assess_risk_legacy(self, questionnaire_data: dict) -> dict:
-        """Legacy risk assessment function."""
+    def _assess_risk_legacy(self, questionnaire_data):
         risk_factors, risk_score = [], 0
         try:
             age = questionnaire_data.get('patient_age', 0)
@@ -377,4 +360,4 @@ class AIProcessor:
             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': []}