Make adjustments for fastai learner

app.py

@@ -5,6 +5,7 @@ import gradio as gr
 
 repo_id = "artificeresearch/spiritvision"
 learner = from_pretrained_fastai(repo_id)
+labels = learner.dls.vocab
 
 
 def predict_fn(img):
@@ -13,9 +14,11 @@ def predict_fn(img):
     :return: prediction and probabilities
     """
     img = img.convert('RGB')
-    prediction, _, probs = learner.predict(img)
     # print(f'{max(100 * probs):.2f}% {prediction} - {img}')
-    return f'{max(100 * probs):.2f}% {prediction} - {img}'
+    pred, pred_idx, probs = learner.predict(img)
+    return {labels[i]: float(probs[i]) for i in range(len(labels))}
 
 
-gr.Interface(predict_fn, gr.inputs.Image(type='pil'), outputs='label').launch()
+gr.Interface(predict_fn,
+             gr.inputs.Image(type='pil', shape=(512, 512)),
+             outputs=gr.outputs.Label(num_top_classes=3)).launch()