Update app.py

app.py

@@ -1,130 +1,125 @@
 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
+# Load the model with enhanced error handling
 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
 except Exception as e:
-    raise RuntimeError(f"❌ Model loading failed: {e}")
+    raise RuntimeError(f"❌ Model loading failed: {str(e)}")
 
-# OpenRouter.ai Configuration
-OPENROUTER_API_KEY = "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"
+# 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"  # Updated model name
+MODEL_NAME = "mistralai/mistral-7b-instruct"
 
-input_shape = (224, 224, 3)
+# Class labels from your classes.txt
+CLASS_LABELS = [
+    "Abrasions", "Bruises", "Burns", "Cut", "Diabetic Wounds", "Gingivitis",
+    "Surgical Wounds", "Venous Wounds", "athlete foot", "cellulitis",
+    "chickenpox", "cutaneous larva migrans", "impetigo", "nail fungus",
+    "ringworm", "shingles", "tooth discoloration", "ulcer"
+]
 
-def preprocess_image(image, target_size):
-    """Preprocess the input image for the model."""
+# Verify class labels match model output
+assert len(CLASS_LABELS) == model.output_shape[-1], \
+       f"Class labels mismatch: {len(CLASS_LABELS)} vs {model.output_shape[-1]}"
+
+def preprocess_image(image, target_size=(224, 224)):
+    """Enhanced image preprocessing with validation"""
+    if not image:
+        raise ValueError("🖼️ No image provided")
+    
     try:
-        if image is None:
-            raise ValueError("No image provided")
-        image = image.convert("RGB")
-        image = image.resize(target_size)
+        image = image.convert("RGB").resize(target_size)
         return np.array(image) / 255.0
     except Exception as e:
-        print(f"⚠️ Image preprocessing error: {e}")
-        raise
+        raise RuntimeError(f"🖼️ Image processing failed: {str(e)}")
 
 def get_medical_guidelines(wound_type):
-    """Fetch medical guidelines using OpenRouter.ai's API."""
+    """Robust API handler with better error reporting"""
     headers = {
         "Authorization": f"Bearer {OPENROUTER_API_KEY}",
         "Content-Type": "application/json",
-        "HTTP-Referer": "https://huggingface.co/spaces/MahatirTusher/Wound_Treatment",
-        "X-Title": "Wound Treatment Advisor"
+        "HTTP-Referer": "https://huggingface.co/spaces/MahatirTusher/Wound_Treatment"
     }
     
-    prompt = f"""As a medical professional, provide detailed guidelines for treating a {wound_type} wound. 
-    Include: 
-    1. First aid steps
-    2. Precautions
-    3. When to seek professional help
-    Output in markdown with clear sections."""
-    
-    data = {
-        "model": MODEL_NAME,
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.5
-    }
+    prompt = f"""As a medical professional, provide concise guidelines for {wound_type}:
+    - First aid steps
+    - Precautions
+    - When to seek help
+    Avoid markdown, use simple language."""
     
     try:
-        print(f"🚀 Sending request to OpenRouter API for {wound_type}...")
-        response = requests.post(OPENROUTER_API_URL, headers=headers, json=data, timeout=10)
-        response.raise_for_status()
+        response = requests.post(
+            OPENROUTER_API_URL,
+            headers=headers,
+            json={
+                "model": MODEL_NAME,
+                "messages": [{"role": "user", "content": prompt}],
+                "temperature": 0.5
+            },
+            timeout=15
+        )
         
-        response_json = response.json()
-        print("🔧 Raw API response:", json.dumps(response_json, indent=2))
+        response.raise_for_status()
+        result = response.json()
         
-        if "choices" not in response_json:
-            return "⚠️ API response format unexpected. Please check logs."
+        if not result.get("choices"):
+            return f"⚠️ API Error: Unexpected response format"
             
-        return response_json["choices"][0]["message"]["content"]
+        return result["choices"][0]["message"]["content"]
         
-    except requests.exceptions.HTTPError as e:
-        print(f"❌ HTTP Error: {e.response.status_code} - {e.response.text}")
-        return f"API Error: {e.response.status_code} - Check console for details"
+    except requests.exceptions.RequestException as e:
+        return f"🔌 Connection Error: {str(e)}"
     except Exception as e:
-        print(f"⚠️ General API error: {str(e)}")
-        return f"Error: {str(e)}"
+        return f"⚠️ Processing Error: {str(e)}"
 
 def predict(image):
-    """Main prediction function."""
+    """Main prediction pipeline with validation"""
     try:
         # Preprocess image
-        input_data = preprocess_image(image, (input_shape[0], input_shape[1]))
-        input_data = np.expand_dims(input_data, axis=0)
-        print("🖼️ Image preprocessed successfully")
-
-        # Load class labels
-        try:
-            with open('classes.txt', 'r') as file:
-                class_labels = file.read().splitlines()
-            print("📋 Class labels loaded:", class_labels)
-        except Exception as e:
-            raise RuntimeError(f"Class labels loading failed: {e}")
-
-        # Verify model compatibility
-        if len(class_labels) != model.output_shape[-1]:
-            raise ValueError(f"Model expects {model.output_shape[-1]} classes, found {len(class_labels)}")
-
+        processed_img = preprocess_image(image)
+        input_tensor = np.expand_dims(processed_img, axis=0)
+        
         # Make prediction
-        predictions = model.predict(input_data)
-        print("📊 Raw predictions:", predictions)
+        predictions = model.predict(input_tensor)[0]
+        sorted_indices = np.argsort(predictions)[::-1]  # Descending order
+        
+        # Format results
+        results = {
+            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]]
+        guidelines = get_medical_guidelines(top_class)
         
-        results = {class_labels[i]: float(predictions[0][i]) 
-                 for i in range(len(class_labels))}
-        predicted_class = max(results, key=results.get)
-        print(f"🏆 Predicted class: {predicted_class}")
-
-        # Get medical guidelines
-        guidelines = get_medical_guidelines(predicted_class)
-        print("📜 Guidelines generated successfully")
-
         return results, guidelines
-
+        
     except Exception as e:
-        print(f"🔥 Critical error in prediction: {str(e)}")
-        return {"Error": str(e)}, ""
+        return {f"🚨 Error": str(e)}, ""
 
-# Gradio Interface
+# Gradio interface configuration
 iface = gr.Interface(
-    fn=predict, 
-    inputs=gr.Image(type="pil", label="Upload Wound Image"), 
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload Wound Image"),
     outputs=[
-        gr.Label(num_top_classes=3, label="Classification Results"),
-        gr.Markdown(label="Medical Guidelines")
+        gr.Label(label="Top Predictions", num_top_classes=3),
+        gr.Textbox(label="Treatment Guidelines", lines=5)
     ],
-    live=False,
-    title="Wound Classification & Treatment Advisor",
-    description="Upload a wound image for AI-powered classification and treatment guidelines.",
+    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"
 )
 
-iface.launch(server_name="0.0.0.0", server_port=7860)
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)