Update app.py

app.py

@@ -15,13 +15,21 @@ def predict(image):
     with torch.no_grad():
         outputs = model(**inputs)
     logits = outputs.logits
+    # Calculate the confidence values
+    softmax = torch.nn.functional.softmax(logits, dim=1)
+    confidences = softmax.squeeze().tolist()
     # Get the predicted class
-    predicted_class = logits.argmax(-1).item()
-    # You may need to adjust the following line based on your class labels
+    predicted_class_index = logits.argmax(-1).item()
     class_names = ["glioma", "meningioma", "notumor", "pituitary"]
-    return class_names[predicted_class]
+    predicted_class = class_names[predicted_class_index]
+    # Create a dictionary to return both the predicted class and the confidence values
+    result = {
+        "predicted_class": predicted_class,
+        "confidences": {class_names[i]: confidences[i] for i in range(len(class_names))}
+    }
+    return result
 
 # Set up the Gradio interface
 image_cp = gr.Image(type="pil", label='Brain')
-interface = gr.Interface(fn=predict, inputs=image_cp, outputs="text")
+interface = gr.Interface(fn=predict, inputs=image_cp, outputs="json")
 interface.launch()