| from transformers import pipeline | |
| class VisualQA(object): | |
| def __init__(self, model_name='nflechas/VQArt', tokenizer_name='dandelin/vilt-b32-finetuned-vqa'): | |
| self.model_name = model_name | |
| self.tokenizer_name = tokenizer_name | |
| self.__load_model() | |
| def __load_model(self): | |
| self.model = pipeline('vqa', model=self.model_name, tokenizer=self.tokenizer_name) | |
| def answer_question(self, query, image): | |
| return self.model(question=query, image=image, top_k=1)[0]['answer'] | |