dominguezdaniel commited on
Commit
9105495
1 Parent(s): b3f3e2f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -4,7 +4,8 @@ from transformers import pipeline
4
  def predict(image):
5
  model_id = "google/vit-base-patch16-224"
6
  classifier = pipeline("image-classification", model=model_id)
7
- return classifier(image)
 
8
 
9
  title = "Image Classifier"
10
  description = "A demo that recognizes and classifies images using the model from Hugging Face's 'google/vit-base-patch16-224'."
 
4
  def predict(image):
5
  model_id = "google/vit-base-patch16-224"
6
  classifier = pipeline("image-classification", model=model_id)
7
+ predictions = classifier(image)
8
+ return {prediction['label']: prediction['score'] for prediction in predictions}
9
 
10
  title = "Image Classifier"
11
  description = "A demo that recognizes and classifies images using the model from Hugging Face's 'google/vit-base-patch16-224'."