Update app.py

app.py

@@ -1,35 +1,23 @@
-import streamlit as st
-from PIL import Image
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
-import torch
+from transformers import pipeline
+import gradio as gr
+def alz_mri_classification(image):
+    classifier = pipeline("image-classification", model="dewifaj/alzheimer_classification")
+    result = classifier(image)
+    # extract the highest score
+    prediction = result[0]
+    score = prediction['score']
+    label = prediction['label']
+    return {"score": score, "label": label}
+    
+example_image_paths = ["Very_Mild_Demented.png", 
+                       "Mild_Demented.png", 
+                       "Moderate_Demented.png", 
+                       "Non_Demented.png"]
 
-# Load the model and feature extractor
-model_name = "dewifaj/resnet18_alzheimer_classifier"  
-model = AutoModelForImageClassification.from_pretrained(model_name)
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-
-# Define the label mapping
-label_mapping = model.config.id2label
-
-def predict(image):
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        predicted_class_idx = logits.argmax(-1).item()
-    return label_mapping[predicted_class_idx]
-
-# Streamlit app
-st.title("Alzheimer Image Classification")
-st.write("Upload an image to classify the stage of Alzheimer's disease.")
-
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
-    st.write("")
-    st.write("Classifying...")
-
-    label = predict(image)
-    st.write(f"The model predicts: **{label}**")
+image_input = gr.Image(type="pil", label="Upload Image")
+iface = gr.Interface(fn=alz_mri_classification, 
+                     inputs=image_input, 
+                     outputs="json", 
+                     examples=example_image_paths,
+                     title="Alzheimer Recognition from MRI")
+iface.launch()