Update src/ai_processor.py

src/ai_processor.py

@@ -1,611 +1,611 @@
-import os
-import logging
-import cv2
-import numpy as np
-from PIL import Image
-import torch
-import json
-from datetime import datetime
-import tensorflow as tf
-from transformers import pipeline
-from ultralytics import YOLO
-from tensorflow.keras.models import load_model
-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
-
-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()
-
-    def _initialize_models(self):
-        """Initialize all AI models including real-time models"""
-        try:
-            # Set HuggingFace token
-            if self.config.HF_TOKEN:
-                HfFolder.save_token(self.config.HF_TOKEN)
-                logging.info("HuggingFace token set successfully")
-
-            # Initialize MedGemma pipeline for medical text generation
-            try:
-                self.models_cache["medgemma_pipe"] = pipeline(
-                    "image-text-to-text",
-                    model="google/medgemma-4b-it",
-                    torch_dtype=torch.bfloat16,
-                    offload_folder="offload",
-                    device_map="auto",
-                    token=self.config.HF_TOKEN
-                )
-                logging.info("✅ MedGemma pipeline loaded successfully")
-            except Exception as e:
-                logging.warning(f"MedGemma pipeline not available: {e}")
-
-            # Initialize YOLO model for wound detection
-            try:
-                self.models_cache["det"] = YOLO(self.config.YOLO_MODEL_PATH)
-                logging.info("✅ YOLO detection model loaded successfully")
-            except Exception as e:
-                logging.warning(f"YOLO model not available: {e}")
-
-            # Initialize segmentation model
-            try:
-                self.models_cache["seg"] = load_model(self.config.SEG_MODEL_PATH, compile=False)
-                logging.info("✅ Segmentation model loaded successfully")
-            except Exception as e:
-                logging.warning(f"Segmentation model not available: {e}")
-
-            # Initialize wound classification model
-            try:
-                self.models_cache["cls"] = pipeline(
-                    "image-classification", 
-                    model="Hemg/Wound-classification", 
-                    token=self.config.HF_TOKEN, 
-                    device="cpu"
-                )
-                logging.info("✅ Wound classification model loaded successfully")
-            except Exception as e:
-                logging.warning(f"Wound classification model not available: {e}")
-
-            # Initialize embedding model for knowledge base
-            try:
-                self.models_cache["embedding_model"] = HuggingFaceEmbeddings(
-                    model_name="sentence-transformers/all-MiniLM-L6-v2", 
-                    model_kwargs={'device': 'cpu'}
-                )
-                logging.info("✅ Embedding model loaded successfully")
-            except Exception as e:
-                logging.warning(f"Embedding model not available: {e}")
-
-            logging.info("✅ All models loaded.")
-            self._load_knowledge_base()
-
-        except Exception as e:
-            logging.error(f"Error initializing AI models: {e}")
-
-    def _load_knowledge_base(self):
-        """Load knowledge base from PDF guidelines"""
-        try:
-            documents = []
-            for pdf_path in self.config.GUIDELINE_PDFS:
-                if os.path.exists(pdf_path):
-                    loader = PyPDFLoader(pdf_path)
-                    docs = loader.load()
-                    documents.extend(docs)
-                    logging.info(f"Loaded PDF: {pdf_path}")
-
-            if documents and 'embedding_model' in self.models_cache:
-                # Split documents into chunks
-                text_splitter = RecursiveCharacterTextSplitter(
-                    chunk_size=1000,
-                    chunk_overlap=100
-                )
-                chunks = text_splitter.split_documents(documents)
-                
-                # Create vector store
-                vectorstore = FAISS.from_documents(chunks, self.models_cache['embedding_model'])
-                self.knowledge_base_cache['vectorstore'] = vectorstore
-                logging.info(f"✅ Knowledge base loaded with {len(chunks)} chunks")
-            else:
-                self.knowledge_base_cache['vectorstore'] = None
-                logging.warning("Knowledge base not available - no PDFs found or embedding model unavailable")
-
-        except Exception as e:
-            logging.warning(f"Knowledge base loading error: {e}")
-            self.knowledge_base_cache['vectorstore'] = None
-
-
-    def perform_visual_analysis(self, image_pil):
-        """Perform comprehensive visual analysis of wound image."""
-        try:
-            # Convert PIL to OpenCV format
-            image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
-    
-            # YOLO detection
-            if 'det' not in self.models_cache:
-                raise ValueError("YOLO detection model not available.")
-    
-            results = self.models_cache['det'].predict(image_cv, verbose=False, device="cpu")
-    
-            if not results or not results[0].boxes:
-                raise ValueError("No wound detected in the image.")
-    
-            # Extract bounding box
-            box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
-            x1, y1, x2, y2 = box
-            region_cv = image_cv[y1:y2, x1:x2]
-    
-            # Save detection image
-            detection_image_cv = image_cv.copy()
-            cv2.rectangle(detection_image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            os.makedirs(os.path.join(self.config.UPLOADS_DIR, "analysis"), exist_ok=True)
-            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-            detection_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"detection_{timestamp}.png")
-            cv2.imwrite(detection_image_path, detection_image_cv)
-            detection_image_pil = Image.fromarray(cv2.cvtColor(detection_image_cv, cv2.COLOR_BGR2RGB))
-    
-            # Initialize outputs
-            length = breadth = area = 0
-            segmentation_image_pil = None
-            segmentation_image_path = None
-    
-            # Segmentation (optional)
-            if 'seg' in self.models_cache:
-                input_size = self.models_cache['seg'].input_shape[1:3]  # (height, width)
-                resized_region = cv2.resize(region_cv, (input_size[1], input_size[0]))
-    
-                seg_input = np.expand_dims(resized_region / 255.0, 0)
-                mask_pred = self.models_cache['seg'].predict(seg_input, verbose=0)[0]
-                mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
-    
-                # Resize mask back to original region size
-                mask_resized = cv2.resize(mask_np, (region_cv.shape[1], region_cv.shape[0]), interpolation=cv2.INTER_NEAREST)
-    
-                # Overlay mask on region for visualization
-                overlay = region_cv.copy()
-                overlay[mask_resized == 1] = [0, 0, 255]  # Red overlay
-    
-                # Blend overlay for final output
-                segmented_visual = cv2.addWeighted(region_cv, 0.7, overlay, 0.3, 0)
-    
-                # Save segmentation image
-                segmentation_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"segmentation_{timestamp}.png")
-                cv2.imwrite(segmentation_image_path, segmented_visual)
-                segmentation_image_pil = Image.fromarray(cv2.cvtColor(segmented_visual, cv2.COLOR_BGR2RGB))
-    
-                # Wound measurements from resized mask
-                contours, _ = cv2.findContours(mask_resized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-                if contours:
-                    cnt = max(contours, key=cv2.contourArea)
-                    x, y, w, h = cv2.boundingRect(cnt)
-                    length = round(h / self.px_per_cm, 2)
-                    breadth = round(w / self.px_per_cm, 2)
-                    area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
-    
-            # Classification (optional)
-            wound_type = "Unknown"
-            if 'cls' in self.models_cache:
-                try:
-                    region_pil = Image.fromarray(cv2.cvtColor(region_cv, cv2.COLOR_BGR2RGB))
-                    cls_result = self.models_cache['cls'](region_pil)
-                    wound_type = max(cls_result, key=lambda x: x['score'])['label']
-                except Exception as e:
-                    logging.warning(f"Wound classification error: {e}")
-    
-            return {
-                'wound_type': wound_type,
-                'length_cm': length,
-                'breadth_cm': breadth,
-                'surface_area_cm2': area,
-                'detection_confidence': float(results[0].boxes[0].conf.cpu().item()),
-                'bounding_box': box.tolist(),
-                'detection_image_path': detection_image_path,
-                'detection_image_pil': detection_image_pil,
-                'segmentation_image_path': segmentation_image_path,
-                'segmentation_image_pil': segmentation_image_pil
-            }
-    
-        except Exception as e:
-            logging.error(f"Visual analysis error: {e}")
-            raise ValueError(f"Visual analysis failed: {str(e)}")
-
-
-    def query_guidelines(self, query: str):
-        """Query the knowledge base for relevant guidelines"""
-        try:
-            vector_store = self.knowledge_base_cache.get("vectorstore")
-            if not vector_store:
-                return "Knowledge base unavailable - clinical guidelines not loaded"
-
-            # Retrieve relevant documents
-            retriever = vector_store.as_retriever(search_kwargs={"k": 10})
-            docs = retriever.invoke(query)
-            
-            if not docs:
-                return "No relevant guidelines found for the query"
-
-            # Format the results
-            formatted_results = []
-            for doc in docs:
-                source = doc.metadata.get('source', 'Unknown')
-                page = doc.metadata.get('page', 'N/A')
-                content = doc.page_content.strip()
-                formatted_results.append(f"Source: {source}, Page: {page}\nContent: {content}")
-
-            return "\n\n".join(formatted_results)
-
-        except Exception as e:
-            logging.error(f"Guidelines query error: {e}")
-            return f"Error querying guidelines: {str(e)}"
-
-    
-    def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
-        """Generate comprehensive medical report using MedGemma"""
-        try:
-            if 'medgemma_pipe' not in self.models_cache:
-                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
-
-            max_tokens = max_new_tokens or self.config.MAX_NEW_TOKENS
-            
-            # Get detection and segmentation images if available
-            detection_image = visual_results.get('detection_image_pil', None)
-            segmentation_image = visual_results.get('segmentation_image_pil', None)
-            
-            # Create image paths for report
-            detection_path = visual_results.get('detection_image_path', '')
-            segmentation_path = visual_results.get('segmentation_image_path', '')
-            
-            # Create detailed prompt for medical analysis with image paths
-            prompt = f"""
-# Wound Care Report
-
-## Patient Information
-{patient_info}
-
-## Visual Analysis Summary
-- Wound Type: {visual_results.get('wound_type', 'Unknown')}
-- Length: {visual_results.get('length_cm', 0)} cm
-- Breadth: {visual_results.get('breadth_cm', 0)} cm
-- Surface Area: {visual_results.get('surface_area_cm2', 0)} cm²
-- Detection Confidence: {visual_results.get('detection_confidence', 0):.2f}
-
-## Clinical Reference
-{guideline_context}
-
-You are SmartHeal-AI Agent, a world-class wound care AI specialist trained in clinical wound assessment and guideline-based treatment planning.
-Your task is to process the following structured inputs (patient data, wound measurements, clinical guidelines, and image) and perform **clinical reasoning and decision-making** to generate a complete wound care report.
----
-🔍 **YOUR PROCESS — FOLLOW STRICTLY:**
-### Step 1: Clinical Reasoning (Chain-of-Thought)
-Use the provided information to think step-by-step about:
-- Patient’s risk factors (e.g. diabetes, age, healing limitations)
-- Wound characteristics (size, tissue appearance, moisture, infection signs)
-- Visual clues from the image (location, granulation, maceration, inflammation, surrounding skin)
-
----
--Step 2: Structured Clinical Report
-Generate the following report sections using markdown and medical terminology:
-**1. Clinical Summary**
-- Describe wound appearance and tissue types (e.g., slough, necrotic, granulating, epithelializing)
-- Include size, wound bed condition, peri-wound skin, and signs of infection or biofilm
-- Mention inferred location (e.g., heel, forefoot) if image allows
-- Summarize patient's systemic risk profile
-**2. Medicinal & Dressing Recommendations**
-Based on your analysis:
-- Recommend specific **wound care dressings** (e.g., hydrocolloid, alginate, foam, antimicrobial silver, etc.) suitable to wound moisture level and infection risk
-- Propose **topical or systemic agents** ONLY if relevant — include name classes (e.g., antiseptic: povidone iodine, antibiotic ointments, enzymatic debriders)
-- Mention **techniques** (e.g., sharp debridement, NPWT, moisture balance, pressure offloading, dressing frequency)
-- Avoid repeating guidelines — **apply them**
-**3. Key Risk Factors**
-Explain how the patient’s condition (e.g., diabetic, poor circulation, advanced age, poor hygiene) may affect wound healing
-**4. Prognosis & Monitoring Advice**
-- Mention how often wound should be reassessed
-- Indicate signs to monitor for deterioration or improvement
-- Include when escalation to specialist is necessary
-
-**Note:** Every dressing change is a chance for wound reassessment. Always perform a thorough wound evaluation at each dressing change.
-"""
-
-            # Prepare messages for MedGemma with all available images
-            content_list = [{"type": "text", "text": prompt}]
-            
-            # Add original image
-            if image_pil:
-                content_list.insert(0, {"type": "image", "image": image_pil})
-                
-            # Add detection image if available
-            if detection_image:
-                content_list.insert(1, {"type": "image", "image": detection_image})
-                
-            # Add segmentation image if available
-            if segmentation_image:
-                content_list.insert(2, {"type": "image", "image": segmentation_image})
-            
-            messages = [
-                {
-                    "role": "system",
-                    "content": [{"type": "text", "text": "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."}],
-                },
-                {
-                    "role": "user",
-                    "content": content_list
-                }
-            ]
-
-            # Generate report using MedGemma
-            output = self.models_cache['medgemma_pipe'](
-                text=messages,
-                max_new_tokens=1024,
-                do_sample=False,
-            )
-            
-            generated_content = output[0]['generated_text'][-1].get('content', '').strip()
-            
-            # Include image paths in the final report for display in UI
-            if generated_content:
-                # Add image paths to the report for frontend display
-                image_paths_section = f"""
-## Analysis Images
-- Original Image: {image_pil}
-- Detection Image: {detection_path}
-- Segmentation Image: {segmentation_path}
-"""
-                generated_content = image_paths_section + generated_content
-                
-            return generated_content if generated_content else self._generate_fallback_report(patient_info, visual_results, guideline_context)
-
-        except Exception as e:
-            logging.error(f"MedGemma report generation error: {e}")
-            return self._generate_fallback_report(patient_info, visual_results, guideline_context)
-
-    def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
-        """Generate a fallback report when MedGemma is not available"""
-        # Get image paths for report
-        detection_path = visual_results.get('detection_image_path', 'Not available')
-        segmentation_path = visual_results.get('segmentation_image_path', 'Not available')
-        
-        report = f"""
-# Wound Analysis Report
-## Patient Information
-{patient_info}
-
-## Visual Analysis Results
-- **Wound Type**: {visual_results.get('wound_type', 'Unknown')}
-- **Dimensions**: {visual_results.get('length_cm', 0)} cm × {visual_results.get('breadth_cm', 0)} cm
-- **Surface Area**: {visual_results.get('surface_area_cm2', 0)} cm²
-- **Detection Confidence**: {visual_results.get('detection_confidence', 0):.2f}
-
-## Analysis Images
-- **Detection Image**: {detection_path}
-- **Segmentation Image**: {segmentation_path}
-
-## Assessment
-Based on the visual analysis, this appears to be a {visual_results.get('wound_type', 'wound')} with measurable dimensions. 
-
-## Recommendations
-- Continue monitoring wound healing progress
-- Maintain proper wound hygiene
-- Follow appropriate dressing protocols
-- Seek medical attention if signs of infection develop
-
-## Clinical Guidelines
-{guideline_context[:500]}...
-
-*Note: This is an automated analysis. Please consult with a healthcare professional for definitive diagnosis and treatment.*
-"""
-        return report
-
-    def save_and_commit_image(self, image_pil):
-        """Save image locally and optionally upload to HuggingFace dataset"""
-        try:
-            # Ensure uploads directory exists
-            os.makedirs(self.config.UPLOADS_DIR, exist_ok=True)
-            
-            # Generate filename with timestamp
-            filename = f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.png"
-            local_path = os.path.join(self.config.UPLOADS_DIR, filename)
-            
-            # Save image locally
-            image_pil.convert("RGB").save(local_path)
-            logging.info(f"Image saved locally: {local_path}")
-
-            # Upload to HuggingFace dataset if configured
-            if self.config.HF_TOKEN and self.config.DATASET_ID:
-                try:
-                    api = HfApi()
-                    api.upload_file(
-                        path_or_fileobj=local_path,
-                        path_in_repo=f"images/{filename}",
-                        repo_id=self.config.DATASET_ID,
-                        repo_type="dataset",
-                        commit_message=f"Upload wound image: {filename}"
-                    )
-                    logging.info("✅ Image uploaded to HuggingFace dataset")
-                except Exception as e:
-                    logging.warning(f"HuggingFace upload failed: {e}")
-
-            return local_path
-
-        except Exception as e:
-            logging.error(f"Image saving error: {e}")
-            return None
-
-    @spaces.GPU(enable_queue=True, duration=120) if SPACES_AVAILABLE else lambda x: x
-    def full_analysis_pipeline(self, image, questionnaire_data):
-        """Complete analysis pipeline with real-time models"""
-        try:
-            # Save the image
-            saved_path = self.save_and_commit_image(image)
-            
-            # Perform visual analysis
-            visual_results = self.perform_visual_analysis(image)
-            
-            # Format patient information
-            patient_info = ", ".join([f"{k}: {v}" for k, v in questionnaire_data.items() if v])
-            
-            # Create query for guidelines
-            wound_type = visual_results.get('wound_type', 'wound')
-            moisture = questionnaire_data.get('moisture', 'unknown')
-            infection = questionnaire_data.get('infection', 'unknown')
-            diabetic = questionnaire_data.get('diabetic', 'unknown')
-            
-            query = f"best practices for managing a {wound_type} with moisture level '{moisture}' and signs of infection '{infection}' in a patient who is diabetic '{diabetic}'"
-            
-            # Query guidelines
-            guideline_context = self.query_guidelines(query)
-            
-            # Generate final report
-            final_report = self.generate_final_report(patient_info, visual_results, guideline_context, image)
-            
-            return {
-                'success': True,
-                'visual_analysis': visual_results,
-                'report': final_report,
-                'saved_image_path': saved_path,
-                'timestamp': datetime.now().isoformat()
-            }
-
-        except Exception as e:
-            logging.error(f"Full analysis pipeline error: {e}")
-            return {
-                'success': False,
-                'error': str(e),
-                'timestamp': datetime.now().isoformat()
-            }
-
-    # Legacy methods for backward compatibility
-    def analyze_wound(self, image, questionnaire_data):
-        """Legacy method for backward compatibility"""
-        try:
-            # Convert string path to PIL Image if needed
-            if isinstance(image, str):
-                try:
-                    from PIL import Image
-                    image = Image.open(image)
-                    logging.info(f"Converted string path to PIL Image: {image}")
-                except Exception as e:
-                    logging.error(f"Error converting string path to image: {e}")
-            
-            # Ensure we have a PIL Image object
-            if not isinstance(image, Image.Image):
-                try:
-                    from PIL import Image
-                    import io
-                    
-                    # If it's a file-like object
-                    if hasattr(image, 'read'):
-                        # Reset file pointer if possible
-                        if hasattr(image, 'seek'):
-                            image.seek(0)
-                        image = Image.open(image)
-                        logging.info("Converted file-like object to PIL Image")
-                except Exception as e:
-                    logging.error(f"Error ensuring image is PIL Image: {e}")
-                    raise ValueError(f"Invalid image format: {type(image)}")
-            
-            result = self.full_analysis_pipeline(image, questionnaire_data)
-            
-            if result['success']:
-                return {
-                    'timestamp': result['timestamp'],
-                    'summary': f"Analysis completed for {questionnaire_data.get('patient_name', 'patient')}",
-                    'recommendations': result['report'],
-                    'wound_detection': {
-                        'status': 'success',
-                        'detections': [result['visual_analysis']],
-                        'total_wounds': 1
-                    },
-                    'segmentation_result': {
-                        'status': 'success',
-                        'wound_area_percentage': result['visual_analysis'].get('surface_area_cm2', 0)
-                    },
-                    'risk_assessment': self._assess_risk_legacy(questionnaire_data),
-                    'guideline_recommendations': [result['report'][:200] + "..."]
-                }
-            else:
-                return {
-                    'timestamp': result['timestamp'],
-                    'summary': f"Analysis failed: {result['error']}",
-                    'recommendations': "Please consult with a healthcare professional.",
-                    'wound_detection': {'status': 'error', 'message': result['error']},
-                    'segmentation_result': {'status': 'error', 'message': result['error']},
-                    'risk_assessment': {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []},
-                    'guideline_recommendations': ["Analysis unavailable due to error"]
-                }
-                
-        except Exception as e:
-            logging.error(f"Legacy analyze_wound error: {e}")
-            return {
-                'timestamp': datetime.now().isoformat(),
-                'summary': f"Analysis error: {str(e)}",
-                'recommendations': "Please consult with a healthcare professional.",
-                'wound_detection': {'status': 'error', 'message': str(e)},
-                'segmentation_result': {'status': 'error', 'message': str(e)},
-                'risk_assessment': {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []},
-                'guideline_recommendations': ["Analysis unavailable due to error"]
-            }
-
-    def _assess_risk_legacy(self, questionnaire_data):
-        """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
-            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
-            }
-            
-        except Exception as e:
-            logging.error(f"Risk assessment error: {e}")
+import os
+import logging
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+import json
+from datetime import datetime
+import tensorflow as tf
+from transformers import pipeline
+from ultralytics import YOLO
+from tensorflow.keras.models import load_model
+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
+
+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()
+
+    def _initialize_models(self):
+        """Initialize all AI models including real-time models"""
+        try:
+            # Set HuggingFace token
+            if self.config.HF_TOKEN:
+                HfFolder.save_token(self.config.HF_TOKEN)
+                logging.info("HuggingFace token set successfully")
+
+            # Initialize MedGemma pipeline for medical text generation
+            try:
+                self.models_cache["medgemma_pipe"] = pipeline(
+                    "image-text-to-text",
+                    model="google/medgemma-4b-it",
+                    torch_dtype=torch.bfloat16,
+                    offload_folder="offload",
+                    device_map="auto",
+                    token=self.config.HF_TOKEN
+                )
+                logging.info("✅ MedGemma pipeline loaded successfully")
+            except Exception as e:
+                logging.warning(f"MedGemma pipeline not available: {e}")
+
+            # Initialize YOLO model for wound detection
+            try:
+                self.models_cache["det"] = YOLO(self.config.YOLO_MODEL_PATH)
+                logging.info("✅ YOLO detection model loaded successfully")
+            except Exception as e:
+                logging.warning(f"YOLO model not available: {e}")
+
+            # Initialize segmentation model
+            try:
+                self.models_cache["seg"] = load_model(self.config.SEG_MODEL_PATH, compile=False)
+                logging.info("✅ Segmentation model loaded successfully")
+            except Exception as e:
+                logging.warning(f"Segmentation model not available: {e}")
+
+            # Initialize wound classification model
+            try:
+                self.models_cache["cls"] = pipeline(
+                    "image-classification", 
+                    model="Hemg/Wound-classification", 
+                    token=self.config.HF_TOKEN, 
+                    device="cpu"
+                )
+                logging.info("✅ Wound classification model loaded successfully")
+            except Exception as e:
+                logging.warning(f"Wound classification model not available: {e}")
+
+            # Initialize embedding model for knowledge base
+            try:
+                self.models_cache["embedding_model"] = HuggingFaceEmbeddings(
+                    model_name="sentence-transformers/all-MiniLM-L6-v2", 
+                    model_kwargs={'device': 'cpu'}
+                )
+                logging.info("✅ Embedding model loaded successfully")
+            except Exception as e:
+                logging.warning(f"Embedding model not available: {e}")
+
+            logging.info("✅ All models loaded.")
+            self._load_knowledge_base()
+
+        except Exception as e:
+            logging.error(f"Error initializing AI models: {e}")
+
+    def _load_knowledge_base(self):
+        """Load knowledge base from PDF guidelines"""
+        try:
+            documents = []
+            for pdf_path in self.config.GUIDELINE_PDFS:
+                if os.path.exists(pdf_path):
+                    loader = PyPDFLoader(pdf_path)
+                    docs = loader.load()
+                    documents.extend(docs)
+                    logging.info(f"Loaded PDF: {pdf_path}")
+
+            if documents and 'embedding_model' in self.models_cache:
+                # Split documents into chunks
+                text_splitter = RecursiveCharacterTextSplitter(
+                    chunk_size=1000,
+                    chunk_overlap=100
+                )
+                chunks = text_splitter.split_documents(documents)
+                
+                # Create vector store
+                vectorstore = FAISS.from_documents(chunks, self.models_cache['embedding_model'])
+                self.knowledge_base_cache['vectorstore'] = vectorstore
+                logging.info(f"✅ Knowledge base loaded with {len(chunks)} chunks")
+            else:
+                self.knowledge_base_cache['vectorstore'] = None
+                logging.warning("Knowledge base not available - no PDFs found or embedding model unavailable")
+
+        except Exception as e:
+            logging.warning(f"Knowledge base loading error: {e}")
+            self.knowledge_base_cache['vectorstore'] = None
+
+
+    def perform_visual_analysis(self, image_pil):
+        """Perform comprehensive visual analysis of wound image."""
+        try:
+            # Convert PIL to OpenCV format
+            image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
+    
+            # YOLO detection
+            if 'det' not in self.models_cache:
+                raise ValueError("YOLO detection model not available.")
+    
+            results = self.models_cache['det'].predict(image_cv, verbose=False, device="cpu")
+    
+            if not results or not results[0].boxes:
+                raise ValueError("No wound detected in the image.")
+    
+            # Extract bounding box
+            box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
+            x1, y1, x2, y2 = box
+            region_cv = image_cv[y1:y2, x1:x2]
+    
+            # Save detection image
+            detection_image_cv = image_cv.copy()
+            cv2.rectangle(detection_image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            os.makedirs(os.path.join(self.config.UPLOADS_DIR, "analysis"), exist_ok=True)
+            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
+            detection_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"detection_{timestamp}.png")
+            cv2.imwrite(detection_image_path, detection_image_cv)
+            detection_image_pil = Image.fromarray(cv2.cvtColor(detection_image_cv, cv2.COLOR_BGR2RGB))
+    
+            # Initialize outputs
+            length = breadth = area = 0
+            segmentation_image_pil = None
+            segmentation_image_path = None
+    
+            # Segmentation (optional)
+            if 'seg' in self.models_cache:
+                input_size = self.models_cache['seg'].input_shape[1:3]  # (height, width)
+                resized_region = cv2.resize(region_cv, (input_size[1], input_size[0]))
+    
+                seg_input = np.expand_dims(resized_region / 255.0, 0)
+                mask_pred = self.models_cache['seg'].predict(seg_input, verbose=0)[0]
+                mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
+    
+                # Resize mask back to original region size
+                mask_resized = cv2.resize(mask_np, (region_cv.shape[1], region_cv.shape[0]), interpolation=cv2.INTER_NEAREST)
+    
+                # Overlay mask on region for visualization
+                overlay = region_cv.copy()
+                overlay[mask_resized == 1] = [0, 0, 255]  # Red overlay
+    
+                # Blend overlay for final output
+                segmented_visual = cv2.addWeighted(region_cv, 0.7, overlay, 0.3, 0)
+    
+                # Save segmentation image
+                segmentation_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"segmentation_{timestamp}.png")
+                cv2.imwrite(segmentation_image_path, segmented_visual)
+                segmentation_image_pil = Image.fromarray(cv2.cvtColor(segmented_visual, cv2.COLOR_BGR2RGB))
+    
+                # Wound measurements from resized mask
+                contours, _ = cv2.findContours(mask_resized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                if contours:
+                    cnt = max(contours, key=cv2.contourArea)
+                    x, y, w, h = cv2.boundingRect(cnt)
+                    length = round(h / self.px_per_cm, 2)
+                    breadth = round(w / self.px_per_cm, 2)
+                    area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
+    
+            # Classification (optional)
+            wound_type = "Unknown"
+            if 'cls' in self.models_cache:
+                try:
+                    region_pil = Image.fromarray(cv2.cvtColor(region_cv, cv2.COLOR_BGR2RGB))
+                    cls_result = self.models_cache['cls'](region_pil)
+                    wound_type = max(cls_result, key=lambda x: x['score'])['label']
+                except Exception as e:
+                    logging.warning(f"Wound classification error: {e}")
+    
+            return {
+                'wound_type': wound_type,
+                'length_cm': length,
+                'breadth_cm': breadth,
+                'surface_area_cm2': area,
+                'detection_confidence': float(results[0].boxes[0].conf.cpu().item()),
+                'bounding_box': box.tolist(),
+                'detection_image_path': detection_image_path,
+                'detection_image_pil': detection_image_pil,
+                'segmentation_image_path': segmentation_image_path,
+                'segmentation_image_pil': segmentation_image_pil
+            }
+    
+        except Exception as e:
+            logging.error(f"Visual analysis error: {e}")
+            raise ValueError(f"Visual analysis failed: {str(e)}")
+
+
+    def query_guidelines(self, query: str):
+        """Query the knowledge base for relevant guidelines"""
+        try:
+            vector_store = self.knowledge_base_cache.get("vectorstore")
+            if not vector_store:
+                return "Knowledge base unavailable - clinical guidelines not loaded"
+
+            # Retrieve relevant documents
+            retriever = vector_store.as_retriever(search_kwargs={"k": 10})
+            docs = retriever.invoke(query)
+            
+            if not docs:
+                return "No relevant guidelines found for the query"
+
+            # Format the results
+            formatted_results = []
+            for doc in docs:
+                source = doc.metadata.get('source', 'Unknown')
+                page = doc.metadata.get('page', 'N/A')
+                content = doc.page_content.strip()
+                formatted_results.append(f"Source: {source}, Page: {page}\nContent: {content}")
+
+            return "\n\n".join(formatted_results)
+
+        except Exception as e:
+            logging.error(f"Guidelines query error: {e}")
+            return f"Error querying guidelines: {str(e)}"
+
+    
+    def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
+        """Generate comprehensive medical report using MedGemma"""
+        try:
+            if 'medgemma_pipe' not in self.models_cache:
+                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+
+            max_tokens = max_new_tokens or self.config.MAX_NEW_TOKENS
+            
+            # Get detection and segmentation images if available
+            detection_image = visual_results.get('detection_image_pil', None)
+            segmentation_image = visual_results.get('segmentation_image_pil', None)
+            
+            # Create image paths for report
+            detection_path = visual_results.get('detection_image_path', '')
+            segmentation_path = visual_results.get('segmentation_image_path', '')
+            
+            # Create detailed prompt for medical analysis with image paths
+            prompt = f"""
+# Wound Care Report
+
+## Patient Information
+{patient_info}
+
+## Visual Analysis Summary
+- Wound Type: {visual_results.get('wound_type', 'Unknown')}
+- Length: {visual_results.get('length_cm', 0)} cm
+- Breadth: {visual_results.get('breadth_cm', 0)} cm
+- Surface Area: {visual_results.get('surface_area_cm2', 0)} cm²
+- Detection Confidence: {visual_results.get('detection_confidence', 0):.2f}
+
+## Clinical Reference
+{guideline_context}
+
+You are SmartHeal-AI Agent, a world-class wound care AI specialist trained in clinical wound assessment and guideline-based treatment planning.
+Your task is to process the following structured inputs (patient data, wound measurements, clinical guidelines, and image) and perform **clinical reasoning and decision-making** to generate a complete wound care report.
+---
+🔍 **YOUR PROCESS — FOLLOW STRICTLY:**
+### Step 1: Clinical Reasoning (Chain-of-Thought)
+Use the provided information to think step-by-step about:
+- Patient’s risk factors (e.g. diabetes, age, healing limitations)
+- Wound characteristics (size, tissue appearance, moisture, infection signs)
+- Visual clues from the image (location, granulation, maceration, inflammation, surrounding skin)
+
+---
+-Step 2: Structured Clinical Report
+Generate the following report sections using markdown and medical terminology:
+**1. Clinical Summary**
+- Describe wound appearance and tissue types (e.g., slough, necrotic, granulating, epithelializing)
+- Include size, wound bed condition, peri-wound skin, and signs of infection or biofilm
+- Mention inferred location (e.g., heel, forefoot) if image allows
+- Summarize patient's systemic risk profile
+**2. Medicinal & Dressing Recommendations**
+Based on your analysis:
+- Recommend specific **wound care dressings** (e.g., hydrocolloid, alginate, foam, antimicrobial silver, etc.) suitable to wound moisture level and infection risk
+- Propose **topical or systemic agents** ONLY if relevant — include name classes (e.g., antiseptic: povidone iodine, antibiotic ointments, enzymatic debriders)
+- Mention **techniques** (e.g., sharp debridement, NPWT, moisture balance, pressure offloading, dressing frequency)
+- Avoid repeating guidelines — **apply them**
+**3. Key Risk Factors**
+Explain how the patient’s condition (e.g., diabetic, poor circulation, advanced age, poor hygiene) may affect wound healing
+**4. Prognosis & Monitoring Advice**
+- Mention how often wound should be reassessed
+- Indicate signs to monitor for deterioration or improvement
+- Include when escalation to specialist is necessary
+
+**Note:** Every dressing change is a chance for wound reassessment. Always perform a thorough wound evaluation at each dressing change.
+"""
+
+            # Prepare messages for MedGemma with all available images
+            content_list = [{"type": "text", "text": prompt}]
+            
+            # Add original image
+            if image_pil:
+                content_list.insert(0, {"type": "image", "image": image_pil})
+                
+            # Add detection image if available
+            if detection_image:
+                content_list.insert(1, {"type": "image", "image": detection_image})
+                
+            # Add segmentation image if available
+            if segmentation_image:
+                content_list.insert(2, {"type": "image", "image": segmentation_image})
+            
+            messages = [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": "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."}],
+                },
+                {
+                    "role": "user",
+                    "content": content_list
+                }
+            ]
+
+            # Generate report using MedGemma
+            output = self.models_cache['medgemma_pipe'](
+                text=messages,
+                max_new_tokens=1024,
+                do_sample=False,
+            )
+            
+            generated_content = output[0]['generated_text'][-1].get('content', '').strip()
+            
+            # Include image paths in the final report for display in UI
+            if generated_content:
+                # Add image paths to the report for frontend display
+                image_paths_section = f"""
+## Analysis Images
+- Original Image: {image_pil}
+- Detection Image: {detection_path}
+- Segmentation Image: {segmentation_path}
+"""
+                generated_content = image_paths_section + generated_content
+                
+            return generated_content if generated_content else self._generate_fallback_report(patient_info, visual_results, guideline_context)
+
+        except Exception as e:
+            logging.error(f"MedGemma report generation error: {e}")
+            return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+
+    def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
+        """Generate a fallback report when MedGemma is not available"""
+        # Get image paths for report
+        detection_path = visual_results.get('detection_image_path', 'Not available')
+        segmentation_path = visual_results.get('segmentation_image_path', 'Not available')
+        
+        report = f"""
+# Wound Analysis Report
+## Patient Information
+{patient_info}
+
+## Visual Analysis Results
+- **Wound Type**: {visual_results.get('wound_type', 'Unknown')}
+- **Dimensions**: {visual_results.get('length_cm', 0)} cm × {visual_results.get('breadth_cm', 0)} cm
+- **Surface Area**: {visual_results.get('surface_area_cm2', 0)} cm²
+- **Detection Confidence**: {visual_results.get('detection_confidence', 0):.2f}
+
+## Analysis Images
+- **Detection Image**: {detection_path}
+- **Segmentation Image**: {segmentation_path}
+
+## Assessment
+Based on the visual analysis, this appears to be a {visual_results.get('wound_type', 'wound')} with measurable dimensions. 
+
+## Recommendations
+- Continue monitoring wound healing progress
+- Maintain proper wound hygiene
+- Follow appropriate dressing protocols
+- Seek medical attention if signs of infection develop
+
+## Clinical Guidelines
+{guideline_context[:500]}...
+
+*Note: This is an automated analysis. Please consult with a healthcare professional for definitive diagnosis and treatment.*
+"""
+        return report
+
+    def save_and_commit_image(self, image_pil):
+        """Save image locally and optionally upload to HuggingFace dataset"""
+        try:
+            # Ensure uploads directory exists
+            os.makedirs(self.config.UPLOADS_DIR, exist_ok=True)
+            
+            # Generate filename with timestamp
+            filename = f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.png"
+            local_path = os.path.join(self.config.UPLOADS_DIR, filename)
+            
+            # Save image locally
+            image_pil.convert("RGB").save(local_path)
+            logging.info(f"Image saved locally: {local_path}")
+
+            # Upload to HuggingFace dataset if configured
+            if self.config.HF_TOKEN and self.config.DATASET_ID:
+                try:
+                    api = HfApi()
+                    api.upload_file(
+                        path_or_fileobj=local_path,
+                        path_in_repo=f"images/{filename}",
+                        repo_id=self.config.DATASET_ID,
+                        repo_type="dataset",
+                        commit_message=f"Upload wound image: {filename}"
+                    )
+                    logging.info("✅ Image uploaded to HuggingFace dataset")
+                except Exception as e:
+                    logging.warning(f"HuggingFace upload failed: {e}")
+
+            return local_path
+
+        except Exception as e:
+            logging.error(f"Image saving error: {e}")
+            return None
+
+    @spaces.GPU(enable_queue=True, duration=120) 
+    def full_analysis_pipeline(self, image, questionnaire_data):
+        """Complete analysis pipeline with real-time models"""
+        try:
+            # Save the image
+            saved_path = self.save_and_commit_image(image)
+            
+            # Perform visual analysis
+            visual_results = self.perform_visual_analysis(image)
+            
+            # Format patient information
+            patient_info = ", ".join([f"{k}: {v}" for k, v in questionnaire_data.items() if v])
+            
+            # Create query for guidelines
+            wound_type = visual_results.get('wound_type', 'wound')
+            moisture = questionnaire_data.get('moisture', 'unknown')
+            infection = questionnaire_data.get('infection', 'unknown')
+            diabetic = questionnaire_data.get('diabetic', 'unknown')
+            
+            query = f"best practices for managing a {wound_type} with moisture level '{moisture}' and signs of infection '{infection}' in a patient who is diabetic '{diabetic}'"
+            
+            # Query guidelines
+            guideline_context = self.query_guidelines(query)
+            
+            # Generate final report
+            final_report = self.generate_final_report(patient_info, visual_results, guideline_context, image)
+            
+            return {
+                'success': True,
+                'visual_analysis': visual_results,
+                'report': final_report,
+                'saved_image_path': saved_path,
+                'timestamp': datetime.now().isoformat()
+            }
+
+        except Exception as e:
+            logging.error(f"Full analysis pipeline error: {e}")
+            return {
+                'success': False,
+                'error': str(e),
+                'timestamp': datetime.now().isoformat()
+            }
+
+    # Legacy methods for backward compatibility
+    def analyze_wound(self, image, questionnaire_data):
+        """Legacy method for backward compatibility"""
+        try:
+            # Convert string path to PIL Image if needed
+            if isinstance(image, str):
+                try:
+                    from PIL import Image
+                    image = Image.open(image)
+                    logging.info(f"Converted string path to PIL Image: {image}")
+                except Exception as e:
+                    logging.error(f"Error converting string path to image: {e}")
+            
+            # Ensure we have a PIL Image object
+            if not isinstance(image, Image.Image):
+                try:
+                    from PIL import Image
+                    import io
+                    
+                    # If it's a file-like object
+                    if hasattr(image, 'read'):
+                        # Reset file pointer if possible
+                        if hasattr(image, 'seek'):
+                            image.seek(0)
+                        image = Image.open(image)
+                        logging.info("Converted file-like object to PIL Image")
+                except Exception as e:
+                    logging.error(f"Error ensuring image is PIL Image: {e}")
+                    raise ValueError(f"Invalid image format: {type(image)}")
+            
+            result = self.full_analysis_pipeline(image, questionnaire_data)
+            
+            if result['success']:
+                return {
+                    'timestamp': result['timestamp'],
+                    'summary': f"Analysis completed for {questionnaire_data.get('patient_name', 'patient')}",
+                    'recommendations': result['report'],
+                    'wound_detection': {
+                        'status': 'success',
+                        'detections': [result['visual_analysis']],
+                        'total_wounds': 1
+                    },
+                    'segmentation_result': {
+                        'status': 'success',
+                        'wound_area_percentage': result['visual_analysis'].get('surface_area_cm2', 0)
+                    },
+                    'risk_assessment': self._assess_risk_legacy(questionnaire_data),
+                    'guideline_recommendations': [result['report'][:200] + "..."]
+                }
+            else:
+                return {
+                    'timestamp': result['timestamp'],
+                    'summary': f"Analysis failed: {result['error']}",
+                    'recommendations': "Please consult with a healthcare professional.",
+                    'wound_detection': {'status': 'error', 'message': result['error']},
+                    'segmentation_result': {'status': 'error', 'message': result['error']},
+                    'risk_assessment': {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []},
+                    'guideline_recommendations': ["Analysis unavailable due to error"]
+                }
+                
+        except Exception as e:
+            logging.error(f"Legacy analyze_wound error: {e}")
+            return {
+                'timestamp': datetime.now().isoformat(),
+                'summary': f"Analysis error: {str(e)}",
+                'recommendations': "Please consult with a healthcare professional.",
+                'wound_detection': {'status': 'error', 'message': str(e)},
+                'segmentation_result': {'status': 'error', 'message': str(e)},
+                'risk_assessment': {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []},
+                'guideline_recommendations': ["Analysis unavailable due to error"]
+            }
+
+    def _assess_risk_legacy(self, questionnaire_data):
+        """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
+            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
+            }
+            
+        except Exception as e:
+            logging.error(f"Risk assessment error: {e}")
             return {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []}