Update app.py

app.py

@@ -1,14 +1,14 @@
 import gradio as gr
 from PIL import Image
 from transformers import ViTImageProcessor, ViTForImageClassification
-
+import torch
 
 
 processor = ViTImageProcessor.from_pretrained('Rageshhf/fine-tuned-model')
 
 id2label = {0: 'Mild_Demented', 1: 'Moderate_Demented', 2: 'Non_Demented', 3: 'Very_Mild_Demented'}
 label2id = {'Mild_Demented': 0, 'Moderate_Demented': 1, 'Non_Demented': 2, 'Very_Mild_Demented': 3}
-
+labels = ['Mild_Demented', 'Moderate_Demented', 'Non_Demented', 'Very_Mild_Demented']
 
 model = ViTForImageClassification.from_pretrained(
     'Rageshhf/fine-tuned-model',
@@ -23,15 +23,20 @@ description = """Trained to classify disease based on image data."""
 
 
 def predict(image):
-
     inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
+
     logits = outputs.logits
-    # model predicts one of the 1000 ImageNet classes
-    predicted_class_idx = logits.argmax(-1).item()
-    return(model.config.id2label[predicted_class_idx])
+    prediction = torch.nn.functional.softmax(logits, dim=1)
+    probabilities = prediction[0].tolist()
+
+    output = {}
+    for i, prob in enumerate(probabilities):
+        output[labels[i]] = prob
+
+    return output
 
-demo = gr.Interface(fn=predict, inputs="image", outputs=gr.Label(num_top_classes=2), title=title, examples=["examples/image_1.png", "examples/image_2.png", "examples/image_3.png"],
+demo = gr.Interface(fn=predict, inputs="image", outputs=gr.Label(num_top_classes=3), title=title, examples=["examples/image_1.png", "examples/image_2.png", "examples/image_3.png"],
     description=description,).launch()
 
 # demo.launch(debug=True)