Update src/ai_processor.py

src/ai_processor.py

@@ -36,14 +36,19 @@ 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."""
+    import tensorflow as tf
+    tf.config.set_visible_devices([], 'GPU')  # Force CPU for TensorFlow
     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",
@@ -53,6 +58,7 @@ def load_classification_pipeline(hf_token):
     )
 
 def load_embedding_model():
+    """Load embedding model for knowledge base."""
     return HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2",
         model_kwargs={"device": "cpu"}
@@ -60,28 +66,34 @@ def load_embedding_model():
 
 # =============== MODEL INITIALIZATION ===============
 def initialize_cpu_models():
+    """Initialize all CPU-only models once."""
     global models_cache
+    
     if HF_TOKEN:
         HfFolder.save_token(HF_TOKEN)
         logging.info("✅ HuggingFace token set")
+    
     if "det" not in models_cache:
         try:
             models_cache["det"] = load_yolo_model(YOLO_MODEL_PATH)
             logging.info("✅ YOLO model loaded (CPU only)")
         except Exception as e:
             logging.error(f"YOLO load failed: {e}")
+
     if "seg" not in models_cache:
         try:
             models_cache["seg"] = load_segmentation_model(SEG_MODEL_PATH)
             logging.info("✅ Segmentation model loaded (CPU)")
         except Exception as e:
             logging.warning(f"Segmentation model not available: {e}")
+
     if "cls" not in models_cache:
         try:
             models_cache["cls"] = load_classification_pipeline(HF_TOKEN)
             logging.info("✅ Classification pipeline loaded (CPU)")
         except Exception as e:
             logging.warning(f"Classification pipeline not available: {e}")
+
     if "embedding_model" not in models_cache:
         try:
             models_cache["embedding_model"] = load_embedding_model()
@@ -90,9 +102,11 @@ def initialize_cpu_models():
             logging.warning(f"Embedding model not available: {e}")
 
 def setup_knowledge_base():
+    """Load PDF documents and create FAISS vector store."""
     global knowledge_base_cache
     if "vector_store" in knowledge_base_cache:
         return
+
     docs = []
     for pdf_path in GUIDELINE_PDFS:
         if os.path.exists(pdf_path):
@@ -102,6 +116,7 @@ def setup_knowledge_base():
                 logging.info(f"Loaded PDF: {pdf_path}")
             except Exception as e:
                 logging.warning(f"Failed to load PDF {pdf_path}: {e}")
+
     if docs and "embedding_model" in models_cache:
         splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
         chunks = splitter.split_documents(docs)
@@ -111,7 +126,7 @@ def setup_knowledge_base():
         knowledge_base_cache["vector_store"] = None
         logging.warning("Knowledge base unavailable")
 
-# Initialize models at startup
+# Initialize models on app startup
 initialize_cpu_models()
 setup_knowledge_base()
 
@@ -125,6 +140,8 @@ def generate_medgemma_report(
     segmentation_image_path,
     max_new_tokens=None,
 ):
+    """GPU-only function for MedGemma report generation."""
+    # Import GPU libraries ONLY here
     import torch
     from transformers import pipeline
     from PIL import Image
@@ -141,6 +158,7 @@ def generate_medgemma_report(
         "patient context."
     )
 
+    # Lazy-load MedGemma pipeline on GPU
     if not hasattr(generate_medgemma_report, "_pipe"):
         try:
             generate_medgemma_report._pipe = pipeline(
@@ -158,16 +176,44 @@ def generate_medgemma_report(
 
     pipe = generate_medgemma_report._pipe
 
+    # Load the original image that was analyzed
+    original_image = None
+    if detection_image_path and os.path.exists(detection_image_path.replace('detection_', 'original_')):
+        original_image = Image.open(detection_image_path.replace('detection_', 'original_'))
+    elif segmentation_image_path and os.path.exists(segmentation_image_path.replace('segmentation_', 'original_')):
+        original_image = Image.open(segmentation_image_path.replace('segmentation_', 'original_'))
+
+    # Compose messages
     msgs = [
         {"role": "system", "content": [{"type": "text", "text": default_system_prompt}]},
         {"role": "user", "content": []},
     ]
 
-    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 images if available
+    if original_image:
+        msgs[1]["content"].append({"type": "image", "image": original_image})
+    else:
+        # Fallback to detection or segmentation 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)})
+                break
+
+    # Attach text prompt
+    prompt = f"""## Patient Information
+{patient_info}
+
+## Visual Analysis Results
+- Wound Type: {visual_results.get('wound_type','Unknown')}
+- Dimensions: {visual_results.get('length_cm', 0)} x {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 Guidelines Context
+{guideline_context[:1500]}...
+
+Please provide a comprehensive wound care assessment and treatment recommendations based on the image and provided information."""
 
-    prompt = f"## Patient\n{patient_info}\n## Wound Type: {visual_results.get('wound_type','Unknown')}"
     msgs[1]["content"].append({"type": "text", "text": prompt})
 
     try:
@@ -188,176 +234,492 @@ class AIProcessor:
         self.hf_token = HF_TOKEN
 
     def perform_visual_analysis(self, image_pil: Image.Image) -> dict:
-        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")
+        """Performs the full visual analysis pipeline."""
+        try:
+            # Convert PIL to OpenCV format
+            image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
+            
+            # YOLO Detection
+            yolo_model = self.models_cache.get("det")
+            if yolo_model is None:
+                raise RuntimeError("YOLO model ('det') not loaded")
 
-        res = yolo.predict(img_cv, verbose=False, device="cpu")[0]
-        if not res.boxes:
-            raise ValueError("No wound detected")
+            results = yolo_model.predict(image_cv, verbose=False, device="cpu")
+            
+            if not results or not results[0].boxes or len(results[0].boxes) == 0:
+                raise ValueError("No wound detected in the image")
 
-        # Safely unpack detection boxes
-        try:
-            xyxy = res.boxes.xyxy.cpu().numpy()
-            if xyxy.shape[0] == 0:
-                raise ValueError("No detection boxes found")
-            x1, y1, x2, y2 = xyxy[0]
-        except Exception as e:
-            logging.warning(f"Error unpacking detection boxes: {e}")
-            raise
-
-        region = img_cv[int(y1):int(y2), int(x1):int(x2)]
-
-        # Save detection overlay
-        det_vis = img_cv.copy()
-        cv2.rectangle(det_vis, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
-        os.makedirs(f"{self.uploads_dir}/analysis", exist_ok=True)
-        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-        det_path = f"{self.uploads_dir}/analysis/detection_{ts}.png"
-        cv2.imwrite(det_path, det_vis)
-
-        # Segmentation
-        length, breadth, area = 0, 0, 0
-        seg_path = None
-        seg_model = self.models_cache.get("seg")
-        if seg_model:
-            try:
-                h, w = seg_model.input_shape[1:3]
-                inp = cv2.resize(region, (w, h)) / 255.0
-                mask_pred = seg_model.predict(inp[None])
-                if mask_pred.shape[1:3] != (h, w):
-                    # Resize if needed
-                    mask_pred = np.squeeze(mask_pred)
-                mask = (mask_pred[0, :, :, 0] > 0.5).astype(np.uint8)
-                mask_rs = cv2.resize(mask, (region.shape[1], region.shape[0]), interpolation=cv2.INTER_NEAREST)
-                # Save segmentation visualization
-                ov = region.copy()
-                ov[mask_rs == 1] = [0, 0, 255]
-                seg_vis = cv2.addWeighted(region, 0.7, ov, 0.3, 0)
-                seg_path = f"{self.uploads_dir}/analysis/segmentation_{ts}.png"
-                cv2.imwrite(seg_path, seg_vis)
-
-                # Find contours
-                cnts, _ = cv2.findContours(mask_rs, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-                if cnts:
-                    cnt = max(cnts, key=cv2.contourArea)
-                    x, y, w_box, h_box = cv2.boundingRect(cnt)
-                    length = round(h_box / self.px_per_cm, 2)
-                    breadth = round(w_box / self.px_per_cm, 2)
-                    area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
-            except Exception as e:
-                logging.warning(f"Segmentation processing error: {e}")
+            # Extract bounding box - handle different output formats
+            boxes_data = results[0].boxes.xyxy.cpu().numpy()
+            
+            if len(boxes_data.shape) == 1:
+                # Single detection case
+                if len(boxes_data) != 4:
+                    raise ValueError(f"Expected 4 coordinates, got {len(boxes_data)}")
+                x1, y1, x2, y2 = boxes_data.astype(int)
+            else:
+                # Multiple detections - take the first one
+                if boxes_data.shape[1] != 4:
+                    raise ValueError(f"Expected 4 coordinates per box, got {boxes_data.shape[1]}")
+                x1, y1, x2, y2 = boxes_data[0].astype(int)
+
+            # Validate coordinates
+            if x1 >= x2 or y1 >= y2 or x1 < 0 or y1 < 0:
+                raise ValueError("Invalid bounding box coordinates")
+
+            # Extract wound region
+            detected_region_cv = image_cv[y1:y2, x1:x2]
+            
+            if detected_region_cv.size == 0:
+                raise ValueError("Detected region is empty")
+
+            # Save detection visualization
+            det_vis = image_cv.copy()
+            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)
+            
+            # Save original image for reference
+            original_path = f"{self.uploads_dir}/analysis/original_{ts}.png"
+            cv2.imwrite(original_path, image_cv)
 
-        # Classification
-        wound_type = "Unknown"
-        cls_pipe = self.models_cache.get("cls")
-        if cls_pipe:
-            try:
-                preds = cls_pipe(Image.fromarray(cv2.cvtColor(region, cv2.COLOR_BGR2RGB)))
-                if preds:
-                    wound_type = max(preds, key=lambda x: x["score"])["label"]
-            except Exception as e:
-                logging.warning(f"Classification error: {e}")
-
-        return {
-            "wound_type": wound_type,
-            "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,
-        }
+            # Segmentation Analysis
+            length = breadth = area = 0
+            seg_path = None
+            
+            seg_model = self.models_cache.get("seg")
+            if seg_model is not None:
+                try:
+                    # Get input shape from model
+                    input_shape = seg_model.input_shape
+                    if len(input_shape) >= 3:
+                        h, w = input_shape[1:3]
+                    else:
+                        h, w = 256, 256  # Default fallback
+
+                    # Prepare input for segmentation
+                    resized = cv2.resize(detected_region_cv, (w, h))
+                    normalized_input = np.expand_dims(resized / 255.0, 0)
+                    
+                    # Predict mask
+                    mask_pred = seg_model.predict(normalized_input, verbose=0)
+                    
+                    # Handle different output formats
+                    if len(mask_pred.shape) == 4:
+                        mask_np = (mask_pred[0, :, :, 0] > 0.5).astype(np.uint8)
+                    elif len(mask_pred.shape) == 3:
+                        mask_np = (mask_pred[0, :, :] > 0.5).astype(np.uint8)
+                    else:
+                        raise ValueError(f"Unexpected segmentation output shape: {mask_pred.shape}")
+                    
+                    # Resize mask back to detection region size
+                    mask_resized = cv2.resize(
+                        mask_np * 255, 
+                        (detected_region_cv.shape[1], detected_region_cv.shape[0]), 
+                        interpolation=cv2.INTER_NEAREST
+                    )
+                    mask_resized = (mask_resized > 127).astype(np.uint8)
+
+                    # Create segmentation visualization
+                    overlay = detected_region_cv.copy()
+                    overlay[mask_resized == 1] = [0, 0, 255]  # Red overlay for wound area
+                    seg_vis = cv2.addWeighted(detected_region_cv, 0.7, overlay, 0.3, 0)
+                    
+                    seg_path = f"{self.uploads_dir}/analysis/segmentation_{ts}.png"
+                    cv2.imwrite(seg_path, seg_vis)
+
+                    # Calculate measurements
+                    contours, _ = cv2.findContours(mask_resized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                    if contours:
+                        # Get the largest contour
+                        largest_contour = max(contours, key=cv2.contourArea)
+                        
+                        # Calculate bounding rectangle
+                        bbox = cv2.boundingRect(largest_contour)
+                        if len(bbox) == 4:
+                            x, y, w_box, h_box = bbox
+                            length = round(h_box / self.px_per_cm, 2)
+                            breadth = round(w_box / self.px_per_cm, 2)
+                            area = round(cv2.contourArea(largest_contour) / (self.px_per_cm ** 2), 2)
+                        else:
+                            logging.warning(f"Unexpected bounding rect format: {bbox}")
+                    else:
+                        logging.info("No contours found in segmentation mask")
+                        
+                except Exception as seg_error:
+                    logging.error(f"Segmentation processing error: {seg_error}")
+                    seg_path = None
+
+            # Wound Classification
+            wound_type = "Unknown"
+            cls_pipeline = self.models_cache.get("cls")
+            if cls_pipeline is not None:
+                try:
+                    detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
+                    predictions = cls_pipeline(detected_image_pil)
+                    if predictions and len(predictions) > 0:
+                        best_pred = max(predictions, key=lambda x: x.get("score", 0))
+                        wound_type = best_pred.get("label", "Unknown")
+                except Exception as cls_error:
+                    logging.warning(f"Classification failed: {cls_error}")
+
+            # Extract confidence score
+            confidence = 0.0
+            if results[0].boxes.conf is not None and len(results[0].boxes.conf) > 0:
+                confidence = float(results[0].boxes.conf[0].cpu().item())
+
+            return {
+                "wound_type": wound_type,
+                "length_cm": length,
+                "breadth_cm": breadth,
+                "surface_area_cm2": area,
+                "detection_confidence": confidence,
+                "detection_image_path": det_path,
+                "segmentation_image_path": seg_path,
+                "original_image_path": original_path
+            }
+
+        except Exception as e:
+            logging.error(f"Visual analysis failed: {e}")
+            raise e
 
     def query_guidelines(self, query: str) -> str:
-        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
-        )
-
-    def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
-        det_path = visual_results.get("detection_image_path", "")
-        seg_path = visual_results.get("segmentation_image_path", "")
-        report = generate_medgemma_report(patient_info, visual_results, guideline_context, det_path, seg_path, max_new_tokens)
-        if report:
-            return report
-        return self._generate_fallback_report(patient_info, visual_results, guideline_context)
-
-    def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
-        return (
-            f"# Fallback Report\n{patient_info}\n"
-            f"Type: {visual_results.get('wound_type','Unknown')}\n"
-            f"Detection Image: {visual_results.get('detection_image_path','N/A')}\n"
-            f"Segmentation Image: {visual_results.get('segmentation_image_path','N/A')}\n"
-            f"Guidelines: {guideline_context[:200]}..."
-        )
+        """Query the knowledge base for relevant information."""
+        try:
+            vector_store = self.knowledge_base_cache.get("vector_store")
+            if not vector_store:
+                return "Clinical guidelines unavailable - knowledge base not loaded"
+            
+            retriever = vector_store.as_retriever(search_kwargs={"k": 10})
+            docs = retriever.invoke(query)
+            
+            if not docs:
+                return "No relevant guidelines found for the query"
+            
+            context = "\n\n".join([
+                f"Source: {doc.metadata.get('source', 'Unknown')}, Page: {doc.metadata.get('page', 'N/A')}\n{doc.page_content}" 
+                for doc in docs
+            ])
+            
+            return context
+            
+        except Exception as e:
+            logging.error(f"Guidelines query failed: {e}")
+            return f"Guidelines query failed: {str(e)}"
+
+    def generate_final_report(
+        self, patient_info: str, visual_results: dict, guideline_context: str, 
+        image_pil: Image.Image, max_new_tokens: int = None
+    ) -> str:
+        """Generate final report using MedGemma GPU pipeline."""
+        try:
+            det_path = visual_results.get("detection_image_path", "")
+            seg_path = visual_results.get("segmentation_image_path", "")
+            
+            report = generate_medgemma_report(
+                patient_info, visual_results, guideline_context, 
+                det_path, seg_path, max_new_tokens
+            )
+            
+            if report and report.strip():
+                return report
+            else:
+                return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+                
+        except Exception as e:
+            logging.error(f"MedGemma report generation failed: {e}")
+            return self._generate_fallback_report(patient_info, visual_results, guideline_context)
+
+    def _generate_fallback_report(
+        self, patient_info: str, visual_results: dict, guideline_context: str
+    ) -> str:
+        """Generate fallback report if MedGemma fails."""
+        
+        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**: {visual_results.get('detection_image_path', 'N/A')}
+- **Segmentation Image**: {visual_results.get('segmentation_image_path', 'N/A')}
+
+## Clinical Guidelines Context
+{guideline_context[:1000]}{'...' if len(guideline_context) > 1000 else ''}
+
+## Assessment Summary
+Based on the automated visual analysis, the wound has been classified as **{visual_results.get('wound_type', 'Unknown')}** with measurable dimensions. The detection confidence indicates the reliability of the automated assessment.
+
+## Recommendations
+1. **Clinical Evaluation**: This automated analysis should be supplemented with professional clinical assessment
+2. **Documentation**: Regular monitoring and documentation of wound progression is recommended
+3. **Treatment Planning**: Develop appropriate treatment protocol based on wound characteristics and patient factors
+4. **Follow-up**: Schedule appropriate follow-up intervals based on wound severity and healing progress
+
+## Important Notes
+- This is an automated analysis and should not replace professional medical judgment
+- All measurements are estimates based on computer vision algorithms
+- Clinical correlation is essential for proper diagnosis and treatment planning
+- Consider patient-specific factors not captured in this automated assessment
+
+## Disclaimer
+This automated analysis is provided for informational purposes only and does not constitute medical advice. Always consult with qualified healthcare professionals for proper diagnosis and treatment.
+"""
+        return report
 
     def save_and_commit_image(self, image_pil: Image.Image) -> str:
-        os.makedirs(self.uploads_dir, exist_ok=True)
-        fn = f"{datetime.now():%Y%m%d_%H%M%S}.png"
-        path = os.path.join(self.uploads_dir, fn)
-        image_pil.convert("RGB").save(path)
-        if self.hf_token and self.dataset_id:
-            try:
-                HfApi().upload_file(
-                    path_or_fileobj=path,
-                    path_in_repo=f"images/{fn}",
-                    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
+        """Save image locally and optionally commit to HF dataset."""
+        try:
+            os.makedirs(self.uploads_dir, exist_ok=True)
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filename = f"{timestamp}.png"
+            path = os.path.join(self.uploads_dir, filename)
+            
+            # Save image
+            image_pil.convert("RGB").save(path)
+            logging.info(f"✅ Image saved locally: {path}")
+            
+            # Upload to HuggingFace dataset if configured
+            if self.hf_token and self.dataset_id:
+                try:
+                    api = HfApi()
+                    api.upload_file(
+                        path_or_fileobj=path,
+                        path_in_repo=f"images/{filename}",
+                        repo_id=self.dataset_id,
+                        repo_type="dataset",
+                        token=self.hf_token,
+                        commit_message=f"Upload wound image: {filename}"
+                    )
+                    logging.info("✅ Image committed to HF dataset")
+                except Exception as e:
+                    logging.warning(f"HF upload failed: {e}")
+            
+            return path
+            
+        except Exception as e:
+            logging.error(f"Failed to save image: {e}")
+            return ""
 
-    def full_analysis_pipeline(self, image, questionnaire_data):
+    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)
-            vis = self.perform_visual_analysis(image)
-            info = ", ".join(f"{k}:{v}" for k,v in questionnaire_data.items() if v)
-            gc = self.query_guidelines(info)
-            report = self.generate_final_report(info, vis, gc, image)
-            return {'success': True, 'visual_analysis': vis, 'report': report, 'saved_image_path': saved}
+            # Save image first
+            saved_path = self.save_and_commit_image(image_pil)
+            logging.info(f"Image saved: {saved_path}")
+            
+            # Perform visual analysis
+            visual_results = self.perform_visual_analysis(image_pil)
+            logging.info(f"Visual analysis completed: {visual_results}")
+            
+            # Process questionnaire data
+            patient_info = ", ".join(f"{k}: {v}" for k, v in questionnaire_data.items() if v)
+            if not patient_info:
+                patient_info = "No patient information provided"
+            
+            # Query guidelines
+            query = f"wound care treatment for {visual_results.get('wound_type', 'wound')} "
+            if questionnaire_data.get('diabetic') == 'Yes':
+                query += "diabetic patient "
+            if questionnaire_data.get('infection') == 'Yes':
+                query += "with infection signs "
+                
+            guideline_context = self.query_guidelines(query)
+            logging.info("Guidelines queried successfully")
+            
+            # Generate final report
+            report = self.generate_final_report(patient_info, visual_results, guideline_context, image_pil)
+            logging.info("Report generated successfully")
+            
+            return {
+                'success': True, 
+                'visual_analysis': visual_results, 
+                'report': report, 
+                'saved_image_path': saved_path,
+                'guideline_context': guideline_context[:500] + "..." if len(guideline_context) > 500 else guideline_context
+            }
+            
         except Exception as e:
             logging.error(f"Pipeline error: {e}")
-            return {'success': False, 'error': str(e)}
-
-    def analyze_wound(self, image, questionnaire_data):
-        if isinstance(image, str):
-            image = Image.open(image)
-        return self.full_analysis_pipeline(image, questionnaire_data)
-
-    def _assess_risk_legacy(self, questionnaire_data):
-        risk_factors, risk_score = [], 0
+            return {
+                'success': False, 
+                'error': str(e),
+                'visual_analysis': {},
+                'report': f"Analysis failed: {str(e)}",
+                'saved_image_path': None,
+                'guideline_context': ""
+            }
+
+    def analyze_wound(self, image, questionnaire_data: dict) -> dict:
+        """Main analysis entry point - maintains original function name."""
+        try:
+            # Handle different image input formats
+            if isinstance(image, str):
+                if os.path.exists(image):
+                    image_pil = Image.open(image)
+                else:
+                    raise ValueError(f"Image file not found: {image}")
+            elif isinstance(image, Image.Image):
+                image_pil = image
+            elif isinstance(image, np.ndarray):
+                image_pil = Image.fromarray(image)
+            else:
+                raise ValueError(f"Unsupported image type: {type(image)}")
+            
+            return self.full_analysis_pipeline(image_pil, questionnaire_data)
+            
+        except Exception as e:
+            logging.error(f"Wound analysis error: {e}")
+            return {
+                'success': False, 
+                'error': str(e),
+                'visual_analysis': {},
+                'report': f"Analysis initialization failed: {str(e)}",
+                'saved_image_path': None,
+                'guideline_context': ""
+            }
+
+    def _assess_risk_legacy(self, questionnaire_data: dict) -> dict:
+        """Legacy risk assessment function - maintains original function name."""
+        risk_factors = []
+        risk_score = 0
+        
         try:
+            # Age assessment
             age = questionnaire_data.get('patient_age', 0)
+            if isinstance(age, str):
+                try:
+                    age = int(age)
+                except ValueError:
+                    age = 0
+                    
             if age > 65:
-                risk_factors.append("Advanced age (>65)"); risk_score += 2
+                risk_factors.append("Advanced age (>65)")
+                risk_score += 2
             elif age > 50:
-                risk_factors.append("Older adult (50-65)"); risk_score += 1
+                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
+            # Wound duration assessment
+            duration = str(questionnaire_data.get('wound_duration', '')).lower()
+            if any(term in duration for term in ['month', 'months', 'year', 'years']):
+                risk_factors.append("Chronic wound (>4 weeks)")
+                risk_score += 3
+            elif any(term in duration for term in ['week', 'weeks']):
+                # Try to extract number of weeks
+                import re
+                weeks_match = re.search(r'(\d+)\s*week', duration)
+                if weeks_match and int(weeks_match.group(1)) > 4:
+                    risk_factors.append("Chronic wound (>4 weeks)")
+                    risk_score += 3
             
+            # Pain level assessment
             pain = questionnaire_data.get('pain_level', 0)
+            if isinstance(pain, str):
+                try:
+                    pain = float(pain)
+                except ValueError:
+                    pain = 0
+                    
             if pain >= 7:
-                risk_factors.append("High pain level"); risk_score += 2
-            
-            hist = questionnaire_data.get('medical_history','').lower()
-            if 'diabetes' in hist:
-                risk_factors.append("Diabetes mellitus"); risk_score += 3
-            if 'vascular' in hist:
-                risk_factors.append("Vascular issues"); risk_score += 2
-            if 'immune' in hist:
-                risk_factors.append("Immune compromise"); risk_score += 2
-            
-            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}
+                risk_factors.append("High pain level (≥7/10)")
+                risk_score += 2
+            elif pain >= 5:
+                risk_factors.append("Moderate pain level (5-6/10)")
+                risk_score += 1
+            
+            # Medical history assessment
+            medical_history = str(questionnaire_data.get('medical_history', '')).lower()
+            diabetic_status = str(questionnaire_data.get('diabetic', '')).lower()
+            
+            if 'diabetes' in medical_history or 'yes' in diabetic_status:
+                risk_factors.append("Diabetes mellitus")
+                risk_score += 3
+            
+            if any(term in medical_history for term in ['vascular', 'circulation', 'arterial', 'venous']):
+                risk_factors.append("Vascular disease")
+                risk_score += 2
+                
+            if any(term in medical_history for term in ['immune', 'immunocompromised', 'steroid', 'chemotherapy']):
+                risk_factors.append("Immune system compromise")
+                risk_score += 2
+                
+            if any(term in medical_history for term in ['smoking', 'smoker', 'tobacco']):
+                risk_factors.append("Smoking history")
+                risk_score += 2
+            
+            # Infection signs
+            infection_signs = str(questionnaire_data.get('infection', '')).lower()
+            if 'yes' in infection_signs:
+                risk_factors.append("Signs of infection present")
+                risk_score += 3
+            
+            # Moisture level
+            moisture = str(questionnaire_data.get('moisture', '')).lower()
+            if any(term in moisture for term in ['wet', 'heavy', 'excessive']):
+                risk_factors.append("Excessive wound exudate")
+                risk_score += 1
+            
+            # Determine risk level
+            if risk_score >= 8:
+                risk_level = "Very High"
+            elif risk_score >= 6:
+                risk_level = "High"
+            elif risk_score >= 3:
+                risk_level = "Moderate"
+            else:
+                risk_level = "Low"
+                
+            return {
+                'risk_score': risk_score, 
+                'risk_level': risk_level, 
+                'risk_factors': risk_factors,
+                'recommendations': self._get_risk_recommendations(risk_level, 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': [],
+                'recommendations': ["Unable to assess risk due to data processing error"]
+            }
+    
+    def _get_risk_recommendations(self, risk_level: str, risk_factors: list) -> list:
+        """Generate risk-based recommendations."""
+        recommendations = []
+        
+        if risk_level in ["High", "Very High"]:
+            recommendations.append("Urgent referral to wound care specialist recommended")
+            recommendations.append("Consider daily wound monitoring")
+            recommendations.append("Implement aggressive wound care protocol")
+        elif risk_level == "Moderate":
+            recommendations.append("Regular wound care follow-up every 2-3 days")
+            recommendations.append("Monitor for signs of deterioration")
+        else:
+            recommendations.append("Standard wound care monitoring")
+            recommendations.append("Weekly assessment recommended")
+        
+        # Specific recommendations based on risk factors
+        if "Diabetes mellitus" in risk_factors:
+            recommendations.append("Strict glycemic control essential")
+            recommendations.append("Monitor for diabetic complications")
+        
+        if "Signs of infection present" in risk_factors:
+            recommendations.append("Consider antibiotic therapy")
+            recommendations.append("Increase wound cleaning frequency")
+        
+        if "Excessive wound exudate" in risk_factors:
+            recommendations.append("Use high-absorption dressings")
+            recommendations.append("More frequent dressing changes may be needed")
+        
+        return recommendations