| import gradio as gr |
| from transformers import ViltProcessor, ViltForQuestionAnswering |
| import torch |
|
|
| torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg') |
| torch.hub.download_url_to_file('https://cdn.openai.com/dall-e-2/demos/text2im/astronaut/horse/photo/0.jpg', 'astronaut.jpg') |
|
|
| processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa") |
| model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa") |
|
|
| def getAnswer(image,text): |
| encoding = processor(image, text, return_tensors="pt") |
| |
| |
| with torch.no_grad(): |
| outputs = model(**encoding) |
| |
| logits = outputs.logits |
| idx = logits.argmax(-1).item() |
| predicted_answer = model.config.id2label[idx] |
| |
| return predicted_answer |
|
|
| image = gr.inputs.Image(type="pil") |
| question = gr.inputs.Textbox(label="Question about the image") |
| answer = gr.outputs.Textbox(label="Predicted answer") |
| examples = [["cats.jpg", "How many cats are there?"], ["astronaut.jpg", "What's the astronaut riding on?"]] |
|
|
|
|
| title="Visual question and answering" |
|
|
| iface = gr.Interface(fn=getAnswer, |
| inputs=[image, question], |
| outputs=answer, |
| examples=examples, |
| title=title, |
| enable_queue=True) |
| iface.launch(debug=True ) |
|
|
