Update app.py

app.py

@@ -17,17 +17,16 @@ OPENROUTER_API_KEY = "sk-or-v1-cf4abd8adde58255d49e31d05fbe3f87d2bbfcdb50eb1dbef
 OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
 MODEL_NAME = "mistralai/mistral-small-24b-instruct-2501:free"
 
-# Define input shape
 input_shape = (224, 224, 3)
 
 def preprocess_image(image, target_size):
     """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.convert("RGB")
     image = image.resize(target_size)
     image_array = np.array(image)
-    image_array = image_array / 255.0  # Normalize
+    image_array = image_array / 255.0
     return image_array
 
 def get_medical_guidelines(wound_type):
@@ -36,7 +35,7 @@ def get_medical_guidelines(wound_type):
         "Authorization": f"Bearer {OPENROUTER_API_KEY}",
         "Content-Type": "application/json",
         "HTTP-Referer": "https://huggingface.co/spaces/MahatirTusher/Wound_Treatment",
-        "X-Title": "Wound_Treatment"
+        "X-Title": "Wound Treatment Advisor"
     }
     
     prompt = f"""As a medical professional, provide detailed guidelines for treating a {wound_type} wound. 
@@ -63,7 +62,7 @@ def predict(image):
         input_data = np.expand_dims(input_data, axis=0)
 
         # Load class labels
-        with open('./classes.txt', 'r') as file:
+        with open('classes.txt', 'r') as file:
             class_labels = file.read().splitlines()
 
         if len(class_labels) != model.output_shape[-1]:
@@ -72,15 +71,15 @@ def predict(image):
         # 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)]
+        predicted_class = max(results, key=results.get)
 
         # Get medical guidelines
         guidelines = get_medical_guidelines(predicted_class)
 
-        return {"predictions": results, "treatment_guidelines": guidelines}
+        return results, guidelines  # Return as tuple instead of dict
 
     except Exception as e:
-        return {"error": str(e)}
+        return {f"Error: {str(e)}"}, ""
 
 # Gradio Interface
 iface = gr.Interface(