PaulMest commited on
Commit
8d98eff
1 Parent(s): a0b05da

Introducing Catan categorization

Browse files
app.py CHANGED
@@ -4,24 +4,35 @@ import timm
4
  import dill
5
  import os
6
 
7
- learn = load_learner('./models/catan-model-paperspace-5.pkl', pickle_module=dill)
8
- # learn = load_learner('catan-model.pkl', pickle_module=dill)
9
-
10
- # categories = learn.dls.vocab
11
- categories = ('Not Catan', 'Catan')
12
 
 
13
 
14
  def classify_image(img):
15
- pred, idx, probs = learn.predict(img)
16
- return dict(zip(categories, map(float, probs)))
 
 
 
 
 
 
 
 
 
 
 
17
 
18
 
19
  # Cell
20
  image = gr.inputs.Image(shape=(192, 192))
21
  label = gr.outputs.Label()
 
 
22
  examples_dir_path = './examples/'
23
  examples = [(examples_dir_path + filename) for filename in os.listdir(examples_dir_path) if filename[:1] != '.']
24
 
25
  # Cell
26
- intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
27
  intf.launch()
 
4
  import dill
5
  import os
6
 
7
+ is_catan_learn = load_learner('./models/catan-model-paperspace-2022-11-29-03-28-12.pkl', pickle_module=dill)
8
+ catan_category_learn = load_learner('./models/categories-of-catan-3.pkl', pickle_module=dill)
 
 
 
9
 
10
+ # learn = load_learner('catan-model.pkl', pickle_module=dill)
11
 
12
  def classify_image(img):
13
+ pred, pred_idx, probs = is_catan_learn.predict(img)
14
+ if float(probs[1]) < 0.2:
15
+ # categories = learn.dls.vocab
16
+ categories = ('Not Catan', 'Catan')
17
+ 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.'
18
+ details = dict(zip(categories, map(float, probs)))
19
+ else:
20
+ pred, pred_idx, probs = catan_category_learn.predict(img)
21
+ message = f'Prediction: *{pred}*; Probability: *{probs[pred_idx]:.04f}%*'
22
+ categories = catan_category_learn.dls.vocab
23
+ details = dict(zip(categories, map(float, probs)))
24
+
25
+ return details, message
26
 
27
 
28
  # Cell
29
  image = gr.inputs.Image(shape=(192, 192))
30
  label = gr.outputs.Label()
31
+ description = gr.Markdown()
32
+
33
  examples_dir_path = './examples/'
34
  examples = [(examples_dir_path + filename) for filename in os.listdir(examples_dir_path) if filename[:1] != '.']
35
 
36
  # Cell
37
+ intf = gr.Interface(fn=classify_image, inputs=image, outputs=[label, description], examples=examples)
38
  intf.launch()
models/catan-model-paperspace-2022-11-29-03-28-12.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:55809db865c9df20cfe4bb0a9f280763475bbc274de21fc6c505bb3687cdf9ad
3
+ size 87464943
models/categories-of-catan-3.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2975e66420b60a4b32e4948e699dd73a4ad21c0d7c09f74d6f5c46763e5538b0
3
+ size 46968545