Update src/ai_processor.py

src/ai_processor.py

@@ -1,62 +1,171 @@
 import os
-import io
-import base64
 import logging
-import numpy as np
 import cv2
+import numpy as np
 from PIL import Image
 from datetime import datetime
+import gradio as gr
+import spaces
+
+from huggingface_hub import HfApi, HfFolder
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-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, "
-    "evidence-based medical assessments focusing on: (1) Precise wound "
-    "classification based on tissue type and appearance, (2) Specific "
-    "treatment recommendations with exact product names or interventions when "
-    "appropriate, (3) Objective evaluation of healing progression or deterioration "
-    "indicators, and (4) Clear follow-up timelines. Avoid general statements and "
-    "prioritize actionable insights based on the visual analysis measurements and "
-    "patient context."
-)
-
-# No torch or transformers-related imports at top-level!
 
+# =============== LOGGING SETUP ===============
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# =============== CONFIGURATION ===============
+UPLOADS_DIR = "uploads"
+if not os.path.exists(UPLOADS_DIR):
+    os.makedirs(UPLOADS_DIR)
+    logging.info(f"Created uploads directory: {UPLOADS_DIR}")
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+YOLO_MODEL_PATH = "best.pt"
+SEG_MODEL_PATH = "segmentation_model.h5"
+GUIDELINE_PDFS = ["eHealth in Wound Care.pdf", "IWGDF Guideline.pdf", "evaluation.pdf"]
+DATASET_ID = "SmartHeal/wound-image-uploads"
+MAX_NEW_TOKENS = 2048
+PIXELS_PER_CM = 38
+
+# =============== GLOBAL CACHES ===============
+models_cache = {}
+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",
+        model="Hemg/Wound-classification",
+        token=hf_token,
+        device="cpu"
+    )
+
+def load_embedding_model():
+    """Load embedding model for knowledge base."""
+    return HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        model_kwargs={"device": "cpu"}
+    )
+
+# =============== 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()
+            logging.info("✅ Embedding model loaded (CPU)")
+        except Exception as e:
+            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):
+            try:
+                loader = PyPDFLoader(pdf_path)
+                docs.extend(loader.load())
+                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)
+        knowledge_base_cache["vector_store"] = FAISS.from_documents(chunks, models_cache["embedding_model"])
+        logging.info(f"✅ Knowledge base ready with {len(chunks)} chunks")
+    else:
+        knowledge_base_cache["vector_store"] = None
+        logging.warning("Knowledge base unavailable")
+
+# Initialize models on app startup
+initialize_cpu_models()
+setup_knowledge_base()
+
+# =============== GPU-DECORATED MEDGEMMA FUNCTION ===============
 @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:
-    # --- All GPU-related imports and model loading here! ---
+    patient_info,
+    visual_results,
+    guideline_context,
+    detection_image_path,
+    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
 
-    # System prompt as before
-    global default_system_prompt
+    default_system_prompt = (
+        "You are a world-class medical AI assistant specializing in wound care "
+        "with expertise in wound assessment and treatment. Provide concise, "
+        "evidence-based medical assessments focusing on: (1) Precise wound "
+        "classification based on tissue type and appearance, (2) Specific "
+        "treatment recommendations with exact product names or interventions when "
+        "appropriate, (3) Objective evaluation of healing progression or deterioration "
+        "indicators, and (4) Clear follow-up timelines. Avoid general statements and "
+        "prioritize actionable insights based on the visual analysis measurements and "
+        "patient context."
+    )
 
     # Lazy-load MedGemma pipeline on GPU
     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='cuda',  # Explicitly on GPU
-                torch_dtype='auto',
-                offload_folder='offload',
-                token=cfg.HF_TOKEN
+                "image-text-to-text",
+                model="google/medgemma-4b-it",
+                device="cuda",
+                torch_dtype=torch.bfloat16,
+                offload_folder="offload",
+                token=HF_TOKEN,
             )
             logging.info("✅ MedGemma pipeline loaded on GPU")
         except Exception as e:
@@ -67,202 +176,132 @@ def generate_medgemma_report(
 
     # Compose 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 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)})
+            msgs[1]["content"].append({"type": "image", "image": Image.open(path)})
 
-    # Attach text
+    # 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})
-
-    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', '')
+    msgs[1]["content"].append({"type": "text", "text": prompt})
 
+    try:
+        out = pipe(text=msgs, max_new_tokens=max_new_tokens or MAX_NEW_TOKENS, do_sample=False)
+        return out[0]["generated_text"][-1].get("content", "")
+    except Exception as e:
+        logging.error(f"Failed to generate MedGemma report: {e}")
+        return f"❌ An error occurred: {e}"
 
+# =============== AI PROCESSOR CLASS ===============
 class AIProcessor:
     def __init__(self):
-        self.models_cache = {}
-        self.knowledge_base_cache = {}
-        self.config = Config()
-        self.px_per_cm = self.config.PIXELS_PER_CM
-        self._initialize_models()
-        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:
-            from ultralytics import YOLO
-            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:
-            from tensorflow.keras.models import load_model
-            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:
-            from transformers import pipeline
-            self.models_cache['cls'] = pipeline(
-                'image-classification',
-                model='Hemg/Wound-classification',
-                token=self.config.HF_TOKEN,
-                device='cpu'
-            )
-            logging.info("✅ Classification pipeline loaded (CPU)")
-        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'}
-            )
-            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):
-                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']
-            )
-            logging.info(f"✅ Knowledge base loaded ({len(chunks)} chunks)")
-        else:
-            self.knowledge_base_cache['vectorstore'] = None
-            logging.warning("Knowledge base unavailable")
+        self.models_cache = models_cache
+        self.knowledge_base_cache = knowledge_base_cache
+        self.px_per_cm = PIXELS_PER_CM
+        self.uploads_dir = UPLOADS_DIR
+        self.dataset_id = DATASET_ID
+        self.hf_token = HF_TOKEN
 
     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:
+        img_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
+        yolo = self.models_cache.get("det")
+        if yolo is None:
             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, verbose=False)[0]
+        res = yolo.predict(img_cv, verbose=False, device="cpu")[0]
         if not res.boxes:
             raise ValueError("No wound detected")
 
-        x1, y1, x2, y2 = res.boxes.xyxy[0].cpu().numpy().astype(int)
+        x1, y1, x2, y2 = res.boxes.xyxy.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.rectangle(det_vis, (x1, y1), (x2, 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
         seg_path = None
-        if 'seg' in self.models_cache:
-            h, w = self.models_cache['seg'].input_shape[1:3]
+        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 = (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]
+            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.config.UPLOADS_DIR}/analysis/segmentation_{ts}.png"
+            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)
+                area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
 
         # Classification
-        wound_type = 'Unknown'
-        if 'cls' in self.models_cache:
+        wound_type = "Unknown"
+        cls_pipe = self.models_cache.get("cls")
+        if cls_pipe:
             try:
-                preds = self.models_cache['cls'](Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB)))
-                wound_type = max(preds, key=lambda x: x['score'])['label']
+                preds = cls_pipe(Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB)))
+                wound_type = max(preds, 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
+            "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) -> str:
-        vs = self.knowledge_base_cache.get('vectorstore')
+        """Query the knowledge base for relevant information."""
+        vs = self.knowledge_base_cache.get("vector_store")
         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}"
-            for d in docs
+        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}" 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
+        self, patient_info: str, visual_results: dict, guideline_context: str, image_pil: Image.Image, max_new_tokens: int = None
     ) -> str:
-        det = visual_results.get('detection_image_path', '')
-        seg = visual_results.get('segmentation_image_path', '')
-        # This GPU call is safe: it triggers all CUDA/model code *inside* the decorator context.
-        report = generate_medgemma_report(
-            patient_info, visual_results, guideline_context,
-            det, seg, max_new_tokens
-        )
+        """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)
         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
+        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')
         return (
-            f"# Report\n{patient_info}\n"
+            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"
@@ -270,40 +309,46 @@ class AIProcessor:
         )
 
     def save_and_commit_image(self, image_pil: Image.Image) -> str:
-        os.makedirs(self.config.UPLOADS_DIR, exist_ok=True)
+        """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.config.UPLOADS_DIR, fn)
-        image_pil.convert('RGB').save(path)
-        if self.config.HF_TOKEN and getattr(self.config, 'DATASET_ID', None):
+        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(
                     path_or_fileobj=path,
                     path_in_repo=f"images/{fn}",
-                    repo_id=self.config.DATASET_ID,
-                    repo_type='dataset'
+                    repo_id=self.dataset_id,
+                    repo_type="dataset",
                 )
+                logging.info("✅ Image committed to HF dataset")
             except Exception as e:
                 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."""
         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)
+            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)}
 
     def analyze_wound(self, image, questionnaire_data: dict) -> dict:
+        """Main analysis entry point."""
         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."""
         risk_factors, risk_score = [], 0
         try:
             age = questionnaire_data.get('patient_age', 0)
@@ -311,12 +356,15 @@ 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()
             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)
             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
@@ -324,8 +372,9 @@ 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")
             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': []}