File size: 2,691 Bytes
6291dfb 8beaaf0 6291dfb 8beaaf0 6291dfb 8beaaf0 6291dfb 8beaaf0 6291dfb 8beaaf0 6291dfb 8beaaf0 6291dfb 8beaaf0 |
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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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()
|