Update app.py

app.py

@@ -15,67 +15,84 @@ except Exception as e:
 # 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
+MODEL_NAME = "mistralai/mistral-small-24b-instruct-2501:free"
 
+# Define input shape
 input_shape = (224, 224, 3)
 
 def preprocess_image(image, target_size):
-    # ... (keep your existing preprocessing code) ...
+    """Preprocess the input image for the model."""
+    if image is None:
+        raise ValueError("No image provided")
+    image = image.convert("RGB")  # Ensure RGB format
+    image = image.resize(target_size)
+    image_array = np.array(image)
+    image_array = image_array / 255.0  # Normalize
+    return image_array
 
 def get_medical_guidelines(wound_type):
-    """
-    Fetch medical guidelines using OpenRouter.ai's Mistral model.
-    """
+    """Fetch medical guidelines using OpenRouter.ai's API."""
     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
+        "HTTP-Referer": "https://huggingface.co/spaces/MahatirTusher/Wound_Treatment",
+        "X-Title": "Wound_Treatment"
     }
     
-    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.
-    """
+    prompt = f"""As a medical professional, provide detailed guidelines for treating a {wound_type} wound. 
+    Include first aid steps, precautions, and when to seek professional help."""
     
     data = {
         "model": MODEL_NAME,
         "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.7  # Adjust for creativity vs. precision
+        "temperature": 0.7
     }
     
     try:
         response = requests.post(OPENROUTER_API_URL, headers=headers, json=data)
         response.raise_for_status()
-        result = response.json()
-        return result["choices"][0]["message"]["content"]
+        return response.json()["choices"][0]["message"]["content"]
     except Exception as e:
         return f"Error fetching guidelines: {str(e)}"
 
 def predict(image):
+    """Main prediction function."""
     try:
-        # ... (keep your existing preprocessing and prediction code) ...
+        # Preprocess image
+        input_data = preprocess_image(image, (input_shape[0], input_shape[1]))
+        input_data = np.expand_dims(input_data, axis=0)
 
-        # After getting `predicted_class`:
+        # Load class labels
+        with open('./classes.txt', 'r') as file:
+            class_labels = file.read().splitlines()
+
+        if len(class_labels) != model.output_shape[-1]:
+            raise ValueError("Class labels mismatch with model output")
+
+        # Make prediction
+        predictions = model.predict(input_data)
+        results = {class_labels[i]: float(predictions[0][i]) for i in range(len(class_labels))}
+        predicted_class = class_labels[np.argmax(predictions)]
+
+        # Get medical guidelines
         guidelines = get_medical_guidelines(predicted_class)
-        
-        return {
-            "predictions": results,  # Your existing classification results
-            "treatment_guidelines": guidelines  # From OpenRouter.ai
-        }
-    
+
+        return {"predictions": results, "treatment_guidelines": guidelines}
+
     except Exception as e:
         return {"error": str(e)}
 
-# Update Gradio interface to show both outputs
+# Gradio Interface
 iface = gr.Interface(
     fn=predict, 
     inputs=gr.Image(type="pil"), 
     outputs=[
-        gr.Label(num_top_classes=18, label="Classification Results"),
+        gr.Label(num_top_classes=3, label="Classification Results"),
         gr.Textbox(label="Medical Guidelines", lines=5)
     ],
-    live=True
+    live=True,
+    title="Wound Classification & Treatment Advisor",
+    description="Upload a wound image for classification and medical guidelines."
 )
 
 iface.launch(server_name="0.0.0.0", server_port=7860)