Update app.py

app.py

@@ -12,6 +12,8 @@ processor = AutoImageProcessor.from_pretrained(model_name)
 # Get the label names from the model's configuration
 labels = model.config.id2label
 
+print("Labels:", labels)  # Debugging statement to check the labels
+
 # Define the prediction function (with preprocessing)
 def predict_image(image):
     """
@@ -25,32 +27,40 @@ def predict_image(image):
     """
     # Preprocessing steps:
     image = np.array(image)
-    image = np.repeat(image[:, :, np.newaxis], 3, axis=2)  # Convert grayscale to RGB
+    if image.ndim == 2:  # Convert grayscale to RGB if needed
+        image = np.repeat(image[:, :, np.newaxis], 3, axis=2)
+
+    # Resize image if necessary (optional step)
+    if image.shape[0] != 224 or image.shape[1] != 224:
+        image = np.array(Image.fromarray(image).resize((224, 224)))
 
     # Model inference:
     inputs = processor(images=image, return_tensors="pt").to(model.device)
     with torch.no_grad():
         logits = model(**inputs).logits
 
+    print(f"logits shape: {logits.shape}")  # Debugging statement to check shape
+    print(f"logits: {logits}")  # Debugging statement to check content
+
     predicted_label_id = logits.argmax(-1).item()
-    predicted_label = labels[str(predicted_label_id)]
+    predicted_label = labels[predicted_label_id]
 
     # Calculate probabilities using softmax
     probabilities = torch.nn.functional.softmax(logits, dim=-1)
-    confidences = {labels[str(i)]: float(probabilities[0][i]) for i in range(len(labels))}
+    confidences = {labels[i]: float(probabilities[0][i]) for i in range(len(labels))}
 
     return predicted_label, confidences
 
-# Create the Gradio interface (updated)
+# Create the Gradio interface
 iface = gr.Interface(
     fn=predict_image,
-    inputs=gr.Image(type="pil", label="Upload MRI Image"),  # Use gr.Image directly
+    inputs=gr.Image(type="pil", label="Upload MRI Image"),
     outputs=[
         gr.Label(label="Prediction"),
         gr.JSON(label="Confidence Scores")
     ],
     title="Alzheimer's Disease MRI Image Classifier",
-    description="Upload a MRI image to predict the stage of Alzheimer's disease."
+    description="Upload an MRI image to predict the stage of Alzheimer's disease."
 )
 
 iface.launch()