Update app.py

app.py

@@ -8,11 +8,13 @@ repo_id = "MasleK/snails_snakes_slugs"
 
 learner = from_pretrained_fastai(repo_id)
 
+categories = learner.learn.dls.vocab
 
 def predict(image):
-    label, index, scores = learner.predict(image)
+    label, index, probs = learner.predict(image)
    
-    return {l: scores[i].item() for i,l in enumerate(learner.dls.vocab)}
+    return dict(zip(categories, map(float,probs)))
+
 title = "Snail, snake, slug Classifier"
 description = f"""<h1> Slug, snake or  snail<h1>
               <p>A  classifier trained on about 300 images. Created as a demo for Gradio and HuggingFace Spaces."""