Update app.py

app.py

@@ -1,17 +1,13 @@
-# import gradio as gr
-
-# gr.load("models/devadethanr/alz_model").launch()
-
-
 import gradio as gr
-from transformers import AutoModelForImageClassification
+from transformers import AutoModelForImageClassification, AutoImageProcessor
 import torch
 import numpy as np
+from PIL import Image
 
-
-# Load the model and image processor from the Hub
-model_name = "devadethanr/alz_model" 
+# Load the model from the Hub
+model_name = "devadethanr/alz_model"
 model = AutoModelForImageClassification.from_pretrained(model_name)
+processor = AutoImageProcessor.from_pretrained(model_name)
 
 # Get the label names from the model's configuration
 labels = model.config.id2label
@@ -27,28 +23,31 @@ def predict_image(image):
     Returns:
         The predicted label with its corresponding probability.
     """
+    # Preprocessing steps:
+    image = np.array(image)
+    image = np.repeat(image[:, :, np.newaxis], 3, axis=2)  # Convert grayscale to RGB
 
-    image = model.preprocess_image(image, return_tensors="pt").to(model.device)
+    # Model inference:
+    inputs = processor(images=image, return_tensors="pt").to(model.device)
     with torch.no_grad():
-        logits = model(**image).logits
+        logits = model(**inputs).logits
 
     predicted_label_id = logits.argmax(-1).item()
-    predicted_label = labels[predicted_label_id]
+    predicted_label = labels[str(predicted_label_id)]
 
     # Calculate probabilities using softmax
     probabilities = torch.nn.functional.softmax(logits, dim=-1)
-    confidences = {label: float(probabilities[0][i]) for i, label in enumerate(labels)}
+    confidences = {labels[str(i)]: float(probabilities[0][i]) for i in range(len(labels))}
 
     return predicted_label, confidences
 
-
-# Create the Gradio interface (same as before)
+# Create the Gradio interface (updated)
 iface = gr.Interface(
     fn=predict_image,
-    inputs=gr.inputs.Image(type="pil", label="Upload MRI Image"),
+    inputs=gr.Image(type="pil", label="Upload MRI Image"),  # Use gr.Image directly
     outputs=[
-        gr.outputs.Label(label="Prediction"),
-        gr.outputs.JSON(label="Confidence Scores")
+        gr.Label(label="Prediction"),
+        gr.JSON(label="Confidence Scores")
     ],
     title="Alzheimer's Disease MRI Image Classifier",
     description="Upload an MRI image to predict the stage of Alzheimer's disease."