dbetm commited on
Commit
85a7dd8
1 Parent(s): 47bfb2e

add model v2 - more painters supported

Browse files
Files changed (2) hide show
  1. app.py +17 -3
  2. model_v2.pkl +3 -0
app.py CHANGED
@@ -10,21 +10,35 @@ CATS_MAP = {
10
  "rembrandt": "Rembrandt",
11
  }
12
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
  # load pre-trained model
14
- model = load_learner("model.pkl")
15
 
16
  # get classes name in right order
17
- full_name_cats = [CATS_MAP[key_class] for key_class in model.dls.vocab]
18
 
19
 
20
  def classify_image(img) -> dict:
21
  category, idx, probs = model.predict(img)
 
22
  return dict(zip(full_name_cats, map(float, probs)))
23
 
24
 
25
  # Gradio control
26
  image = gr.inputs.Image(shape=(224, 224))
27
- label = gr.outputs.Label()
28
  examples = [
29
  f"images_examples/{filename}"
30
  for filename in ("mona_lisa.jpg", "starry_night.jpg", "le_reve.jpg")
 
10
  "rembrandt": "Rembrandt",
11
  }
12
 
13
+ CATS_MAP_V2 = {
14
+ "picasso": "Pablo Picasso",
15
+ "vanGogh": "Vincent van Gogh",
16
+ "dali": "Salvador Dalí",
17
+ "daVinci": "Leonard da Vinci",
18
+ "rembrandt": "Rembrandt",
19
+ "monet": "Claude Monet",
20
+ "caruso": "Santiago Caruso",
21
+ "renoir": "Pierre-Auguste Renoir",
22
+ "oKeeffe": "Georgia O’Keeffe",
23
+ "krasner": "Lee Krasner",
24
+ }
25
+
26
  # load pre-trained model
27
+ model = load_learner("model_v2.pkl")
28
 
29
  # get classes name in right order
30
+ full_name_cats = [CATS_MAP_V2[key_class] for key_class in model.dls.vocab]
31
 
32
 
33
  def classify_image(img) -> dict:
34
  category, idx, probs = model.predict(img)
35
+
36
  return dict(zip(full_name_cats, map(float, probs)))
37
 
38
 
39
  # Gradio control
40
  image = gr.inputs.Image(shape=(224, 224))
41
+ label = gr.outputs.Label(num_top_classes=4)
42
  examples = [
43
  f"images_examples/{filename}"
44
  for filename in ("mona_lisa.jpg", "starry_night.jpg", "le_reve.jpg")
model_v2.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:cb83103d6b4a409ecde5f5c0a17e7570beed3f8907821d91ba91b97ec4347d4f
3
+ size 46999859