Update app.py

app.py

@@ -3,6 +3,8 @@ import tensorflow as tf
 from tensorflow.keras.models import load_model
 import numpy as np
 from PIL import Image
+import requests
+import json
 
 # Load the model
 try:
@@ -10,59 +12,70 @@ try:
 except Exception as e:
     raise RuntimeError(f"Error loading model: {e}")
 
-# Define the input shape
+# OpenRouter.ai Configuration
+OPENROUTER_API_KEY = "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"
+OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
+MODEL_NAME = "mistralai/mistral-small-24b-instruct-2501:free"  # Mistral model via OpenRouter
+
 input_shape = (224, 224, 3)
 
 def preprocess_image(image, target_size):
+    # ... (keep your existing preprocessing code) ...
+
+def get_medical_guidelines(wound_type):
     """
-    Preprocess the input image for the model.
-    - Resize the image to the target size.
-    - Normalize pixel values to the range [0, 1].
+    Fetch medical guidelines using OpenRouter.ai's Mistral model.
     """
-    if image is None:
-        raise ValueError("No image provided")
-    image = image.convert("RGB")  # Ensure the image is in RGB mode
-    image = image.resize(target_size)
-    image_array = np.array(image)
-    image_array = image_array / 255.0  # Normalize the image
-    return image_array
+    headers = {
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+        "Content-Type": "application/json",
+        "HTTP-Referer": "https://your-huggingface-space-url.com",  # Optional
+        "X-Title": "Wound Classifier"  # Optional
+    }
+    
+    prompt = f"""
+    As a medical professional, provide detailed guidelines for treating a {wound_type} wound. 
+    Include steps for first aid, precautions, and when to seek professional help.
+    """
+    
+    data = {
+        "model": MODEL_NAME,
+        "messages": [{"role": "user", "content": prompt}],
+        "temperature": 0.7  # Adjust for creativity vs. precision
+    }
+    
+    try:
+        response = requests.post(OPENROUTER_API_URL, headers=headers, json=data)
+        response.raise_for_status()
+        result = response.json()
+        return result["choices"][0]["message"]["content"]
+    except Exception as e:
+        return f"Error fetching guidelines: {str(e)}"
 
 def predict(image):
-    """
-    Predict the class probabilities for the input image.
-    - Preprocess the image.
-    - Predict using the loaded model.
-    - Return results as a dictionary with class labels and probabilities.
-    """
     try:
-        # Preprocess the image
-        input_data = preprocess_image(image, (input_shape[0], input_shape[1]))
-        input_data = np.expand_dims(input_data, axis=0)  # Add batch dimension
+        # ... (keep your existing preprocessing and prediction code) ...
 
-        # Load class labels
-        try:
-            with open('./classes.txt', 'r') as file:
-                class_labels = file.read().splitlines()
-        except FileNotFoundError:
-            raise RuntimeError("Class labels file 'classes.txt' not found.")
+        # After getting `predicted_class`:
+        guidelines = get_medical_guidelines(predicted_class)
         
-        if len(class_labels) != model.output_shape[-1]:
-            raise ValueError("Mismatch between model output and class labels.")
-
-        # Predict probabilities
-        predictions = model.predict(input_data)
-        results = {class_labels[i]: float(predictions[0][i]) for i in range(len(class_labels))}
-        return results
+        return {
+            "predictions": results,  # Your existing classification results
+            "treatment_guidelines": guidelines  # From OpenRouter.ai
+        }
+    
     except Exception as e:
         return {"error": str(e)}
 
-# Create a Gradio interface
+# Update Gradio interface to show both outputs
 iface = gr.Interface(
     fn=predict, 
     inputs=gr.Image(type="pil"), 
-    outputs=gr.Label(num_top_classes=18),  # Adjust num_top_classes as needed
+    outputs=[
+        gr.Label(num_top_classes=18, label="Classification Results"),
+        gr.Textbox(label="Medical Guidelines", lines=5)
+    ],
     live=True
 )
 
-# Launch the Gradio interface
-iface.launch(server_name="0.0.0.0", server_port=7860)
+iface.launch(server_name="0.0.0.0", server_port=7860)