Update src/ai_processor.py

src/ai_processor.py

@@ -1,3 +1,7 @@
+# Disable GPU for all CPU-only model loading to avoid triggering CUDA init in the main process
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = ''
+
 import io
 import base64
 import logging
@@ -15,7 +19,6 @@ from langchain_community.vectorstores import FAISS
 from huggingface_hub import HfApi, HfFolder
 import spaces
 from .config import Config
-import os
 
 default_system_prompt = (
     "You are a world-class medical AI assistant specializing in wound care "
@@ -38,37 +41,28 @@ class AIProcessor:
         self._initialize_models()
 
     def _initialize_models(self):
-        """Initialize CPU-only AI models; MedGemma is loaded on demand within GPU context."""
-        # Set HF token
+        """Initialize all CPU-only models here; GPU models loaded later in GPU context."""
+        # HuggingFace token
         if self.config.HF_TOKEN:
             HfFolder.save_token(self.config.HF_TOKEN)
-            logging.info("HuggingFace token set successfully")
+            logging.info("✅ HuggingFace token set")
 
-        # YOLO detection on CPU (force CPU)
+        # YOLO detection (CPU only)
         try:
-            self.models_cache['det'] = YOLO(
-                self.config.YOLO_MODEL_PATH,
-                device='cpu'
-            )
-            logging.info("✅ YOLO detection model loaded on CPU")
+            self.models_cache['det'] = YOLO(self.config.YOLO_MODEL_PATH)
+            logging.info("✅ YOLO model loaded (CPU only)")
         except Exception as 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
+            logging.error(f"Failed to load YOLO model: {e}")
             raise
 
-        # Segmentation model on CPU
+        # Segmentation model (CPU)
         try:
-            self.models_cache['seg'] = load_model(
-                self.config.SEG_MODEL_PATH,
-                compile=False
-            )
-            logging.info("✅ Segmentation model loaded on CPU")
+            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 on CPU
+        # Classification model (CPU)
         try:
             self.models_cache['cls'] = pipeline(
                 'image-classification',
@@ -76,21 +70,21 @@ class AIProcessor:
                 token=self.config.HF_TOKEN,
                 device='cpu'
             )
-            logging.info("✅ Wound classification model loaded on CPU")
+            logging.info("✅ Classification pipeline loaded (CPU)")
         except Exception as e:
-            logging.warning(f"Wound classification model not available: {e}")
+            logging.warning(f"Classification pipeline not available: {e}")
 
-        # Embeddings for KB on CPU
+        # Embedding model (CPU)
         try:
             self.models_cache['embedding_model'] = HuggingFaceEmbeddings(
                 model_name='sentence-transformers/all-MiniLM-L6-v2',
                 model_kwargs={'device': 'cpu'}
             )
-            logging.info("✅ Embedding model loaded on CPU")
+            logging.info("✅ Embedding model loaded (CPU)")
         except Exception as e:
             logging.warning(f"Embedding model not available: {e}")
 
-        # Load PDF guidelines
+        # Load PDF guidelines into FAISS
         self._load_knowledge_base()
 
     def _load_knowledge_base(self):
@@ -99,118 +93,80 @@ class AIProcessor:
             if os.path.exists(pdf):
                 loader = PyPDFLoader(pdf)
                 docs.extend(loader.load())
-                logging.info(f"Loaded PDF: {pdf}")
+                logging.info(f"Loaded guideline 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)")
+            logging.info(f"✅ Knowledge base loaded with {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."""
+        """Detect & segment on CPU; return metrics and file paths."""
         if 'det' not in self.models_cache:
-            raise RuntimeError(
-                "YOLO detection model ('det') not loaded; cannot perform visual analysis"
-            )
+            raise RuntimeError("YOLO model ('det') not loaded")
 
-        try:
-            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")
+        img_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
+        res = self.models_cache['det'].predict(img_cv, device='cpu', 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]
+        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)
-            ts = datetime.now().strftime('%Y%m%d_%H%M%S')
-            det_path = (
-                f"{self.config.UPLOADS_DIR}/analysis/detection_{ts}.png"
-            )
-            cv2.imwrite(det_path, det_vis)
+        # 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)
+        ts = datetime.now().strftime('%Y%m%d_%H%M%S')
+        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 (if available)
+        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
+            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"
+            cv2.imwrite(seg_path, seg_vis)
 
-            # 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
-                )
-                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"
-                )
-                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)
 
-                cnts, _ = cv2.findContours(
-                    mask_rs,
-                    cv2.RETR_EXTERNAL,
-                    cv2.CHAIN_APPROX_SIMPLE
+        # Classification
+        wound_type = 'Unknown'
+        if 'cls' in self.models_cache:
+            try:
+                preds = self.models_cache['cls'](
+                    Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB))
                 )
-                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
-                    )
+                wound_type = max(preds, key=lambda x: x['score'])['label']
+            except Exception:
+                pass
 
-            # Classification
-            wound_type = 'Unknown'
-            if 'cls' in self.models_cache:
-                try:
-                    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_image_path': det_path,
-                'segmentation_image_path': seg_path
-            }
-
-        except Exception as e:
-            logging.error(f"Visual analysis error: {e}")
-            raise
+        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_image_path': det_path,
+            'segmentation_image_path': seg_path
+        }
 
     def query_guidelines(self, query: str):
         vs = self.knowledge_base_cache.get('vectorstore')
@@ -218,15 +174,12 @@ class AIProcessor:
             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"
-            f"{d.page_content}"
+            f"Source: {d.metadata.get('source','?')}, Page: {d.metadata.get('page','?')}\n{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."""
         if 'medgemma_pipe' not in self.models_cache:
             try:
@@ -234,16 +187,14 @@ class AIProcessor:
                     'image-text-to-text',
                     model='google/medgemma-4b-it',
                     device='auto',
-                    torch_dtype=torch.bfloat16,
+                    torch_dtype='auto',
                     offload_folder='offload',
                     token=self.config.HF_TOKEN
                 )
                 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
-                )
+                logging.warning(f"MedGemma pipeline load failed: {e}")
+                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
 
         msgs = [
             {'role':'system','content':[{'type':'text','text':default_system_prompt}]},
@@ -254,13 +205,8 @@ class AIProcessor:
         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)}
-                )
-        prompt = (
-            f"## Patient\n{patient_info}\n"
-            f"## Wound Type: {visual_results['wound_type']}"
-        )
+                msgs[1]['content'].append({'type':'image','image':Image.open(p)})
+        prompt = f"## Patient\n{patient_info}\n## Wound Type: {visual_results['wound_type']}"
         msgs[1]['content'].append({'type':'text','text':prompt})
 
         out = self.models_cache['medgemma_pipe'](
@@ -269,18 +215,14 @@ class AIProcessor:
             do_sample=False
         )
         report = out[0]['generated_text'][-1].get('content','')
-        return report or self._generate_fallback_report(
-            patient_info, visual_results, guideline_context
-        )
+        return report or self._generate_fallback_report(patient_info, visual_results, guideline_context)
 
     def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
         dp = visual_results.get('detection_image_path','N/A')
         sp = visual_results.get('segmentation_image_path','N/A')
         return (
-            f"# Report\n{patient_info}\n"
-            f"Type: {visual_results['wound_type']}\n"
-            f"Detection Image: {dp}\n"
-            f"Segmentation Image: {sp}\n"
+            f"# Report\n{patient_info}\nType: {visual_results['wound_type']}\n"
+            f"Detection Image: {dp}\nSegmentation Image: {sp}\n"
             f"Guidelines: {guideline_context[:200]}..."
         )
 
@@ -289,7 +231,7 @@ class AIProcessor:
         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 getattr(self.config,'DATASET_ID', None):
+        if self.config.HF_TOKEN and getattr(self.config, 'DATASET_ID', None):
             try:
                 api = HfApi()
                 api.upload_file(
@@ -306,17 +248,10 @@ class AIProcessor:
         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)}
@@ -327,27 +262,58 @@ class AIProcessor:
         return self.full_analysis_pipeline(image, questionnaire_data)
 
     def _assess_risk_legacy(self, questionnaire_data):
-        risk_factors, risk_score = [], 0
+        """Legacy risk assessment for backward compatibility"""
+        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 = 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}
+                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
+            }
+            
         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': []}