Pipeline VQA: Add support for list of images and questions as pipeline input (#31217)

* Add list check for image and question

* Handle passing two lists and update docstring

* Add tests

* Add support for dataset

* Add test for dataset as input

* fixup

* fix unprotected import

* fix unprotected import

* fix import again

* fix param type
This commit is contained in:
Vu Huy Nguyen
2024-06-06 15:50:45 +02:00
committed by GitHub
parent 4c82102523
commit f9296249a3
2 changed files with 102 additions and 4 deletions

View File

@@ -14,6 +14,8 @@
import unittest
from datasets import load_dataset
from transformers import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING, is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
@@ -34,6 +36,8 @@ from .test_pipelines_common import ANY
if is_torch_available():
import torch
from transformers.pipelines.pt_utils import KeyDataset
if is_vision_available():
from PIL import Image
@@ -172,6 +176,65 @@ class VisualQuestionAnsweringPipelineTests(unittest.TestCase):
outputs = vqa_pipeline([{"image": image, "question": question}, {"image": image, "question": question}])
self.assertEqual(outputs, [[{"answer": "two"}]] * 2)
@require_torch
def test_small_model_pt_image_list(self):
vqa_pipeline = pipeline("visual-question-answering", model="hf-internal-testing/tiny-vilt-random-vqa")
images = [
"./tests/fixtures/tests_samples/COCO/000000039769.png",
"./tests/fixtures/tests_samples/COCO/000000004016.png",
]
outputs = vqa_pipeline(image=images, question="How many cats are there?", top_k=1)
self.assertEqual(
outputs, [[{"score": ANY(float), "answer": ANY(str)}], [{"score": ANY(float), "answer": ANY(str)}]]
)
@require_torch
def test_small_model_pt_question_list(self):
vqa_pipeline = pipeline("visual-question-answering", model="hf-internal-testing/tiny-vilt-random-vqa")
image = "./tests/fixtures/tests_samples/COCO/000000039769.png"
questions = ["How many cats are there?", "Are there any dogs?"]
outputs = vqa_pipeline(image=image, question=questions, top_k=1)
self.assertEqual(
outputs, [[{"score": ANY(float), "answer": ANY(str)}], [{"score": ANY(float), "answer": ANY(str)}]]
)
@require_torch
def test_small_model_pt_both_list(self):
vqa_pipeline = pipeline("visual-question-answering", model="hf-internal-testing/tiny-vilt-random-vqa")
images = [
"./tests/fixtures/tests_samples/COCO/000000039769.png",
"./tests/fixtures/tests_samples/COCO/000000004016.png",
]
questions = ["How many cats are there?", "Are there any dogs?"]
outputs = vqa_pipeline(image=images, question=questions, top_k=1)
self.assertEqual(
outputs,
[
[{"score": ANY(float), "answer": ANY(str)}],
[{"score": ANY(float), "answer": ANY(str)}],
[{"score": ANY(float), "answer": ANY(str)}],
[{"score": ANY(float), "answer": ANY(str)}],
],
)
@require_torch
def test_small_model_pt_dataset(self):
vqa_pipeline = pipeline("visual-question-answering", model="hf-internal-testing/tiny-vilt-random-vqa")
dataset = load_dataset("hf-internal-testing/dummy_image_text_data", split="train[:2]")
question = "What's in the image?"
outputs = vqa_pipeline(image=KeyDataset(dataset, "image"), question=question, top_k=1)
self.assertEqual(
outputs,
[
[{"score": ANY(float), "answer": ANY(str)}],
[{"score": ANY(float), "answer": ANY(str)}],
],
)
@require_tf
@unittest.skip("Visual question answering not implemented in TF")
def test_small_model_tf(self):