Spaces:
Running
Running
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) |