Test / image.py
SilvusTV's picture
create function on image.py and call it on app.py
438833d
raw
history blame
No virus
764 Bytes
from transformers import ViltProcessor, ViltForQuestionAnswering
import requests
from PIL import Image
def image(url, text):
image = Image.open(requests.get(url, stream=True).raw)
processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
# prepare inputs
encoding = processor(image, text, return_tensors="pt")
# forward pass
outputs = model(**encoding)
logits = outputs.logits
idx = logits.argmax(-1).item()
print("question asked:", text)
print("image link:", url)
print("Predicted answer:", model.config.id2label[idx])
return "Predicted answer:", model.config.id2label[idx]
# prepare image + question