init

transformers/tests/pipelines/test_pipelines_image_feature_extraction.py

@@ -0,0 +1,139 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import pytest
+
+from transformers import (
+    MODEL_MAPPING,
+    TOKENIZER_MAPPING,
+    ImageFeatureExtractionPipeline,
+    is_torch_available,
+    is_vision_available,
+    pipeline,
+)
+from transformers.testing_utils import is_pipeline_test, nested_simplify, require_torch
+
+
+if is_torch_available():
+    import torch
+
+if is_vision_available():
+    from PIL import Image
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+    image = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
+    return image
+
+
+@is_pipeline_test
+class ImageFeatureExtractionPipelineTests(unittest.TestCase):
+    model_mapping = MODEL_MAPPING
+
+    @require_torch
+    def test_small_model_pt(self):
+        feature_extractor = pipeline(task="image-feature-extraction", model="hf-internal-testing/tiny-random-vit")
+        img = prepare_img()
+        outputs = feature_extractor(img)
+        self.assertEqual(
+            nested_simplify(outputs[0][0]),
+            [-1.417, -0.392, -1.264, -1.196, 1.648, 0.885, 0.56, -0.606, -1.175, 0.823, 1.912, 0.081, -0.053, 1.119, -0.062, -1.757, -0.571, 0.075, 0.959, 0.118, 1.201, -0.672, -0.498, 0.364, 0.937, -1.623, 0.228, 0.19, 1.697, -1.115, 0.583, -0.981])  # fmt: skip
+
+    @require_torch
+    def test_small_model_w_pooler_pt(self):
+        feature_extractor = pipeline(
+            task="image-feature-extraction", model="hf-internal-testing/tiny-random-vit-w-pooler"
+        )
+        img = prepare_img()
+        outputs = feature_extractor(img, pool=True)
+        self.assertEqual(
+            nested_simplify(outputs[0]),
+            [-0.056,  0.083,  0.021,  0.038,  0.242, -0.279, -0.033, -0.003, 0.200, -0.192,  0.045, -0.095, -0.077,  0.017, -0.058, -0.063, -0.029, -0.204,  0.014,  0.042,  0.305, -0.205, -0.099,  0.146, -0.287,  0.020,  0.168, -0.052,  0.046,  0.048, -0.156,  0.093])  # fmt: skip
+
+    @require_torch
+    def test_image_processing_small_model_pt(self):
+        feature_extractor = pipeline(task="image-feature-extraction", model="hf-internal-testing/tiny-random-vit")
+
+        # test with image processor parameters
+        image_processor_kwargs = {"size": {"height": 300, "width": 300}}
+        img = prepare_img()
+        with pytest.raises(ValueError):
+            # Image doesn't match model input size
+            feature_extractor(img, image_processor_kwargs=image_processor_kwargs)
+
+        image_processor_kwargs = {"image_mean": [0, 0, 0], "image_std": [1, 1, 1]}
+        img = prepare_img()
+        outputs = feature_extractor(img, image_processor_kwargs=image_processor_kwargs)
+        self.assertEqual(np.squeeze(outputs).shape, (226, 32))
+
+        # Test pooling option
+        outputs = feature_extractor(img, pool=True)
+        self.assertEqual(np.squeeze(outputs).shape, (32,))
+
+    @require_torch
+    def test_return_tensors_pt(self):
+        feature_extractor = pipeline(task="image-feature-extraction", model="hf-internal-testing/tiny-random-vit")
+        img = prepare_img()
+        outputs = feature_extractor(img, return_tensors=True)
+        self.assertTrue(torch.is_tensor(outputs))
+
+    def get_test_pipeline(
+        self,
+        model,
+        tokenizer=None,
+        image_processor=None,
+        feature_extractor=None,
+        processor=None,
+        dtype="float32",
+    ):
+        if image_processor is None:
+            self.skipTest(reason="No image processor")
+
+        elif type(model.config) in TOKENIZER_MAPPING:
+            self.skipTest(
+                reason="This is a bimodal model, we need to find a more consistent way to switch on those models."
+            )
+
+        elif model.config.is_encoder_decoder:
+            self.skipTest(
+                """encoder_decoder models are trickier for this pipeline.
+                Do we want encoder + decoder inputs to get some features?
+                Do we want encoder only features ?
+                For now ignore those.
+                """
+            )
+
+        feature_extractor_pipeline = ImageFeatureExtractionPipeline(
+            model=model,
+            tokenizer=tokenizer,
+            feature_extractor=feature_extractor,
+            image_processor=image_processor,
+            processor=processor,
+            dtype=dtype,
+        )
+        img = prepare_img()
+        return feature_extractor_pipeline, [img, img]
+
+    def run_pipeline_test(self, feature_extractor, examples):
+        imgs = examples
+        outputs = feature_extractor(imgs[0])
+
+        self.assertEqual(len(outputs), 1)
+
+        outputs = feature_extractor(imgs)
+        self.assertEqual(len(outputs), 2)