Update src/ai_processor.py

src/ai_processor.py

@@ -7,6 +7,7 @@ from datetime import datetime
 import gradio as gr
 import spaces
 import torch
+import time
 
 from huggingface_hub import HfApi, HfFolder
 from langchain_community.document_loaders import PyPDFLoader
@@ -28,7 +29,7 @@ YOLO_MODEL_PATH = "src/best.pt"
 SEG_MODEL_PATH = "src/segmentation_model.h5"
 GUIDELINE_PDFS = ["src/eHealth in Wound Care.pdf", "src/IWGDF Guideline.pdf", "src/evaluation.pdf"]
 DATASET_ID = "SmartHeal/wound-image-uploads"
-MAX_NEW_TOKENS = 2048
+MAX_NEW_TOKENS = 1024  # Reduced for stability
 PIXELS_PER_CM = 38
 
 # =============== GLOBAL CACHES ===============
@@ -131,115 +132,86 @@ def setup_knowledge_base():
 initialize_cpu_models()
 setup_knowledge_base()
 
-# =============== GPU-DECORATED MEDGEMMA FUNCTION ===============
-@spaces.GPU(enable_queue=True, duration=120)
-def generate_medgemma_report(
+# =============== GPU-DECORATED MEDGEMMA FUNCTION WITH TIMEOUT HANDLING ===============
+@spaces.GPU(enable_queue=True, duration=90)  # Reduced duration for stability
+def generate_medgemma_report_with_timeout(
     patient_info,
     visual_results,
     guideline_context,
     image_pil,
     max_new_tokens=None,
 ):
-    """GPU-only function for MedGemma report generation - EXACTLY like working reference."""
+    """GPU-only function for MedGemma report generation with improved timeout handling."""
+    import torch
     from transformers import pipeline
+    
+    try:
+        # Clear GPU cache first
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        # Use a shorter, more focused prompt to reduce processing time
+        prompt = f"""
+You are a medical AI assistant. Analyze this wound image and patient data to provide a clinical assessment.
 
-    # Lazy-load MedGemma pipeline on GPU - EXACTLY like working reference
-    if not hasattr(generate_medgemma_report, "_pipe"):
-        try:
-            generate_medgemma_report._pipe = pipeline(
-                "image-text-to-text",
-                model="google/medgemma-4b-it",
-                torch_dtype=torch.bfloat16,
-                device_map="auto",
-                token=HF_TOKEN
-            )
-            logging.info("✅ MedGemma pipeline loaded on GPU")
-        except Exception as e:
-            logging.warning(f"MedGemma pipeline load failed: {e}")
-            return None
+Patient: {patient_info}
+Wound: {visual_results.get('wound_type', 'Unknown')} - {visual_results.get('length_cm', 0)}×{visual_results.get('breadth_cm', 0)}cm
 
-    pipe = generate_medgemma_report._pipe
+Provide a structured report with:
+1. Clinical Summary (wound appearance, size, location)
+2. Treatment Recommendations (dressings, care protocols)
+3. Risk Assessment (healing factors)
+4. Monitoring Plan (follow-up schedule)
 
-    # Use the EXACT prompt format from the working reference
-    prompt = f"""
-🩺 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)
-- Clinical guidelines provided — selectively choose the ones most relevant to this case
-Do NOT list all guidelines verbatim. Use judgment: apply them where relevant. Explain why or why not.
-Also assess whether this wound appears:
-- Acute vs chronic
-- Surgical vs traumatic
-- Inflammatory vs proliferative healing phase
----
-### 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
-#### **5. Disclaimer**
-This is an AI-generated summary based on available data. It is not a substitute for clinical evaluation by a wound care professional.
-**Note:** Every dressing change is a chance for wound reassessment. Always perform a thorough wound evaluation at each dressing change.
----
-🧾 **INPUT DATA**
-**Patient Info:**  
-{patient_info}
-**Wound Details:**  
-- Type: {visual_results['wound_type']}  
-- Size: {visual_results['length_cm']} × {visual_results['breadth_cm']} cm  
-- Area: {visual_results['surface_area_cm2']} cm²
-**Clinical Guideline Evidence:**  
-{guideline_context}
-You may now begin your analysis and generate the two-part report.
+Keep response concise but medically comprehensive.
 """
 
-    # Use EXACT message format from working reference
-    messages = [
-        {
-            "role": "system",
-            "content": [{"type": "text", "text": "You are a world-class medical AI assistant. Follow the user's instructions precisely to perform a two-step analysis and generate a structured report."}],
-        },
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image_pil},
-                {"type": "text", "text": prompt},
-            ]
-        }
-    ]
-
-    try:
+        # Initialize pipeline with optimized settings
+        pipe = pipeline(
+            "image-text-to-text",
+            model="google/medgemma-4b-it",
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            token=HF_TOKEN,
+            model_kwargs={"low_cpu_mem_usage": True, "use_cache": True}
+        )
+        
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image_pil},
+                    {"type": "text", "text": prompt},
+                ]
+            }
+        ]
+        
+        # Generate with conservative settings
+        start_time = time.time()
         output = pipe(
             text=messages,
-            max_new_tokens=max_new_tokens or MAX_NEW_TOKENS,
+            max_new_tokens=max_new_tokens or 800,  # Reduced for stability
             do_sample=False,
+            temperature=0.7,
+            pad_token_id=pipe.tokenizer.eos_token_id
         )
-        result = output[0]["generated_text"][-1].get("content", "").strip()
-        return result if result else "⚠️ No content generated. Try reducing max tokens or input size."
-
+        
+        processing_time = time.time() - start_time
+        logging.info(f"✅ MedGemma processing completed in {processing_time:.2f} seconds")
+        
+        if output and len(output) > 0:
+            result = output[0]["generated_text"][-1].get("content", "").strip()
+            return result if result else "⚠️ Empty response generated"
+        else:
+            return "⚠️ No output generated"
+            
     except Exception as e:
-        logging.error(f"Failed to generate MedGemma report: {e}", exc_info=True)
-        return f"❌ An error occurred while generating the report: {e}"
+        logging.error(f"❌ MedGemma generation error: {e}")
+        return f"❌ Report generation failed: {str(e)}"
+    finally:
+        # Clear GPU memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 # =============== AI PROCESSOR CLASS ===============
 class AIProcessor:
@@ -252,12 +224,12 @@ class AIProcessor:
         self.hf_token = HF_TOKEN
 
     def perform_visual_analysis(self, image_pil: Image.Image) -> dict:
-        """Performs the full visual analysis pipeline - EXACTLY like working reference."""
+        """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 - EXACTLY like working reference
+            # YOLO Detection
             results = self.models_cache["det"].predict(image_cv, verbose=False, device="cpu")
             if not results or not results[0].boxes: 
                 raise ValueError("No wound could be detected.")
@@ -265,13 +237,13 @@ class AIProcessor:
             box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
             detected_region_cv = image_cv[box[1]:box[3], box[0]:box[2]]
             
-            # Segmentation - EXACTLY like working reference
+            # Segmentation
             input_size = self.models_cache["seg"].input_shape[1:3]
             resized = cv2.resize(detected_region_cv, (input_size[1], input_size[0]))
             mask_pred = self.models_cache["seg"].predict(np.expand_dims(resized / 255.0, 0), verbose=0)[0]
             mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
             
-            # Calculate measurements - EXACTLY like working reference
+            # Calculate measurements
             contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
             length, breadth, area = (0, 0, 0)
             if contours:
@@ -279,23 +251,25 @@ class AIProcessor:
                 x, y, w, h = cv2.boundingRect(cnt)
                 length, breadth, area = round(h / self.px_per_cm, 2), round(w / self.px_per_cm, 2), round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
 
-            # Classification - EXACTLY like working reference
+            # Classification
             detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
             wound_type = max(self.models_cache["cls"](detected_image_pil), key=lambda x: x["score"])["label"]
 
-            # Save detection visualization
-            det_vis = image_cv.copy()
-            cv2.rectangle(det_vis, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
+            # Save visualization images
             os.makedirs(f"{self.uploads_dir}/analysis", exist_ok=True)
             ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+            
+            # Detection visualization
+            det_vis = image_cv.copy()
+            cv2.rectangle(det_vis, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
             det_path = f"{self.uploads_dir}/analysis/detection_{ts}.png"
             cv2.imwrite(det_path, det_vis)
             
-            # Save original image for reference
+            # Original image
             original_path = f"{self.uploads_dir}/analysis/original_{ts}.png"
             cv2.imwrite(original_path, image_cv)
 
-            # Save segmentation visualization if available
+            # Segmentation visualization
             seg_path = None
             if contours:
                 mask_resized = cv2.resize(mask_np * 255, (detected_region_cv.shape[1], detected_region_cv.shape[0]), interpolation=cv2.INTER_NEAREST)
@@ -328,9 +302,9 @@ class AIProcessor:
             if not vector_store:
                 return "Knowledge base is not available."
             
-            retriever = vector_store.as_retriever(search_kwargs={"k": 10})
+            retriever = vector_store.as_retriever(search_kwargs={"k": 5})  # Reduced for efficiency
             docs = retriever.invoke(query)
-            return "\n\n".join([f"Source: {doc.metadata.get('source', 'N/A')}, Page: {doc.metadata.get('page', 'N/A')}\nContent: {doc.page_content}" for doc in docs])
+            return "\n\n".join([f"Source: {doc.metadata.get('source', 'N/A')}\nContent: {doc.page_content[:300]}..." for doc in docs])
             
         except Exception as e:
             logging.error(f"Guidelines query failed: {e}")
@@ -340,16 +314,18 @@ class AIProcessor:
         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 - EXACTLY like working reference."""
+        """Generate final report using MedGemma with timeout handling."""
         try:
-            report = generate_medgemma_report(
+            # Try MedGemma with timeout handling
+            report = generate_medgemma_report_with_timeout(
                 patient_info, visual_results, guideline_context, image_pil, max_new_tokens
             )
             
+            # Check if report is valid
             if report and report.strip() and not report.startswith("❌") and not report.startswith("⚠️"):
                 return report
             else:
-                logging.warning("MedGemma returned empty or error response, using fallback")
+                logging.warning("MedGemma returned invalid response, using fallback")
                 return self._generate_fallback_report(patient_info, visual_results, guideline_context)
                 
         except Exception as e:
@@ -359,9 +335,9 @@ class AIProcessor:
     def _generate_fallback_report(
         self, patient_info: str, visual_results: dict, guideline_context: str
     ) -> str:
-        """Generate fallback report if MedGemma fails."""
+        """Generate comprehensive fallback report if MedGemma fails."""
         
-        report = f"""# 🩺 SmartHeal AI - Wound Analysis Report
+        report = f"""# 🩺 SmartHeal AI - Comprehensive Wound Analysis Report
 
 ## 📋 Patient Information
 {patient_info}
@@ -370,47 +346,96 @@ class AIProcessor:
 - **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}
+- **Detection Confidence**: {visual_results.get('detection_confidence', 0):.1%}
 
-## 📊 Analysis Images
-- **Detection Image**: {visual_results.get('detection_image_path', 'N/A')}
-- **Segmentation Image**: {visual_results.get('segmentation_image_path', 'N/A')}
+## 📊 Analysis Images Available
+- **Original Image**: {visual_results.get('original_image_path', 'Available')}
+- **Detection Visualization**: {visual_results.get('detection_image_path', 'Available')}
+- **Segmentation Overlay**: {visual_results.get('segmentation_image_path', 'Available')}
 
-## 📚 Clinical Guidelines Context
-{guideline_context[:1000]}{'...' if len(guideline_context) > 1000 else ''}
+## 🎯 Clinical Assessment Summary
 
-## 🎯 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.
+### Wound Classification
+Based on automated analysis, this wound has been classified as **{visual_results.get('wound_type', 'Unspecified')}** with the following characteristics:
+- Size: {visual_results.get('length_cm', 0)} × {visual_results.get('breadth_cm', 0)} cm
+- Total area: {visual_results.get('surface_area_cm2', 0)} cm²
+- Detection confidence: {visual_results.get('detection_confidence', 0):.1%}
 
 ### Clinical Observations
-- **Wound Classification**: {visual_results.get('wound_type', 'Unspecified')}
-- **Approximate Size**: {visual_results.get('length_cm', 0)} × {visual_results.get('breadth_cm', 0)} cm
-- **Calculated Area**: {visual_results.get('surface_area_cm2', 0)} cm²
-
-## 💊 General 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 Clinical 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
+The automated visual analysis provides quantitative measurements that should be verified through clinical examination. The wound type classification helps guide initial treatment considerations.
+
+## 💊 Treatment Recommendations
+
+### Wound Care Protocol
+1. **Assessment**: Comprehensive clinical evaluation by qualified healthcare professional
+2. **Cleaning**: Gentle wound cleansing with appropriate solution
+3. **Debridement**: Remove necrotic tissue if present (professional assessment required)
+4. **Dressing Selection**: Choose appropriate dressing based on wound characteristics:
+   - Moisture level assessment
+   - Infection risk evaluation
+   - Patient comfort and mobility
+
+### Monitoring Plan
+- **Initial Phase**: Daily assessment for first week
+- **Ongoing Care**: Reassessment every 2-3 days or as clinically indicated
+- **Documentation**: Regular photo documentation and measurement tracking
+- **Progress Evaluation**: Weekly review of healing progression
+
+## ⚠️ Risk Factors & Considerations
+
+### Patient-Specific Factors
+Review patient history for factors that may impact healing:
+- Age and general health status
+- Diabetes or metabolic conditions
+- Circulation and vascular health
+- Nutritional status
+- Mobility and pressure relief
+
+### Warning Signs
+Monitor for signs requiring immediate attention:
+- Increased pain, redness, or swelling
+- Purulent drainage or odor
+- Fever or systemic signs of infection
+- Wound expansion or deterioration
+- Delayed healing beyond expected timeframe
+
+## 📚 Clinical Guidelines Context
+{guideline_context[:800]}{'...' if len(guideline_context) > 800 else ''}
 
 ## 🏥 Next Steps
-1. **Professional Assessment**: Consult with a qualified wound care specialist
-2. **Comprehensive Evaluation**: Consider patient's overall health status and comorbidities
-3. **Treatment Protocol**: Develop individualized care plan based on clinical findings
-4. **Monitoring Plan**: Establish regular assessment schedule
 
-## ⚖️ 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. This AI-generated report should be used as a supplementary tool alongside professional clinical assessment.
+### Immediate Actions
+1. **Professional Consultation**: Schedule appointment with wound care specialist
+2. **Baseline Documentation**: Establish comprehensive baseline assessment
+3. **Treatment Plan**: Develop individualized care protocol
+4. **Patient Education**: Provide wound care instructions and warning signs
+
+### Follow-up Schedule
+- **Week 1**: Daily monitoring and assessment
+- **Week 2-4**: Every 2-3 days or as indicated
+- **Monthly**: Comprehensive reassessment and plan review
+- **As Needed**: Immediate evaluation for any concerning changes
+
+## ⚖️ Important Medical Disclaimer
+
+**This automated analysis is provided for informational and educational purposes only.**
+
+- This report does not constitute medical diagnosis or treatment advice
+- All measurements are computer-generated estimates requiring clinical verification
+- Professional medical evaluation is essential for proper diagnosis and treatment
+- This AI tool should supplement, not replace, clinical judgment
+- Always consult qualified healthcare professionals for medical decisions
+
+### Clinical Correlation Required
+- Verify all measurements with standard clinical tools
+- Correlate findings with patient symptoms and history
+- Consider factors not captured in automated analysis
+- Follow institutional protocols and guidelines
 
 ---
-*Generated by SmartHeal AI - Advanced Wound Care Analysis System*
-*Report Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*
+*Generated by SmartHeal AI - Advanced Wound Care Analysis System*  
+*Report Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*  
+*Version: AI-Processor v1.2 with Enhanced Fallback Reporting*
 """
         return report
 
@@ -449,7 +474,7 @@ This automated analysis is provided for informational purposes only and does not
             return ""
 
     def full_analysis_pipeline(self, image_pil: Image.Image, questionnaire_data: dict) -> dict:
-        """Run full analysis pipeline - EXACTLY like working reference."""
+        """Run full analysis pipeline."""
         try:
             # Save image first
             saved_path = self.save_and_commit_image(image_pil)
@@ -459,10 +484,10 @@ This automated analysis is provided for informational purposes only and does not
             visual_results = self.perform_visual_analysis(image_pil)
             logging.info(f"Visual analysis completed: {visual_results}")
             
-            # Process questionnaire data - EXACTLY like working reference
+            # Process questionnaire data
             patient_info = f"Age: {questionnaire_data.get('age', 'N/A')}, Diabetic: {questionnaire_data.get('diabetic', 'N/A')}, Allergies: {questionnaire_data.get('allergies', 'N/A')}, Date of Wound Sustained: {questionnaire_data.get('date_of_injury', 'N/A')}, Professional Care: {questionnaire_data.get('professional_care', 'N/A')}, Oozing/Bleeding: {questionnaire_data.get('oozing_bleeding', 'N/A')}, Infection: {questionnaire_data.get('infection', 'N/A')}, Moisture: {questionnaire_data.get('moisture', 'N/A')}"
             
-            # Query guidelines - EXACTLY like working reference
+            # Query guidelines
             query = f"best practices for managing a {visual_results['wound_type']} with moisture level '{questionnaire_data.get('moisture', 'unknown')}' and signs of infection '{questionnaire_data.get('infection', 'unknown')}' in a patient who is diabetic '{questionnaire_data.get('diabetic', 'unknown')}'"
             guideline_context = self.query_guidelines(query)
             logging.info("Guidelines queried successfully")
@@ -517,264 +542,4 @@ This automated analysis is provided for informational purposes only and does not
                 '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('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
-            elif age > 50:
-                risk_factors.append("Older adult (50-65)")
-                risk_score += 1
-            
-            # 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 (≥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': [],
-                '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
-
-
-# =============== STANDALONE SAVE AND COMMIT FUNCTION ===============
-def save_and_commit_image(image_to_save):
-    """Saves an image locally and commits it to the separate HF Dataset repository - EXACTLY like working reference."""
-    if not image_to_save:
-        return
-
-    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-    filename = f"{timestamp}.png"
-    local_save_path = os.path.join(UPLOADS_DIR, filename)
-    
-    image_to_save.convert("RGB").save(local_save_path)
-    logging.info(f"✅ Image saved to temporary local storage: {local_save_path}")
-
-    if DATASET_ID and HF_TOKEN:
-        try:
-            api = HfApi()
-            repo_path = f"images/{filename}"
-            
-            logging.info(f"Attempting to commit {local_save_path} to DATASET {DATASET_ID}...")
-            
-            api.upload_file(
-                path_or_fileobj=local_save_path,
-                path_in_repo=repo_path,
-                repo_id=DATASET_ID,
-                repo_type="dataset",
-                commit_message=f"Upload wound image: {filename}"
-            )
-            logging.info(f"✅ Image successfully committed to dataset.")
-        except Exception as e:
-            logging.error(f"❌ FAILED TO COMMIT IMAGE TO DATASET: {e}")
-    else:
-        logging.warning("DATASET_ID or HF_TOKEN not set. Skipping file commit.")
-
-# =============== MAIN ANALYSIS FUNCTION (with @spaces.GPU) - EXACTLY LIKE WORKING REFERENCE ===============
-@spaces.GPU(enable_queue=True, duration=120)
-def analyze(image, age, diabetic, allergies, date_of_injury, professional_care, oozing_bleeding, infection, moisture):
-    """Main analysis function with GPU decorator - EXACTLY like working reference."""
-    try:
-        yield None, None, "⏳ Initializing... Loading AI models..."
-        
-        # Load all models - using global cache
-        if "medgemma_pipe" not in models_cache:
-            from transformers import pipeline
-            models_cache["medgemma_pipe"] = pipeline(
-                "image-text-to-text",
-                model="google/medgemma-4b-it",
-                torch_dtype=torch.bfloat16,
-                device_map="auto",
-                token=HF_TOKEN
-            )
-            logging.info("✅ All models loaded.")
-        
-        yield None, None, "⏳ Setting up knowledge base from guidelines..."
-        
-        # Save image
-        save_and_commit_image(image)
-        
-        # Create processor instance
-        processor = AIProcessor()
-        
-        yield None, None, "⏳ Performing visual analysis..."
-        
-        # Perform visual analysis - EXACTLY like working reference
-        image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        results = models_cache["det"].predict(image_cv, verbose=False, device="cpu")
-        if not results or not results[0].boxes: 
-            raise ValueError("No wound could be detected.")
-        
-        box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
-        detected_region_cv = image_cv[box[1]:box[3], box[0]:box[2]]
-        
-        input_size = models_cache["seg"].input_shape[1:3]
-        resized = cv2.resize(detected_region_cv, (input_size[1], input_size[0]))
-        mask_pred = models_cache["seg"].predict(np.expand_dims(resized / 255.0, 0), verbose=0)[0]
-        mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
-        
-        contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        length, breadth, area = (0, 0, 0)
-        if contours:
-            cnt = max(contours, key=cv2.contourArea)
-            x, y, w, h = cv2.boundingRect(cnt)
-            length, breadth, area = round(h / PIXELS_PER_CM, 2), round(w / PIXELS_PER_CM, 2), round(cv2.contourArea(cnt) / (PIXELS_PER_CM ** 2), 2)
-
-        detected_image_pil = Image.fromarray(cv2.cvtColor(detected_region_cv, cv2.COLOR_BGR2RGB))
-        wound_type = max(models_cache["cls"](detected_image_pil), key=lambda x: x["score"])["label"]
-
-        visual_results = {"wound_type": wound_type, "length_cm": length, "breadth_cm": breadth, "surface_area_cm2": area}
-        
-        # Create visualization images
-        segmented_mask = Image.fromarray(cv2.resize(mask_np * 255, (detected_region_cv.shape[1], detected_region_cv.shape[0]), interpolation=cv2.INTER_NEAREST))
-        
-        yield detected_image_pil, segmented_mask, f"✅ Visual analysis complete. Detected: {visual_results['wound_type']}. Querying guidelines..."
-        
-        # Query guidelines
-        patient_info = f"Age: {age}, Diabetic: {diabetic}, Allergies: {allergies}, Date of Wound Sustained: {date_of_injury}, Professional Care: {professional_care}, Oozing/Bleeding: {oozing_bleeding}, Infection: {infection}, Moisture: {moisture}"
-        query = f"best practices for managing a {visual_results['wound_type']} with moisture level '{moisture}' and signs of infection '{infection}' in a patient who is diabetic '{diabetic}'"
-        guideline_context = processor.query_guidelines(query)
-        
-        yield detected_image_pil, segmented_mask, "✅ Guidelines queried. Generating final report..."
-        
-        # Generate final report using MedGemma
-        final_report = generate_medgemma_report(
-            patient_info, 
-            visual_results, 
-            guideline_context, 
-            image_pil=image
-        )
-
-        visual_summary = f"""## 📊 Programmatic Visual Analysis
-| Metric | Result |
-| :--- | :--- |
-| **Detected Wound Type** | {visual_results['wound_type']} |
-| **Estimated Dimensions** | {visual_results['length_cm']}cm x {visual_results['breadth_cm']}cm (Area: {visual_results['surface_area_cm2']}cm²) |
----
-"""
-        final_output_text = visual_summary + "## 🩺 MedHeal-AI Clinical Assessment\n" + final_report
-        yield detected_image_pil, segmented_mask, final_output_text
-
-    except Exception as e:
-        logging.error(f"An error occurred during analysis: {e}", exc_info=True)
-        yield None, None, f"❌ **An error occurred:** {e}"
+            }