[WIP] Add OneformerFastImageProcessor (#38343)

* [WIP] OneformerFastImageProcessor

* update init

* Fully working oneformer image processor fast

* change Nearest to Neares exact interpolation where needed

* fix doc

---------

Co-authored-by: yonigozlan <yoni.gozlan@huggingface.co>
Co-authored-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com>

docs/source/en/model_doc/oneformer.md

@@ -69,7 +69,14 @@ The resource should ideally demonstrate something new instead of duplicating an
 
 [[autodoc]] OneFormerImageProcessor
     - preprocess
-    - encode_inputs
+    - post_process_semantic_segmentation
+    - post_process_instance_segmentation
+    - post_process_panoptic_segmentation
+
+## OneFormerImageProcessorFast
+
+[[autodoc]] OneFormerImageProcessorFast
+    - preprocess
     - post_process_semantic_segmentation
     - post_process_instance_segmentation
     - post_process_panoptic_segmentation
@@ -87,4 +94,3 @@ The resource should ideally demonstrate something new instead of duplicating an
 
 [[autodoc]] OneFormerForUniversalSegmentation
     - forward
-    

src/transformers/models/auto/image_processing_auto.py

@@ -130,7 +130,7 @@ else:
             ("mobilevitv2", ("MobileViTImageProcessor", "MobileViTImageProcessorFast")),
             ("nat", ("ViTImageProcessor", "ViTImageProcessorFast")),
             ("nougat", ("NougatImageProcessor", "NougatImageProcessorFast")),
-            ("oneformer", ("OneFormerImageProcessor",)),
+            ("oneformer", ("OneFormerImageProcessor", "OneFormerImageProcessorFast")),
             ("owlv2", ("Owlv2ImageProcessor",)),
             ("owlvit", ("OwlViTImageProcessor", "OwlViTImageProcessorFast")),
             ("paligemma", ("SiglipImageProcessor", "SiglipImageProcessorFast")),

src/transformers/models/conditional_detr/image_processing_conditional_detr_fast.py

@@ -454,10 +454,10 @@ class ConditionalDetrImageProcessorFast(BaseImageProcessorFast):
                 The target size of the image, as returned by the preprocessing `resize` step.
             threshold (`float`, *optional*, defaults to 0.5):
                 The threshold used to binarize the segmentation masks.
-            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST`):
+            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST_EXACT`):
                 The resampling filter to use when resizing the masks.
         """
-        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST
+        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
         ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
 
         new_annotation = {}

src/transformers/models/deformable_detr/image_processing_deformable_detr_fast.py

@@ -445,10 +445,10 @@ class DeformableDetrImageProcessorFast(BaseImageProcessorFast):
                 The target size of the image, as returned by the preprocessing `resize` step.
             threshold (`float`, *optional*, defaults to 0.5):
                 The threshold used to binarize the segmentation masks.
-            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST`):
+            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST_EXACT`):
                 The resampling filter to use when resizing the masks.
         """
-        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST
+        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
         ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
 
         new_annotation = {}

src/transformers/models/detr/image_processing_detr_fast.py

@@ -466,10 +466,10 @@ class DetrImageProcessorFast(BaseImageProcessorFast):
                 The target size of the image, as returned by the preprocessing `resize` step.
             threshold (`float`, *optional*, defaults to 0.5):
                 The threshold used to binarize the segmentation masks.
-            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST`):
+            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST_EXACT`):
                 The resampling filter to use when resizing the masks.
         """
-        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST
+        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
         ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
 
         new_annotation = {}

src/transformers/models/grounding_dino/image_processing_grounding_dino_fast.py

@@ -476,10 +476,10 @@ class GroundingDinoImageProcessorFast(BaseImageProcessorFast):
                 The target size of the image, as returned by the preprocessing `resize` step.
             threshold (`float`, *optional*, defaults to 0.5):
                 The threshold used to binarize the segmentation masks.
-            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST`):
+            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST_EXACT`):
                 The resampling filter to use when resizing the masks.
         """
-        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST
+        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
         ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
 
         new_annotation = {}

src/transformers/models/oneformer/__init__.py

@@ -20,6 +20,7 @@ from ...utils.import_utils import define_import_structure
 if TYPE_CHECKING:
     from .configuration_oneformer import *
     from .image_processing_oneformer import *
+    from .image_processing_oneformer_fast import *
     from .modeling_oneformer import *
     from .processing_oneformer import *
 else:

src/transformers/models/rt_detr/image_processing_rt_detr_fast.py

@@ -264,10 +264,10 @@ class RTDetrImageProcessorFast(BaseImageProcessorFast):
                 The target size of the image, as returned by the preprocessing `resize` step.
             threshold (`float`, *optional*, defaults to 0.5):
                 The threshold used to binarize the segmentation masks.
-            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST`):
+            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST_EXACT`):
                 The resampling filter to use when resizing the masks.
         """
-        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST
+        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
         ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
 
         new_annotation = {}

src/transformers/models/yolos/image_processing_yolos_fast.py

@@ -493,10 +493,10 @@ class YolosImageProcessorFast(BaseImageProcessorFast):
                 The target size of the image, as returned by the preprocessing `resize` step.
             threshold (`float`, *optional*, defaults to 0.5):
                 The threshold used to binarize the segmentation masks.
-            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST`):
+            resample (`InterpolationMode`, defaults to `InterpolationMode.NEAREST_EXACT`):
                 The resampling filter to use when resizing the masks.
         """
-        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST
+        interpolation = interpolation if interpolation is not None else F.InterpolationMode.NEAREST_EXACT
         ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)]
 
         new_annotation = {}

tests/models/oneformer/test_image_processing_oneformer.py

@@ -19,9 +19,10 @@ import tempfile
 import unittest
 
 import numpy as np
+from datasets import load_dataset
 
 from transformers.testing_utils import require_torch, require_vision
-from transformers.utils import is_torch_available, is_vision_available
+from transformers.utils import is_torch_available, is_torchvision_available, is_vision_available
 
 from ...test_image_processing_common import ImageProcessingTestMixin, prepare_image_inputs
 
@@ -31,6 +32,9 @@ if is_torch_available():
 
     if is_vision_available():
         from transformers import OneFormerImageProcessor
+
+    if is_torchvision_available():
+        from transformers import OneFormerImageProcessorFast
         from transformers.models.oneformer.image_processing_oneformer import binary_mask_to_rle, prepare_metadata
         from transformers.models.oneformer.modeling_oneformer import OneFormerForUniversalSegmentationOutput
 
@@ -152,12 +156,24 @@ class OneFormerImageProcessorTester:
         )
 
 
+# Copied from transformers.tests.models.beit.test_image_processing_beit.prepare_semantic_single_inputs
+def prepare_semantic_single_inputs():
+    ds = load_dataset("hf-internal-testing/fixtures_ade20k", split="test")
+    example = ds[0]
+    return example["image"], example["map"]
+
+
+# Copied from transformers.tests.models.beit.test_image_processing_beit.prepare_semantic_batch_inputs
+def prepare_semantic_batch_inputs():
+    ds = load_dataset("hf-internal-testing/fixtures_ade20k", split="test")
+    return list(ds["image"][:2]), list(ds["map"][:2])
+
+
 @require_torch
 @require_vision
 class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
     image_processing_class = OneFormerImageProcessor if (is_vision_available() and is_torch_available()) else None
-    # only for test_image_processing_common.test_image_proc_to_json_string
+    fast_image_processing_class = OneFormerImageProcessorFast if is_torchvision_available() else None
-    image_processing_class = image_processing_class
 
     def setUp(self):
         super().setUp()
@@ -168,7 +184,8 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
         return self.image_processor_tester.prepare_image_processor_dict()
 
     def test_image_proc_properties(self):
-        image_processor = self.image_processing_class(**self.image_processor_dict)
+        for image_processing_class in self.image_processor_list:
+            image_processor = image_processing_class(**self.image_processor_dict)
             self.assertTrue(hasattr(image_processor, "image_mean"))
             self.assertTrue(hasattr(image_processor, "image_std"))
             self.assertTrue(hasattr(image_processor, "do_normalize"))
@@ -182,9 +199,9 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
             self.assertTrue(hasattr(image_processor, "do_reduce_labels"))
 
     def comm_get_image_processor_inputs(
-        self, with_segmentation_maps=False, is_instance_map=False, segmentation_type="np"
+        self, with_segmentation_maps=False, is_instance_map=False, segmentation_type="np", image_processing_class=None
     ):
-        image_processor = self.image_processing_class(**self.image_processor_dict)
+        image_processor = image_processing_class(**self.image_processor_dict)
         # prepare image and target
         num_labels = self.image_processor_tester.num_labels
         annotations = None
@@ -218,8 +235,12 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
 
     def test_call_with_segmentation_maps(self):
         def common(is_instance_map=False, segmentation_type=None):
+            for image_processing_class in self.image_processor_list:
                 inputs = self.comm_get_image_processor_inputs(
-                with_segmentation_maps=True, is_instance_map=is_instance_map, segmentation_type=segmentation_type
+                    with_segmentation_maps=True,
+                    is_instance_map=is_instance_map,
+                    segmentation_type=segmentation_type,
+                    image_processing_class=image_processing_class,
                 )
 
                 mask_labels = inputs["mask_labels"]
@@ -251,7 +272,8 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
         self.assertEqual(rle[1], 45)
 
     def test_post_process_semantic_segmentation(self):
-        fature_extractor = self.image_processing_class(
+        for image_processing_class in self.image_processor_list:
+            fature_extractor = image_processing_class(
                 num_labels=self.image_processor_tester.num_classes,
                 max_seq_length=77,
                 task_seq_length=77,
@@ -278,7 +300,8 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
             self.assertEqual(segmentation[0].shape, target_sizes[0])
 
     def test_post_process_instance_segmentation(self):
-        image_processor = self.image_processing_class(
+        for image_processing_class in self.image_processor_list:
+            image_processor = image_processing_class(
                 num_labels=self.image_processor_tester.num_classes,
                 max_seq_length=77,
                 task_seq_length=77,
@@ -312,7 +335,8 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
                 self.assertEqual(el["segmentation"].shape, (1, 4))
 
     def test_post_process_panoptic_segmentation(self):
-        image_processor = self.image_processing_class(
+        for image_processing_class in self.image_processor_list:
+            image_processor = image_processing_class(
                 num_labels=self.image_processor_tester.num_classes,
                 max_seq_length=77,
                 task_seq_length=77,
@@ -340,7 +364,7 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
             "2": {"isthing": 1, "name": "baz"},
         }
         metadata = prepare_metadata(class_info)
-
+        for image_processing_class in self.image_processor_list:
             with tempfile.TemporaryDirectory() as tmpdirname:
                 metadata_path = os.path.join(tmpdirname, "metadata.json")
                 with open(metadata_path, "w") as f:
@@ -349,7 +373,7 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
                 config_dict = self.image_processor_dict
                 config_dict["class_info_file"] = metadata_path
                 config_dict["repo_path"] = tmpdirname
-            image_processor = self.image_processing_class(**config_dict)
+                image_processor = image_processing_class(**config_dict)
 
             self.assertEqual(image_processor.metadata, metadata)
 
@@ -357,11 +381,74 @@ class OneFormerImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
         image_processor_dict = dict(self.image_processor_dict)
         image_processor_dict.pop("do_reduce_labels", None)
         image_processor_dict["reduce_labels"] = True
-
+        # Only test for OneFormerImageProcessor
+        image_processing_class = self.image_processing_class
         # test we are able to create the image processor with the deprecated kwargs
-        image_processor = self.image_processing_class(**image_processor_dict)
+        image_processor = image_processing_class(**image_processor_dict)
         self.assertEqual(image_processor.do_reduce_labels, True)
 
         # test we still support reduce_labels with config
-        image_processor = self.image_processing_class.from_dict(image_processor_dict)
+        image_processor = image_processing_class.from_dict(image_processor_dict)
         self.assertEqual(image_processor.do_reduce_labels, True)
+
+    def test_slow_fast_equivalence(self):
+        if not self.test_slow_image_processor or not self.test_fast_image_processor:
+            self.skipTest(reason="Skipping slow/fast equivalence test")
+
+        if self.image_processing_class is None or self.fast_image_processing_class is None:
+            self.skipTest(reason="Skipping slow/fast equivalence test as one of the image processors is not defined")
+
+        dummy_image, dummy_map = prepare_semantic_single_inputs()
+
+        image_processor_slow = self.image_processing_class(**self.image_processor_dict)
+        image_processor_fast = self.fast_image_processing_class(**self.image_processor_dict)
+
+        image_encoding_slow = image_processor_slow(dummy_image, segmentation_maps=dummy_map, return_tensors="pt")
+        image_encoding_fast = image_processor_fast(dummy_image, segmentation_maps=dummy_map, return_tensors="pt")
+        self._assert_slow_fast_tensors_equivalence(image_encoding_slow.pixel_values, image_encoding_fast.pixel_values)
+        for mask_label_slow, mask_label_fast in zip(image_encoding_slow.mask_labels, image_encoding_fast.mask_labels):
+            self._assert_slow_fast_tensors_equivalence(mask_label_slow, mask_label_fast)
+        for class_label_slow, class_label_fast in zip(
+            image_encoding_slow.class_labels, image_encoding_fast.class_labels
+        ):
+            self._assert_slow_fast_tensors_equivalence(class_label_slow.float(), class_label_fast.float())
+        self.assertEqual(image_encoding_slow.text_inputs, image_encoding_fast.text_inputs)
+        self.assertEqual(image_encoding_slow.task_inputs, image_encoding_fast.task_inputs)
+
+    def test_slow_fast_equivalence_batched(self):
+        if not self.test_slow_image_processor or not self.test_fast_image_processor:
+            self.skipTest(reason="Skipping slow/fast equivalence test")
+
+        if self.image_processing_class is None or self.fast_image_processing_class is None:
+            self.skipTest(reason="Skipping slow/fast equivalence test as one of the image processors is not defined")
+
+        if hasattr(self.image_processor_tester, "do_center_crop") and self.image_processor_tester.do_center_crop:
+            self.skipTest(
+                reason="Skipping as do_center_crop is True and center_crop functions are not equivalent for fast and slow processors"
+            )
+
+        dummy_images, dummy_maps = prepare_semantic_batch_inputs()
+
+        image_processor_slow = self.image_processing_class(**self.image_processor_dict)
+        image_processor_fast = self.fast_image_processing_class(**self.image_processor_dict)
+
+        encoding_slow = image_processor_slow(
+            dummy_images,
+            segmentation_maps=dummy_maps,
+            task_inputs=["instance"] + ["semantic"] * (len(dummy_images) - 1),
+            return_tensors="pt",
+        )
+        encoding_fast = image_processor_fast(
+            dummy_images,
+            segmentation_maps=dummy_maps,
+            task_inputs=["instance"] + ["semantic"] * (len(dummy_images) - 1),
+            return_tensors="pt",
+        )
+
+        self._assert_slow_fast_tensors_equivalence(encoding_slow.pixel_values, encoding_fast.pixel_values)
+        for mask_label_slow, mask_label_fast in zip(encoding_slow.mask_labels, encoding_fast.mask_labels):
+            self._assert_slow_fast_tensors_equivalence(mask_label_slow, mask_label_fast)
+        for class_label_slow, class_label_fast in zip(encoding_slow.class_labels, encoding_fast.class_labels):
+            self._assert_slow_fast_tensors_equivalence(class_label_slow.float(), class_label_fast.float())
+        self.assertEqual(encoding_slow.text_inputs, encoding_fast.text_inputs)
+        self.assertEqual(encoding_slow.task_inputs, encoding_fast.task_inputs)