Update app.py

app.py

@@ -1,25 +1,23 @@
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # Suppress TensorFlow warnings
 import gradio as gr
 import tensorflow as tf
+from tensorflow.keras.models import load_model
 import numpy as np
 from PIL import Image
 import requests
 import json
-import os
 
-# Load the model with enhanced error handling
+# ================== MODEL LOADING ==================
 try:
     model = load_model('wound_classifier_model_googlenet.h5')
     print("✅ Model loaded successfully")
-    print(f"ℹ️ Model expects {model.output_shape[-1]} output classes")  # Should be 18
+    print(f"ℹ️ Input shape: {model.input_shape}")
+    print(f"ℹ️ Output shape: {model.output_shape}")
 except Exception as e:
     raise RuntimeError(f"❌ Model loading failed: {str(e)}")
 
-# OpenRouter configuration
-OPENROUTER_API_KEY = os.getenv("OPENROUTER_KEY", "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")
-OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
-MODEL_NAME = "mistralai/mistral-7b-instruct"
-
-# Class labels from your classes.txt
+# ================== CLASS LABELS ==================
 CLASS_LABELS = [
     "Abrasions", "Bruises", "Burns", "Cut", "Diabetic Wounds", "Gingivitis",
     "Surgical Wounds", "Venous Wounds", "athlete foot", "cellulitis",
@@ -27,36 +25,46 @@ CLASS_LABELS = [
     "ringworm", "shingles", "tooth discoloration", "ulcer"
 ]
 
-# Verify class labels match model output
+# Verify model compatibility
 assert len(CLASS_LABELS) == model.output_shape[-1], \
-       f"Class labels mismatch: {len(CLASS_LABELS)} vs {model.output_shape[-1]}"
+    f"Class mismatch: Model expects {model.output_shape[-1]} classes, found {len(CLASS_LABELS)}"
+
+# ================== OPENROUTER CONFIG ==================
+OPENROUTER_API_KEY = "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"
+OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
+MODEL_NAME = "mistralai/mistral-7b-instruct"
 
+# ================== IMAGE PROCESSING ==================
 def preprocess_image(image, target_size=(224, 224)):
-    """Enhanced image preprocessing with validation"""
-    if not image:
-        raise ValueError("🖼️ No image provided")
-    
+    """Process and validate input images"""
     try:
+        if not image:
+            raise ValueError("🚨 No image provided")
+        
         image = image.convert("RGB").resize(target_size)
-        return np.array(image) / 255.0
+        array = np.array(image) / 255.0
+        print(f"🖼️ Image processed: Shape {array.shape}")
+        return array
     except Exception as e:
         raise RuntimeError(f"🖼️ Image processing failed: {str(e)}")
 
+# ================== MEDICAL GUIDELINES ==================
 def get_medical_guidelines(wound_type):
-    """Robust API handler with better error reporting"""
+    """Fetch treatment guidelines from OpenRouter API"""
     headers = {
         "Authorization": f"Bearer {OPENROUTER_API_KEY}",
         "Content-Type": "application/json",
         "HTTP-Referer": "https://huggingface.co/spaces/MahatirTusher/Wound_Treatment"
     }
     
-    prompt = f"""As a medical professional, provide concise guidelines for {wound_type}:
+    prompt = f"""As a medical expert, provide treatment guidelines for {wound_type}:
     - First aid steps
     - Precautions
-    - When to seek help
-    Avoid markdown, use simple language."""
+    - When to seek professional help
+    Use clear, simple language without markdown."""
     
     try:
+        print(f"📡 Sending API request for {wound_type}...")
         response = requests.post(
             OPENROUTER_API_URL,
             headers=headers,
@@ -65,24 +73,24 @@ def get_medical_guidelines(wound_type):
                 "messages": [{"role": "user", "content": prompt}],
                 "temperature": 0.5
             },
-            timeout=15
+            timeout=20
         )
         
         response.raise_for_status()
         result = response.json()
+        print("🔧 Raw API response:", json.dumps(result, indent=2))
         
         if not result.get("choices"):
-            return f"⚠️ API Error: Unexpected response format"
+            return "⚠️ API response format unexpected"
             
         return result["choices"][0]["message"]["content"]
         
-    except requests.exceptions.RequestException as e:
-        return f"🔌 Connection Error: {str(e)}"
     except Exception as e:
-        return f"⚠️ Processing Error: {str(e)}"
+        return f"⚠️ Guidelines unavailable: {str(e)}"
 
+# ================== MAIN PREDICTION ==================
 def predict(image):
-    """Main prediction pipeline with validation"""
+    """Complete prediction pipeline"""
     try:
         # Preprocess image
         processed_img = preprocess_image(image)
@@ -97,9 +105,9 @@ def predict(image):
             CLASS_LABELS[i]: float(predictions[i])
             for i in sorted_indices[:3]  # Top 3 predictions
         }
-        
-        # Get guidelines for top prediction
         top_class = CLASS_LABELS[sorted_indices[0]]
+        
+        # Get guidelines
         guidelines = get_medical_guidelines(top_class)
         
         return results, guidelines
@@ -107,18 +115,21 @@ def predict(image):
     except Exception as e:
         return {f"🚨 Error": str(e)}, ""
 
-# Gradio interface configuration
+# ================== GRADIO INTERFACE ==================
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload Wound Image"),
     outputs=[
         gr.Label(label="Top Predictions", num_top_classes=3),
-        gr.Textbox(label="Treatment Guidelines", lines=5)
+        gr.Textbox(label="Treatment Guidelines", lines=8)
     ],
-    title="Advanced Wound Classification System",
-    description="Identifies 18 wound types and provides treatment guidelines",
-    examples=["./example_abrasion.jpg", "./example_burn.jpg"] if os.path.exists("example_abrasion.jpg") else None,
-    allow_flagging="never"
+    title="AI Wound Classification System",
+    description="Identifies 18 wound types and provides treatment recommendations",
+    allow_flagging="never",
+    examples=[
+        f for f in ["abrasion.jpg", "burn.jpg"] 
+        if os.path.exists(f)
+    ]
 )
 
 if __name__ == "__main__":