Update src/ai_processor.py

src/ai_processor.py

@@ -6,6 +6,7 @@ from PIL import Image
 from datetime import datetime
 import gradio as gr
 import spaces
+import torch
 
 from huggingface_hub import HfApi, HfFolder
 from langchain_community.document_loaders import PyPDFLoader
@@ -23,9 +24,9 @@ if not os.path.exists(UPLOADS_DIR):
     logging.info(f"Created uploads directory: {UPLOADS_DIR}")
 
 HF_TOKEN = os.getenv("HF_TOKEN")
-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"]
+YOLO_MODEL_PATH = "best.pt"
+SEG_MODEL_PATH = "segmentation_model.h5"
+GUIDELINE_PDFS = ["eHealth in Wound Care.pdf", "IWGDF Guideline.pdf", "evaluation.pdf"]
 DATASET_ID = "SmartHeal/wound-image-uploads"
 MAX_NEW_TOKENS = 2048
 PIXELS_PER_CM = 38
@@ -136,38 +137,21 @@ def generate_medgemma_report(
     patient_info,
     visual_results,
     guideline_context,
-    detection_image_path,
-    segmentation_image_path,
+    image_pil,
     max_new_tokens=None,
 ):
-    """GPU-only function for MedGemma report generation."""
-    # Import GPU libraries ONLY here
-    import torch
+    """GPU-only function for MedGemma report generation - EXACTLY like working reference."""
     from transformers import pipeline
-    from PIL import Image
-
-    default_system_prompt = (
-        "You are a world-class medical AI assistant specializing in wound care "
-        "with expertise in wound assessment and treatment. Provide concise, "
-        "evidence-based medical assessments focusing on: (1) Precise wound "
-        "classification based on tissue type and appearance, (2) Specific "
-        "treatment recommendations with exact product names or interventions when "
-        "appropriate, (3) Objective evaluation of healing progression or deterioration "
-        "indicators, and (4) Clear follow-up timelines. Avoid general statements and "
-        "prioritize actionable insights based on the visual analysis measurements and "
-        "patient context."
-    )
 
-    # Lazy-load MedGemma pipeline on GPU
+    # 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",
-                device="cuda",
                 torch_dtype=torch.bfloat16,
-                offload_folder="offload",
-                token=HF_TOKEN,
+                device_map="auto",
+                token=HF_TOKEN
             )
             logging.info("✅ MedGemma pipeline loaded on GPU")
         except Exception as e:
@@ -176,52 +160,86 @@ 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": []},
-    ]
-
-    # 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
+    # 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.
+"""
 
-## 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."""
-
-    msgs[1]["content"].append({"type": "text", "text": prompt})
+    # 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:
-        out = pipe(text=msgs, max_new_tokens=max_new_tokens or MAX_NEW_TOKENS, do_sample=False)
-        return out[0]["generated_text"][-1].get("content", "")
+        output = pipe(
+            text=messages,
+            max_new_tokens=max_new_tokens or MAX_NEW_TOKENS,
+            do_sample=False,
+        )
+        result = output[0]["generated_text"][-1].get("content", "").strip()
+        return result if result else "⚠️ No content generated. Try reducing max tokens or input size."
+
     except Exception as e:
-        logging.error(f"Failed to generate MedGemma report: {e}")
-        return f"❌ An error occurred: {e}"
+        logging.error(f"Failed to generate MedGemma report: {e}", exc_info=True)
+        return f"❌ An error occurred while generating the report: {e}"
 
 # =============== AI PROCESSOR CLASS ===============
 class AIProcessor:
@@ -234,48 +252,40 @@ class AIProcessor:
         self.hf_token = HF_TOKEN
 
     def perform_visual_analysis(self, image_pil: Image.Image) -> dict:
-        """Performs the full visual analysis pipeline."""
+        """Performs the full visual analysis pipeline - EXACTLY like working reference."""
         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")
-
-            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")
-
-            # 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")
+            # YOLO Detection - EXACTLY like working reference
+            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.")
+                
+            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
+            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
+            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 / 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
+            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, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.rectangle(det_vis, (box[0], box[1]), (box[2], box[3]), (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"
@@ -285,101 +295,27 @@ class AIProcessor:
             original_path = f"{self.uploads_dir}/analysis/original_{ts}.png"
             cv2.imwrite(original_path, image_cv)
 
-            # Segmentation Analysis
-            length = breadth = area = 0
+            # Save segmentation visualization if available
             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,
+            if contours:
+                mask_resized = cv2.resize(mask_np * 255, (detected_region_cv.shape[1], detected_region_cv.shape[0]), interpolation=cv2.INTER_NEAREST)
+                overlay = detected_region_cv.copy()
+                overlay[mask_resized > 127] = [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)
+
+            visual_results = {
+                "wound_type": wound_type, 
+                "length_cm": length, 
+                "breadth_cm": breadth, 
                 "surface_area_cm2": area,
-                "detection_confidence": confidence,
+                "detection_confidence": float(results[0].boxes.conf[0].cpu().item()) if results[0].boxes.conf is not None else 0.0,
                 "detection_image_path": det_path,
                 "segmentation_image_path": seg_path,
                 "original_image_path": original_path
             }
+            return visual_results
 
         except Exception as e:
             logging.error(f"Visual analysis failed: {e}")
@@ -390,20 +326,11 @@ class AIProcessor:
         try:
             vector_store = self.knowledge_base_cache.get("vector_store")
             if not vector_store:
-                return "Clinical guidelines unavailable - knowledge base not loaded"
+                return "Knowledge base is not available."
             
             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
+            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])
             
         except Exception as e:
             logging.error(f"Guidelines query failed: {e}")
@@ -413,19 +340,16 @@ 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."""
+        """Generate final report using MedGemma GPU pipeline - EXACTLY like working reference."""
         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
+                patient_info, visual_results, guideline_context, image_pil, max_new_tokens
             )
             
-            if report and report.strip():
+            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")
                 return self._generate_fallback_report(patient_info, visual_results, guideline_context)
                 
         except Exception as e:
@@ -437,41 +361,56 @@ class AIProcessor:
     ) -> str:
         """Generate fallback report if MedGemma fails."""
         
-        report = f"""# Wound Analysis Report
+        report = f"""# 🩺 SmartHeal AI - Wound Analysis Report
 
-## Patient Information
+## 📋 Patient Information
 {patient_info}
 
-## Visual Analysis Results
+## 🔍 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
+## 📊 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
+## 📚 Clinical Guidelines Context
 {guideline_context[:1000]}{'...' if len(guideline_context) > 1000 else ''}
 
-## 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.
 
-## Recommendations
+### 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 Notes
+## ⚠️ 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
 
-## 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.
+## 🏥 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.
+
+---
+*Generated by SmartHeal AI - Advanced Wound Care Analysis System*
+*Report Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*
 """
         return report
 
@@ -510,7 +449,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."""
+        """Run full analysis pipeline - EXACTLY like working reference."""
         try:
             # Save image first
             saved_path = self.save_and_commit_image(image_pil)
@@ -520,18 +459,11 @@ 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
-            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"
+            # Process questionnaire data - EXACTLY like working reference
+            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
-            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 "
-                
+            # Query guidelines - EXACTLY like working reference
+            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")
             
@@ -594,7 +526,7 @@ This automated analysis is provided for informational purposes only and does not
         
         try:
             # Age assessment
-            age = questionnaire_data.get('patient_age', 0)
+            age = questionnaire_data.get('age', 0)
             if isinstance(age, str):
                 try:
                     age = int(age)
@@ -722,4 +654,127 @@ This automated analysis is provided for informational purposes only and does not
             recommendations.append("Use high-absorption dressings")
             recommendations.append("More frequent dressing changes may be needed")
         
-        return recommendations
+        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}"