File size: 1,030 Bytes
30f5735
 
 
 
c1d8fa7
 
30f5735
 
 
 
 
 
 
 
 
 
37ce7af
 
12b00d2
30f5735
37ce7af
30f5735
 
 
 
c1d8fa7
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import gradio as gr
from transformers import ViltProcessor, ViltForVisualQuestionAnswering
import torch

torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')

processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
model = ViltForVisualQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")

def answer_question(image, text):
    encoding = processor(image, text, return_tensors="pt")
    
    # forward pass
    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")
answer = gr.outputs.Textbox(label="Predicted answer")
examples = [["cats.jpg"], ["How many cats are there?"]]
gr.Interface(fn=answer_question, inputs=[image, question], outputs=answer, examples=examples, enable_queue=True).launch(debug=True)