Julien Simon commited on
Commit
7539610
1 Parent(s): a83a277

Predict with all models

Browse files
Files changed (1) hide show
  1. app.py +14 -13
app.py CHANGED
@@ -14,34 +14,35 @@ model_names = [
14
  ]
15
 
16
 
17
- def process(image_file, top_k, model_name):
18
- p = pipeline("image-classification", model=model_name)
19
- pred = p(image_file)
20
- return {x["label"]: x["score"] for x in pred[:top_k]}
 
 
 
21
 
22
 
23
  # Inputs
24
  image = gr.Image(type="filepath", label="Upload an image")
25
  top_k = gr.Slider(minimum=1, maximum=5, step=1, value=5, label="Top k classes")
26
- model_selection = gr.Dropdown(
27
- model_names, value="google/vit-base-patch16-224", label="Pick a model"
28
- )
29
 
30
  # Output
31
- labels = gr.Label()
32
 
33
  description = "This Space lets you quickly compare the most popular image classifiers available on the hub, including the recent NAT and DINAT models. All of them have been fine-tuned on the ImageNet-1k dataset. Anecdotally, the three sample images have been generated with a Stable Diffusion model :)"
34
 
35
  iface = gr.Interface(
36
  theme="huggingface",
37
  description=description,
 
38
  fn=process,
39
- inputs=[image, top_k, model_selection],
40
- outputs=[labels],
41
  examples=[
42
- ["bike.jpg", 5, "google/vit-base-patch16-224"],
43
- ["car.jpg", 5, "microsoft/swin-base-patch4-window7-224"],
44
- ["food.jpg", 5, "facebook/convnext-base-224"],
45
  ],
46
  allow_flagging="never",
47
  )
 
14
  ]
15
 
16
 
17
+ def process(image_file, top_k):
18
+ labels = []
19
+ for m in model_names:
20
+ p = pipeline("image-classification", model=m)
21
+ pred = p(image_file)
22
+ labels.append({x["label"]: x["score"] for x in pred[:top_k]})
23
+ return labels
24
 
25
 
26
  # Inputs
27
  image = gr.Image(type="filepath", label="Upload an image")
28
  top_k = gr.Slider(minimum=1, maximum=5, step=1, value=5, label="Top k classes")
 
 
 
29
 
30
  # Output
31
+ labels = [gr.Label(label=m) for m in model_names]
32
 
33
  description = "This Space lets you quickly compare the most popular image classifiers available on the hub, including the recent NAT and DINAT models. All of them have been fine-tuned on the ImageNet-1k dataset. Anecdotally, the three sample images have been generated with a Stable Diffusion model :)"
34
 
35
  iface = gr.Interface(
36
  theme="huggingface",
37
  description=description,
38
+ layout="horizontal",
39
  fn=process,
40
+ inputs=[image, top_k],
41
+ outputs=labels,
42
  examples=[
43
+ ["bike.jpg", 5],
44
+ ["car.jpg", 5],
45
+ ["food.jpg", 5],
46
  ],
47
  allow_flagging="never",
48
  )