|
import gradio as gr |
|
import torch |
|
from transformers import (AutoProcessor, BlipForQuestionAnswering, |
|
ViltForQuestionAnswering) |
|
|
|
torch.hub.download_url_to_file( |
|
'http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg') |
|
torch.hub.download_url_to_file( |
|
'https://huggingface.co/datasets/nielsr/textcaps-sample/resolve/main/stop_sign.png', |
|
'stop_sign.png') |
|
torch.hub.download_url_to_file( |
|
'https://cdn.openai.com/dall-e-2/demos/text2im/astronaut/horse/photo/0.jpg', |
|
'astronaut.jpg') |
|
|
|
blip_processor_large = AutoProcessor.from_pretrained( |
|
'Salesforce/blip-vqa-capfilt-large') |
|
blip_model_large = BlipForQuestionAnswering.from_pretrained( |
|
'Salesforce/blip-vqa-capfilt-large') |
|
|
|
vilt_processor = AutoProcessor.from_pretrained( |
|
'dandelin/vilt-b32-finetuned-vqa') |
|
vilt_model = ViltForQuestionAnswering.from_pretrained( |
|
'dandelin/vilt-b32-finetuned-vqa') |
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
blip_model_large.to(device) |
|
vilt_model.to(device) |
|
|
|
|
|
@torch.inference_mode() |
|
def generate_answer_blip(processor, model, image, question): |
|
inputs = processor(images=image, text=question, |
|
return_tensors='pt').to(device) |
|
generated_ids = model.generate(**inputs, max_length=50) |
|
generated_answer = processor.batch_decode(generated_ids, |
|
skip_special_tokens=True) |
|
return generated_answer[0] |
|
|
|
|
|
@torch.inference_mode() |
|
def generate_answer_vilt(processor, model, image, question): |
|
encoding = processor(images=image, text=question, |
|
return_tensors='pt').to(device) |
|
outputs = model(**encoding) |
|
predicted_class_idx = outputs.logits.argmax(-1).item() |
|
return model.config.id2label[predicted_class_idx] |
|
|
|
|
|
def generate_answers(image, question): |
|
answer_blip_large = generate_answer_blip(blip_processor_large, |
|
blip_model_large, image, question) |
|
answer_vilt = generate_answer_vilt(vilt_processor, vilt_model, image, |
|
question) |
|
return answer_blip_large, answer_vilt |
|
|
|
|
|
demo = gr.Interface( |
|
fn=generate_answers, |
|
inputs=[gr.Image(type='pil'), |
|
gr.Textbox(label='Question')], |
|
outputs=[ |
|
gr.Textbox(label='Answer generated by BLIP-large'), |
|
gr.Textbox(label='Answer generated by ViLT') |
|
], |
|
examples=[ |
|
['cats.jpg', 'How many cats are there?'], |
|
['stop_sign.png', "What's behind the stop sign?"], |
|
['astronaut.jpg', "What's the astronaut riding on?"], |
|
], |
|
title='Interactive demo: comparing visual question answering (VQA) models') |
|
demo.queue().launch() |
|
|
