Use a label

app.py

@@ -6,14 +6,27 @@ from fastai.learner import load_learner
 
 model = load_learner('./mushrooms.pkl')
 
+
+categories = [
+    "Agaricus",
+    "Amanita",
+    "Boletus",
+    "Cortinarius",
+    "Entoloma",
+    "Hygrocybe",
+    "Lactarius",
+    "Russula",
+    "Suillus",
+]
+
 def classify_image(image):
     prediction, _, probs = model.predict(image)
-    return f"{prediction}: {max(probs)}"
+    return dict(zip(categories, map(float, probs)))
 
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
-    outputs="text",
+    outputs=gr.Label(),
     title="Image Classifier",
     description="WHAT IS THE MUSHROOM?"
 )