added categorization

app.py

@@ -18,11 +18,11 @@ learn = load_learner("model-export.pkl")
 
 labels = learn.dls.vocab
 
-
-def predict(img):
+categories = ('Dog', 'Cat')
+def classify_image(img):
     img = PILImage.create(img)
-    pred, pred_idx, probs = learn.predict(img)
-    return {labels[i]: float(probs[i]) for i in range(len(labels))}
+    pred,idx,probs = learn.predict(img)
+    return dict(zip(categories, map(float,probs)))
 
 
 title = "Pet Breed Classifier"
@@ -33,7 +33,7 @@ allowed_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.gif'}
 examples = [file for file in path.iterdir() if file.suffix.lower() in allowed_extensions]
 
 gr.Interface(
-    fn=predict,
+    fn=classify_image,
     inputs="image",
     outputs="label",
     title=title,