Spaces:

kolkhi
/

bearclassifier

Sleeping

George Sergia commited on 9 days ago

Commit

334e182

1 Parent(s): 5150d70

Fix categories

Files changed (1) hide show

app.py CHANGED Viewed

@@ -2,11 +2,14 @@ from fastai.vision.all import *
 from huggingface_hub import push_to_hub_fastai, from_pretrained_fastai
 import gradio as gr
-learner = from_pretrained_fastai("kolkhi/bears")
-categories = ("grizzly", "black", "white", "brown", "teddy")
 def classify_bear(img):
-    pred,idx,probs=learner.predict(img)
     return dict(zip(categories, map(float,probs)))

 from huggingface_hub import push_to_hub_fastai, from_pretrained_fastai
 import gradio as gr
+#learner = from_pretrained_fastai("kolkhi/bears")
+learner = load_learner("bears.pkl")
+categories = learner.dls.vocab
 def classify_bear(img):
+    img_new = PILImage.create(img)
+    img_new.resize((128,128))
+    pred,idx,probs=learner.predict(img_new)
     return dict(zip(categories, map(float,probs)))