rubensmau commited on
Commit
e05fcd2
1 Parent(s): 2ca4905

novo arquivo

Browse files
Files changed (1) hide show
  1. app.py +25 -0
app.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import ImageClassificationPipeline, PerceiverForImageClassificationConvProcessing, PerceiverFeatureExtractor
3
+ import torch
4
+ torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
5
+ torch.hub.download_url_to_file('https://storage.googleapis.com/perceiver_io/dalmation.jpg', 'dog.jpg')
6
+
7
+ feature_extractor = PerceiverFeatureExtractor.from_pretrained("deepmind/vision-perceiver-conv")
8
+ model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
9
+ image_pipe = ImageClassificationPipeline(model=model, feature_extractor=feature_extractor)
10
+ def classify_image(image):
11
+ results = image_pipe(image)
12
+ # convert to format Gradio expects
13
+ output = {}
14
+ for prediction in results:
15
+ predicted_label = prediction['label']
16
+ score = prediction['score']
17
+ output[predicted_label] = score
18
+ return output
19
+ image = gr.inputs.Image(type="pil")
20
+ label = gr.outputs.Label(num_top_classes=5)
21
+ examples = [["cats.jpg"], ["dog.jpg"]]
22
+ title = "Interactive demo: Perceiver for image classification"
23
+ description = "Demo for classifying images with Perceiver IO. To use it, simply upload an image or use the example images below and click 'submit' to let the model predict the 5 most probable ImageNet classes. Results will show up in a few seconds."
24
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2107.14795'>Perceiver IO: A General Architecture for Structured Inputs & Outputs</a> | <a href='https://deepmind.com/blog/article/building-architectures-that-can-handle-the-worlds-data/'>Official blog</a></p>"
25
+ gr.Interface(fn=classify_image, inputs=image, outputs=label, title=title, description=description, examples=examples, enable_queue=True).launch(debug=True)