Update src/ai_processor.py

src/ai_processor.py

@@ -1,5 +1,5 @@
-# Disable GPU for all CPU-only model loading to avoid triggering CUDA init in the main process
 import os
+# Ensure all CPU-only models never touch CUDA
 os.environ['CUDA_VISIBLE_DEVICES'] = ''
 
 import io
@@ -20,6 +20,7 @@ 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, "
@@ -32,6 +33,59 @@ default_system_prompt = (
     "patient context."
 )
 
+@spaces.GPU(enable_queue=True, duration=120)
+def generate_medgemma_report(
+    patient_info: str,
+    visual_results: dict,
+    guideline_context: str,
+    detection_image_path: str,
+    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
+    if not hasattr(generate_medgemma_report, "_pipe"):
+        try:
+            cfg = Config()
+            generate_medgemma_report._pipe = pipeline(
+                'image-text-to-text',
+                model='google/medgemma-4b-it',
+                device='auto',
+                torch_dtype='auto',
+                offload_folder='offload',
+                token=cfg.HF_TOKEN
+            )
+            logging.info("✅ MedGemma pipeline loaded on GPU")
+        except Exception as e:
+            logging.warning(f"MedGemma pipeline load failed: {e}")
+            return None
+
+    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})
+
+    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','')
+
+
 class AIProcessor:
     def __init__(self):
         self.models_cache = {}
@@ -39,30 +93,31 @@ class AIProcessor:
         self.config = Config()
         self.px_per_cm = self.config.PIXELS_PER_CM
         self._initialize_models()
+        self._load_knowledge_base()
 
     def _initialize_models(self):
-        """Initialize all CPU-only models here; GPU models loaded later in GPU context."""
-        # HuggingFace token
+        """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)
+        # YOLO detection (CPU)
         try:
             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: {e}")
+            logging.error(f"YOLO load failed: {e}")
             raise
 
-        # Segmentation model (CPU)
+        # 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 model (CPU)
+        # Classification (CPU)
         try:
             self.models_cache['cls'] = pipeline(
                 'image-classification',
@@ -78,40 +133,38 @@ class AIProcessor:
         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}")
 
-        # Load PDF guidelines into FAISS
-        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):
                 loader = PyPDFLoader(pdf)
                 docs.extend(loader.load())
-                logging.info(f"Loaded guideline PDF: {pdf}")
+                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
-            logging.info(f"✅ Knowledge base loaded with {len(chunks)} chunks")
+            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 metrics and file paths."""
+    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")
 
         img_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
-        res = self.models_cache['det'].predict(img_cv, device='cpu', verbose=False)[0]
+        res = self.models_cache['det'].predict(img_cv, verbose=False)[0]
         if not res.boxes:
             raise ValueError("No wound detected")
 
@@ -126,19 +179,18 @@ class AIProcessor:
         det_path = f"{self.config.UPLOADS_DIR}/analysis/detection_{ts}.png"
         cv2.imwrite(det_path, det_vis)
 
-        # Segmentation (if available)
+        # Segmentation metrics
         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 = (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_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)
@@ -168,7 +220,7 @@ class AIProcessor:
             'segmentation_image_path': seg_path
         }
 
-    def query_guidelines(self, query: str):
+    def query_guidelines(self, query: str) -> str:
         vs = self.knowledge_base_cache.get('vectorstore')
         if not vs:
             return "Clinical guidelines unavailable"
@@ -178,63 +230,55 @@ class AIProcessor:
             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):
-        """Run MedGemma on GPU; return markdown report."""
-        if 'medgemma_pipe' not in self.models_cache:
-            try:
-                self.models_cache['medgemma_pipe'] = pipeline(
-                    'image-text-to-text',
-                    model='google/medgemma-4b-it',
-                    device='auto',
-                    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 load failed: {e}")
-                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
-
-        msgs = [
-            {'role':'system','content':[{'type':'text','text':default_system_prompt}]},
-            {'role':'user','content':[]}
-        ]
-        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)})
-        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'](
-            text=msgs,
-            max_new_tokens=max_new_tokens or self.config.MAX_NEW_TOKENS,
-            do_sample=False
+    def generate_final_report(
+        self,
+        patient_info: str,
+        visual_results: dict,
+        guideline_context: str,
+        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(
+            patient_info,
+            visual_results,
+            guideline_context,
+            det,
+            seg,
+            max_new_tokens
         )
-        report = out[0]['generated_text'][-1].get('content','')
-        return report or self._generate_fallback_report(patient_info, visual_results, guideline_context)
+        if report:
+            return report
+        return self._generate_fallback_report(patient_info, visual_results, guideline_context)
 
-    def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
+    def _generate_fallback_report(
+        self,
+        patient_info: str,
+        visual_results: dict,
+        guideline_context: str
+    ) -> str:
         dp = visual_results.get('detection_image_path','N/A')
         sp = visual_results.get('segmentation_image_path','N/A')
         return (
-            f"# Report\n{patient_info}\nType: {visual_results['wound_type']}\n"
-            f"Detection Image: {dp}\nSegmentation Image: {sp}\n"
+            f"# 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"Guidelines: {guideline_context[:200]}..."
         )
 
-    def save_and_commit_image(self, image_pil):
+    def save_and_commit_image(self, image_pil: Image.Image) -> str:
         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 getattr(self.config, 'DATASET_ID', None):
             try:
-                api = HfApi()
-                api.upload_file(
+                HfApi().upload_file(
                     path_or_fileobj=path,
                     path_in_repo=f"images/{fn}",
                     repo_id=self.config.DATASET_ID,
@@ -244,19 +288,28 @@ class AIProcessor:
                 logging.warning(f"HF upload failed: {e}")
         return path
 
-    def full_analysis_pipeline(self, image_pil, questionnaire_data):
+    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)
+            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)}
 
-    def analyze_wound(self, image, questionnaire_data):
+    def analyze_wound(self, image, questionnaire_data: dict) -> dict:
         if isinstance(image, str):
             image = Image.open(image)
         return self.full_analysis_pipeline(image, questionnaire_data)