Introducing Catan categorization

app.py

@@ -4,24 +4,35 @@ import timm
 import dill
 import os
 
-learn = load_learner('./models/catan-model-paperspace-5.pkl', pickle_module=dill)
-# learn = load_learner('catan-model.pkl', pickle_module=dill)
-
-# categories = learn.dls.vocab
-categories = ('Not Catan', 'Catan')
+is_catan_learn = load_learner('./models/catan-model-paperspace-2022-11-29-03-28-12.pkl', pickle_module=dill)
+catan_category_learn = load_learner('./models/categories-of-catan-3.pkl', pickle_module=dill)
 
+# learn = load_learner('catan-model.pkl', pickle_module=dill)
 
 def classify_image(img):
-    pred, idx, probs = learn.predict(img)
-    return dict(zip(categories, map(float, probs)))
+    pred, pred_idx, probs = is_catan_learn.predict(img)
+    if float(probs[1]) < 0.2:
+        # categories = learn.dls.vocab
+        categories = ('Not Catan', 'Catan')
+        message = f'Did not detect Catan in this upload: *{probs[1]:.4f}%*. Choose another photo with Catan in it and we will categorize what kind of Catan we find.'
+        details = dict(zip(categories, map(float, probs)))
+    else:
+        pred, pred_idx, probs = catan_category_learn.predict(img)
+        message = f'Prediction: *{pred}*; Probability: *{probs[pred_idx]:.04f}%*'
+        categories = catan_category_learn.dls.vocab
+        details = dict(zip(categories, map(float, probs)))
+
+    return details, message
 
 
 # Cell
 image = gr.inputs.Image(shape=(192, 192))
 label = gr.outputs.Label()
+description = gr.Markdown()
+
 examples_dir_path = './examples/'
 examples = [(examples_dir_path + filename) for filename in os.listdir(examples_dir_path) if filename[:1] != '.']
 
 # Cell
-intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
+intf = gr.Interface(fn=classify_image, inputs=image, outputs=[label, description], examples=examples)
 intf.launch()

models/catan-model-paperspace-2022-11-29-03-28-12.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55809db865c9df20cfe4bb0a9f280763475bbc274de21fc6c505bb3687cdf9ad
+size 87464943

models/categories-of-catan-3.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2975e66420b60a4b32e4948e699dd73a4ad21c0d7c09f74d6f5c46763e5538b0
+size 46968545