|
import gradio as gr |
|
from transformers import ImageClassificationPipeline, PerceiverForImageClassificationConvProcessing, PerceiverFeatureExtractor |
|
import torch |
|
|
|
torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg') |
|
torch.hub.download_url_to_file('https://storage.googleapis.com/perceiver_io/dalmation.jpg', 'dog.jpg') |
|
|
|
feature_extractor = PerceiverFeatureExtractor.from_pretrained("deepmind/vision-perceiver-conv") |
|
model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv") |
|
|
|
image_pipe = ImageClassificationPipeline(model=model, feature_extractor=feature_extractor) |
|
|
|
def classify_image(image): |
|
results = image_pipe(image) |
|
|
|
output = {} |
|
for prediction in results: |
|
predicted_label = prediction['label'] |
|
score = prediction['score'] |
|
output[predicted_label] = score |
|
return output |
|
|
|
image = gr.inputs.Image(type="pil") |
|
label = gr.outputs.Label(num_top_classes=5) |
|
examples = [["cats.jpg"], ["dog.jpg"]] |
|
|
|
gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples, enable_queue=True).launch(debug=True) |