vilt-vqa / app.py
nielsr's picture
nielsr HF staff
Fix bug
12b00d2
raw history blame
No virus
852 Bytes
import gradio as gr
from transformers import ViltProcessor, ViltForVisualQuestionAnswering
import torch
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")
gr.Interface(fn=classify_image, inputs=[image, question], outputs=answer, enable_queue=True).launch(debug=True)