Repository: huggingface/transformers
Pull request: Fix MaskFormer/Mask2Former fast image processors

# What does this PR do?

Depends on https://github.com/huggingface/transformers/pull/41391.
These two fast image processors had issues and were not properly tested:
- There was an issue where the processors would crash if `do_resize-False`
- After conversion to binary masks, the grouped masks cannot be stacked anymore, as their channels dimensions are not the same. This fix uses the method introduced in https://github.com/huggingface/transformers/pull/41391 to group the masks according to the shapes of the corresponding images.

This PR fixes the issues and ensure that the integration tests are also ran with the fast image processors

Changed files:
filename: src/transformers/image_transforms.py
patch: @@ -828,7 +828,25 @@ def _cast_tensor_to_float(x):
 
 
 def _group_images_by_shape(nested_images, *paired_inputs, is_nested: bool = False):
-    """Helper function to flatten a single level of nested image and batch structures and group by shape."""
+    """
+    Helper function to flatten a single level of nested image and batch structures and group by shape.
+    Args:
+        nested_images (list):
+            A list of images or a single tensor
+        paired_inputs (Any, *optional*):
+            Zero or more lists that mirror the structure of `nested_images` (flat list, or list of lists when
+            `is_nested=True`). Each element is paired 1:1 with the corresponding image so it can be grouped by the
+            same shape key. These paired values are grouped alongside `nested_images` but are not stacked in the output, so
+            they do not need to be tensors.
+        is_nested (bool, *optional*, defaults to False):
+            Whether the images are nested.
+    Returns:
+        tuple[dict, ...]:
+            - A dictionary with shape as key and list of images with that shape as value
+            - A dictionary with shape as key and list of paired values with that shape as value
+            - A dictionary mapping original indices to (shape, index) tuples
+            - A dictionary mapping original indices to (shape, index) tuples for each paired input
+    """
     grouped_images = defaultdict(list)
     grouped_images_index = {}
     paired_grouped_values = [defaultdict(list) for _ in paired_inputs]
@@ -880,27 +898,20 @@ def _reconstruct_nested_structure(indices, processed_images):
     return result
 
 
-def _disable_grouping_output_nested(images, *paired_inputs):
-    """Build the disable_grouping output tuple for a single-level nested structure."""
-    outer_range = range(len(images))
-    inner_ranges = [range(len(images[i])) for i in outer_range]
-
-    # Precompute all (i, j) pairs
-    ij_pairs = [(i, j) for i in outer_range for j in inner_ranges[i]]
-
-    images_dict = {(i, j): images[i][j].unsqueeze(0) for (i, j) in ij_pairs}
-    paired_dicts = [{(i, j): paired_list[i][j].unsqueeze(0) for (i, j) in ij_pairs} for paired_list in paired_inputs]
-    index_map = {(i, j): ((i, j), 0) for (i, j) in ij_pairs}
-    return images_dict, *paired_dicts, index_map
-
+def _iterate_items(items, is_nested: bool):
+    """
+    Helper function to iterate over items yielding (key, item) pairs.
 
-def _disable_grouping_output_flat(images, *paired_inputs):
-    """Build the disable_grouping output tuple for a flat list structure."""
-    idx_range = range(len(images))
-    images_dict = {i: images[i].unsqueeze(0) for i in idx_range}
-    paired_dicts = [{i: paired_list[i].unsqueeze(0) for i in idx_range} for paired_list in paired_inputs]
-    index_map = {i: (i, 0) for i in idx_range}
-    return images_dict, *paired_dicts, index_map
+    For nested structures, yields ((row_index, col_index), item).
+    For flat structures, yields (index, item).
+    """
+    if is_nested:
+        for i, row in enumerate(items):
+            for j, item in enumerate(row):
+                yield (i, j), item
+    else:
+        for i, item in enumerate(items):
+            yield i, item
 
 
 def group_images_by_shape(
@@ -920,7 +931,7 @@ def group_images_by_shape(
     Args:
         images (Union[list["torch.Tensor"], "torch.Tensor"]):
             A list of images or a single tensor
-        *paired_inputs (Any):
+        paired_inputs (Any, *optional*):
             Zero or more lists that mirror the structure of `images` (flat list, or list of lists when
             `is_nested=True`). Each element is paired 1:1 with the corresponding image so it can be grouped by the
             same shape key. These paired values are grouped alongside `images` but are not stacked in the output, so
@@ -944,10 +955,14 @@ def group_images_by_shape(
         disable_grouping = device == "cpu"
 
     if disable_grouping:
-        if is_nested:
-            return _disable_grouping_output_nested(images, *paired_inputs)
-        else:
-            return _disable_grouping_output_flat(images, *paired_inputs)
+        return (
+            {key: img.unsqueeze(0) for key, img in _iterate_items(images, is_nested)},
+            *[
+                {key: item.unsqueeze(0) for key, item in _iterate_items(paired_list, is_nested)}
+                for paired_list in paired_inputs
+            ],
+            {key: (key, 0) for key, _ in _iterate_items(images, is_nested)},
+        )
 
     # Handle single level nested structure
     grouped_images, *paired_grouped_values, grouped_images_index = _group_images_by_shape(

filename: src/transformers/models/mask2former/image_processing_mask2former_fast.py
patch: @@ -109,9 +109,6 @@ class Mask2FormerImageProcessorFast(BaseImageProcessorFast):
     valid_kwargs = Mask2FormerImageProcessorKwargs
 
     def __init__(self, **kwargs: Unpack[Mask2FormerImageProcessorKwargs]) -> None:
-        if "pad_and_return_pixel_mask" in kwargs:
-            kwargs["do_pad"] = kwargs.pop("pad_and_return_pixel_mask")
-
         size = kwargs.pop("size", None)
         max_size = kwargs.pop("max_size", None)
 
@@ -224,7 +221,7 @@ def pad(
             padding = [0, 0, padding_right, padding_bottom]
             images = F.pad(images, padding, fill=fill)
             if segmentation_maps is not None:
-                segmentation_maps = F.pad(segmentation_maps, padding, fill=ignore_index)
+                segmentation_maps = [F.pad(mask, padding, fill=ignore_index) for mask in segmentation_maps]
 
         # Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.
         pixel_mask = torch.zeros((images.shape[0], *padded_size), dtype=torch.int64, device=images.device)
@@ -318,9 +315,11 @@ def _preprocess(
                 stacked_images = self.resize(
                     image=stacked_images, size=size, size_divisor=size_divisor, interpolation=interpolation
                 )
-                if segmentation_maps is not None:
+            if segmentation_maps is not None:
+                stacked_segmentation_maps = grouped_segmentation_maps[shape]
+                if do_resize:
                     stacked_segmentation_maps = self.resize(
-                        image=grouped_segmentation_maps[shape],
+                        image=stacked_segmentation_maps,
                         size=size,
                         size_divisor=size_divisor,
                         interpolation=F.InterpolationMode.NEAREST_EXACT,
@@ -357,14 +356,18 @@ def _preprocess(
                 mask_labels.append(masks)
                 class_labels.append(classes)
 
-        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
-        processed_images_grouped = {}
-        processed_pixel_masks_grouped = {}
         if segmentation_maps is not None:
-            grouped_segmentation_maps, grouped_segmentation_maps_index = group_images_by_shape(
-                mask_labels, disable_grouping=disable_grouping
+            # group mask_labels as paired inputs and not images so as not to stack them
+            grouped_images, grouped_segmentation_maps, grouped_images_index = group_images_by_shape(
+                resized_images, mask_labels, disable_grouping=disable_grouping
             )
             processed_segmentation_maps_grouped = {}
+        else:
+            grouped_images, grouped_images_index = group_images_by_shape(
+                resized_images, disable_grouping=disable_grouping
+            )
+        processed_images_grouped = {}
+        processed_pixel_masks_grouped = {}
         for shape, stacked_images in grouped_images.items():
             # Fused rescale and normalize
             stacked_images = self.rescale_and_normalize(
@@ -379,7 +382,8 @@ def _preprocess(
             processed_images_grouped[shape] = padded_images
             processed_pixel_masks_grouped[shape] = pixel_masks
             if segmentation_maps is not None:
-                processed_segmentation_maps_grouped[shape] = padded_segmentation_maps.squeeze(1)
+                processed_segmentation_maps_grouped[shape] = padded_segmentation_maps
+
         processed_images = reorder_images(processed_images_grouped, grouped_images_index)
         processed_pixel_masks = reorder_images(processed_pixel_masks_grouped, grouped_images_index)
         encoded_inputs = BatchFeature(
@@ -390,7 +394,7 @@ def _preprocess(
             tensor_type=return_tensors,
         )
         if segmentation_maps is not None:
-            mask_labels = reorder_images(processed_segmentation_maps_grouped, grouped_segmentation_maps_index)
+            mask_labels = reorder_images(processed_segmentation_maps_grouped, grouped_images_index)
             # we cannot batch them since they don't share a common class size
             encoded_inputs["mask_labels"] = mask_labels
             encoded_inputs["class_labels"] = class_labels

filename: src/transformers/models/maskformer/image_processing_maskformer_fast.py
patch: @@ -114,9 +114,6 @@ class MaskFormerImageProcessorFast(BaseImageProcessorFast):
     valid_kwargs = MaskFormerImageProcessorKwargs
 
     def __init__(self, **kwargs: Unpack[MaskFormerImageProcessorKwargs]) -> None:
-        if "pad_and_return_pixel_mask" in kwargs:
-            kwargs["do_pad"] = kwargs.pop("pad_and_return_pixel_mask")
-
         size = kwargs.pop("size", None)
         max_size = kwargs.pop("max_size", None)
 
@@ -229,7 +226,7 @@ def pad(
             padding = [0, 0, padding_right, padding_bottom]
             images = F.pad(images, padding, fill=fill)
             if segmentation_maps is not None:
-                segmentation_maps = F.pad(segmentation_maps, padding, fill=ignore_index)
+                segmentation_maps = [F.pad(mask, padding, fill=ignore_index) for mask in segmentation_maps]
 
         # Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.
         pixel_mask = torch.zeros((images.shape[0], *padded_size), dtype=torch.int64, device=images.device)
@@ -323,9 +320,11 @@ def _preprocess(
                 stacked_images = self.resize(
                     image=stacked_images, size=size, size_divisor=size_divisor, interpolation=interpolation
                 )
-                if segmentation_maps is not None:
+            if segmentation_maps is not None:
+                stacked_segmentation_maps = grouped_segmentation_maps[shape]
+                if do_resize:
                     stacked_segmentation_maps = self.resize(
-                        image=grouped_segmentation_maps[shape],
+                        image=stacked_segmentation_maps,
                         size=size,
                         size_divisor=size_divisor,
                         interpolation=F.InterpolationMode.NEAREST_EXACT,
@@ -362,14 +361,18 @@ def _preprocess(
                 mask_labels.append(masks)
                 class_labels.append(classes)
 
-        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
-        processed_images_grouped = {}
-        processed_pixel_masks_grouped = {}
         if segmentation_maps is not None:
-            grouped_segmentation_maps, grouped_segmentation_maps_index = group_images_by_shape(
-                mask_labels, disable_grouping=disable_grouping
+            # group mask_labels as paired inputs and not images so as not to stack them
+            grouped_images, grouped_segmentation_maps, grouped_images_index = group_images_by_shape(
+                resized_images, mask_labels, disable_grouping=disable_grouping
             )
             processed_segmentation_maps_grouped = {}
+        else:
+            grouped_images, grouped_images_index = group_images_by_shape(
+                resized_images, disable_grouping=disable_grouping
+            )
+        processed_images_grouped = {}
+        processed_pixel_masks_grouped = {}
         for shape, stacked_images in grouped_images.items():
             # Fused rescale and normalize
             stacked_images = self.rescale_and_normalize(
@@ -384,7 +387,8 @@ def _preprocess(
             processed_images_grouped[shape] = padded_images
             processed_pixel_masks_grouped[shape] = pixel_masks
             if segmentation_maps is not None:
-                processed_segmentation_maps_grouped[shape] = padded_segmentation_maps.squeeze(1)
+                processed_segmentation_maps_grouped[shape] = padded_segmentation_maps
+
         processed_images = reorder_images(processed_images_grouped, grouped_images_index)
         processed_pixel_masks = reorder_images(processed_pixel_masks_grouped, grouped_images_index)
         encoded_inputs = BatchFeature(
@@ -395,7 +399,7 @@ def _preprocess(
             tensor_type=return_tensors,
         )
         if segmentation_maps is not None:
-            mask_labels = reorder_images(processed_segmentation_maps_grouped, grouped_segmentation_maps_index)
+            mask_labels = reorder_images(processed_segmentation_maps_grouped, grouped_images_index)
             # we cannot batch them since they don't share a common class size
             encoded_inputs["mask_labels"] = mask_labels
             encoded_inputs["class_labels"] = class_labels

filename: tests/models/mask2former/test_image_processing_mask2former.py
patch: @@ -194,13 +194,14 @@ def test_image_processor_properties(self):
     def comm_get_image_processing_inputs(
         self,
         image_processor_tester,
+        image_processing_class,
         with_segmentation_maps=False,
         is_instance_map=False,
         segmentation_type="np",
         numpify=False,
         input_data_format=None,
     ):
-        image_processing = self.image_processing_class(**image_processor_tester.prepare_image_processor_dict())
+        image_processing = image_processing_class(**image_processor_tester.prepare_image_processor_dict())
         # prepare image and target
         num_labels = image_processor_tester.num_labels
         annotations = None
@@ -228,7 +229,6 @@ def comm_get_image_processing_inputs(
             annotations,
             return_tensors="pt",
             instance_id_to_semantic_id=instance_id_to_semantic_id,
-            pad_and_return_pixel_mask=True,
             input_data_format=input_data_format,
         )
 
@@ -264,25 +264,26 @@ def common(
                 image_mean=[0.5] * num_channels,
                 image_std=[0.5] * num_channels,
             )
+            for image_processing_class in self.image_processor_list:
+                inputs = self.comm_get_image_processing_inputs(
+                    image_processor_tester=image_processor_tester,
+                    image_processing_class=image_processing_class,
+                    with_segmentation_maps=True,
+                    is_instance_map=is_instance_map,
+                    segmentation_type=segmentation_type,
+                    numpify=numpify,
+                    input_data_format=input_data_format,
+                )
 
-            inputs = self.comm_get_image_processing_inputs(
-                image_processor_tester=image_processor_tester,
-                with_segmentation_maps=True,
-                is_instance_map=is_instance_map,
-                segmentation_type=segmentation_type,
-                numpify=numpify,
-                input_data_format=input_data_format,
-            )
-
-            mask_labels = inputs["mask_labels"]
-            class_labels = inputs["class_labels"]
-            pixel_values = inputs["pixel_values"]
+                mask_labels = inputs["mask_labels"]
+                class_labels = inputs["class_labels"]
+                pixel_values = inputs["pixel_values"]
 
-            # check the batch_size
-            for mask_label, class_label in zip(mask_labels, class_labels):
-                self.assertEqual(mask_label.shape[0], class_label.shape[0])
-                # this ensure padding has happened
-                self.assertEqual(mask_label.shape[1:], pixel_values.shape[2:])
+                # check the batch_size
+                for mask_label, class_label in zip(mask_labels, class_labels):
+                    self.assertEqual(mask_label.shape[0], class_label.shape[0])
+                    # this ensure padding has happened
+                    self.assertEqual(mask_label.shape[1:], pixel_values.shape[2:])
 
         common()
         common(is_instance_map=True)
@@ -335,31 +336,32 @@ def get_instance_segmentation_and_mapping(annotation):
         instance_seg2, inst2class2 = get_instance_segmentation_and_mapping(annotation2)
 
         # create a image processor
-        image_processing = Mask2FormerImageProcessor(do_reduce_labels=True, ignore_index=255, size=(512, 512))
-
-        # prepare the images and annotations
-        inputs = image_processing(
-            [image1, image2],
-            [instance_seg1, instance_seg2],
-            instance_id_to_semantic_id=[inst2class1, inst2class2],
-            return_tensors="pt",
-        )
+        for image_processing_class in self.image_processor_list:
+            image_processing = image_processing_class(do_reduce_labels=True, ignore_index=255, size=(512, 512))
+
+            # prepare the images and annotations
+            inputs = image_processing(
+                [image1, image2],
+                [instance_seg1, instance_seg2],
+                instance_id_to_semantic_id=[inst2class1, inst2class2],
+                return_tensors="pt",
+            )
 
-        # verify the pixel values and pixel mask
-        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
-        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
+            # verify the pixel values and pixel mask
+            self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
+            self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
 
-        # verify the class labels
-        self.assertEqual(len(inputs["class_labels"]), 2)
-        torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([30, 55]))
-        torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([4, 4, 23, 55]))
+            # verify the class labels
+            self.assertEqual(len(inputs["class_labels"]), 2)
+            torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([30, 55]))
+            torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([4, 4, 23, 55]))
 
-        # verify the mask labels
-        self.assertEqual(len(inputs["mask_labels"]), 2)
-        self.assertEqual(inputs["mask_labels"][0].shape, (2, 512, 512))
-        self.assertEqual(inputs["mask_labels"][1].shape, (4, 512, 512))
-        self.assertEqual(inputs["mask_labels"][0].sum().item(), 41527.0)
-        self.assertEqual(inputs["mask_labels"][1].sum().item(), 26259.0)
+            # verify the mask labels
+            self.assertEqual(len(inputs["mask_labels"]), 2)
+            self.assertEqual(inputs["mask_labels"][0].shape, (2, 512, 512))
+            self.assertEqual(inputs["mask_labels"][1].shape, (4, 512, 512))
+            self.assertEqual(inputs["mask_labels"][0].sum().item(), 41527.0)
+            self.assertEqual(inputs["mask_labels"][1].sum().item(), 26259.0)
 
     def test_integration_semantic_segmentation(self):
         # load 2 images and corresponding semantic annotations from the hub
@@ -378,30 +380,31 @@ def test_integration_semantic_segmentation(self):
         )
 
         # create a image processor
-        image_processing = Mask2FormerImageProcessor(do_reduce_labels=True, ignore_index=255, size=(512, 512))
+        for image_processing_class in self.image_processor_list:
+            image_processing = image_processing_class(do_reduce_labels=True, ignore_index=255, size=(512, 512))
 
-        # prepare the images and annotations
-        inputs = image_processing(
-            [image1, image2],
-            [annotation1, annotation2],
-            return_tensors="pt",
-        )
+            # prepare the images and annotations
+            inputs = image_processing(
+                [image1, image2],
+                [annotation1, annotation2],
+                return_tensors="pt",
+            )
 
-        # verify the pixel values and pixel mask
-        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
-        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
+            # verify the pixel values and pixel mask
+            self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
+            self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
 
-        # verify the class labels
-        self.assertEqual(len(inputs["class_labels"]), 2)
-        torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([2, 4, 60]))
-        torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([0, 3, 7, 8, 15, 28, 30, 143]))
+            # verify the class labels
+            self.assertEqual(len(inputs["class_labels"]), 2)
+            torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([2, 4, 60]))
+            torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([0, 3, 7, 8, 15, 28, 30, 143]))
 
-        # verify the mask labels
-        self.assertEqual(len(inputs["mask_labels"]), 2)
-        self.assertEqual(inputs["mask_labels"][0].shape, (3, 512, 512))
-        self.assertEqual(inputs["mask_labels"][1].shape, (8, 512, 512))
-        self.assertEqual(inputs["mask_labels"][0].sum().item(), 170200.0)
-        self.assertEqual(inputs["mask_labels"][1].sum().item(), 257036.0)
+            # verify the mask labels
+            self.assertEqual(len(inputs["mask_labels"]), 2)
+            self.assertEqual(inputs["mask_labels"][0].shape, (3, 512, 512))
+            self.assertEqual(inputs["mask_labels"][1].shape, (8, 512, 512))
+            self.assertEqual(inputs["mask_labels"][0].sum().item(), 170200.0)
+            self.assertEqual(inputs["mask_labels"][1].sum().item(), 257036.0)
 
     def test_integration_panoptic_segmentation(self):
         # load 2 images and corresponding panoptic annotations from the hub
@@ -435,34 +438,35 @@ def create_panoptic_map(annotation, segments_info):
         panoptic_map2, inst2class2 = create_panoptic_map(annotation2, segments_info2)
 
         # create a image processor
-        image_processing = Mask2FormerImageProcessor(ignore_index=0, do_resize=False)
-
-        # prepare the images and annotations
-        pixel_values_list = [np.moveaxis(np.array(image1), -1, 0), np.moveaxis(np.array(image2), -1, 0)]
-        inputs = image_processing.encode_inputs(
-            pixel_values_list,
-            [panoptic_map1, panoptic_map2],
-            instance_id_to_semantic_id=[inst2class1, inst2class2],
-            return_tensors="pt",
-        )
+        for image_processing_class in self.image_processor_list:
+            image_processing = image_processing_class(ignore_index=0, do_resize=False)
+
+            # prepare the images and annotations
+            pixel_values_list = [np.moveaxis(np.array(image1), -1, 0), np.moveaxis(np.array(image2), -1, 0)]
+            inputs = image_processing(
+                pixel_values_list,
+                [panoptic_map1, panoptic_map2],
+                instance_id_to_semantic_id=[inst2class1, inst2class2],
+                return_tensors="pt",
+            )
 
-        # verify the pixel values and pixel mask
-        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 711))
-        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 711))
-
-        # verify the class labels
-        self.assertEqual(len(inputs["class_labels"]), 2)
-        expected_class_labels = torch.tensor([4, 17, 32, 42, 42, 42, 42, 42, 42, 42, 32, 12, 12, 12, 12, 12, 42, 42, 12, 12, 12, 42, 12, 12, 12, 12, 12, 3, 12, 12, 12, 12, 42, 42, 42, 12, 42, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 5, 12, 12, 12, 12, 12, 12, 12, 0, 43, 43, 43, 96, 43, 104, 43, 31, 125, 31, 125, 138, 87, 125, 149, 138, 125, 87, 87])  # fmt: skip
-        torch.testing.assert_close(inputs["class_labels"][0], torch.tensor(expected_class_labels))
-        expected_class_labels = torch.tensor([19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 67, 82, 19, 19, 17, 19, 19, 19, 19, 19, 19, 19, 19, 19, 12, 12, 42, 12, 12, 12, 12, 3, 14, 12, 12, 12, 12, 12, 12, 12, 12, 14, 5, 12, 12, 0, 115, 43, 43, 115, 43, 43, 43, 8, 8, 8, 138, 138, 125, 143])  # fmt: skip
-        torch.testing.assert_close(inputs["class_labels"][1], expected_class_labels)
-
-        # verify the mask labels
-        self.assertEqual(len(inputs["mask_labels"]), 2)
-        self.assertEqual(inputs["mask_labels"][0].shape, (79, 512, 711))
-        self.assertEqual(inputs["mask_labels"][1].shape, (61, 512, 711))
-        self.assertEqual(inputs["mask_labels"][0].sum().item(), 315193.0)
-        self.assertEqual(inputs["mask_labels"][1].sum().item(), 350747.0)
+            # verify the pixel values and pixel mask
+            self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 711))
+            self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 711))
+
+            # verify the class labels
+            self.assertEqual(len(inputs["class_labels"]), 2)
+            expected_class_labels = torch.tensor([4, 17, 32, 42, 42, 42, 42, 42, 42, 42, 32, 12, 12, 12, 12, 12, 42, 42, 12, 12, 12, 42, 12, 12, 12, 12, 12, 3, 12, 12, 12, 12, 42, 42, 42, 12, 42, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 5, 12, 12, 12, 12, 12, 12, 12, 0, 43, 43, 43, 96, 43, 104, 43, 31, 125, 31, 125, 138, 87, 125, 149, 138, 125, 87, 87])  # fmt: skip
+            torch.testing.assert_close(inputs["class_labels"][0], torch.tensor(expected_class_labels))
+            expected_class_labels = torch.tensor([19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 67, 82, 19, 19, 17, 19, 19, 19, 19, 19, 19, 19, 19, 19, 12, 12, 42, 12, 12, 12, 12, 3, 14, 12, 12, 12, 12, 12, 12, 12, 12, 14, 5, 12, 12, 0, 115, 43, 43, 115, 43, 43, 43, 8, 8, 8, 138, 138, 125, 143])  # fmt: skip
+            torch.testing.assert_close(inputs["class_labels"][1], expected_class_labels)
+
+            # verify the mask labels
+            self.assertEqual(len(inputs["mask_labels"]), 2)
+            self.assertEqual(inputs["mask_labels"][0].shape, (79, 512, 711))
+            self.assertEqual(inputs["mask_labels"][1].shape, (61, 512, 711))
+            self.assertEqual(inputs["mask_labels"][0].sum().item(), 315193.0)
+            self.assertEqual(inputs["mask_labels"][1].sum().item(), 350747.0)
 
     def test_binary_mask_to_rle(self):
         fake_binary_mask = np.zeros((20, 50))

filename: tests/models/maskformer/test_image_processing_maskformer.py
patch: @@ -188,9 +188,9 @@ def test_image_processor_properties(self):
             self.assertTrue(hasattr(image_processing, "num_labels"))
 
     def comm_get_image_processing_inputs(
-        self, with_segmentation_maps=False, is_instance_map=False, segmentation_type="np"
+        self, image_processing_class, with_segmentation_maps=False, is_instance_map=False, segmentation_type="np"
     ):
-        image_processing = self.image_processing_class(**self.image_processor_dict)
+        image_processing = image_processing_class(**self.image_processor_dict)
         # prepare image and target
         num_labels = self.image_processor_tester.num_labels
         annotations = None
@@ -212,7 +212,6 @@ def comm_get_image_processing_inputs(
             annotations,
             return_tensors="pt",
             instance_id_to_semantic_id=instance_id_to_semantic_id,
-            pad_and_return_pixel_mask=True,
         )
 
         return inputs
@@ -233,19 +232,23 @@ def test_with_size_divisor(self):
 
     def test_call_with_segmentation_maps(self):
         def common(is_instance_map=False, segmentation_type=None):
-            inputs = self.comm_get_image_processing_inputs(
-                with_segmentation_maps=True, is_instance_map=is_instance_map, segmentation_type=segmentation_type
-            )
+            for image_processing_class in self.image_processor_list:
+                inputs = self.comm_get_image_processing_inputs(
+                    image_processing_class=image_processing_class,
+                    with_segmentation_maps=True,
+                    is_instance_map=is_instance_map,
+                    segmentation_type=segmentation_type,
+                )
 
-            mask_labels = inputs["mask_labels"]
-            class_labels = inputs["class_labels"]
-            pixel_values = inputs["pixel_values"]
+                mask_labels = inputs["mask_labels"]
+                class_labels = inputs["class_labels"]
+                pixel_values = inputs["pixel_values"]
 
-            # check the batch_size
-            for mask_label, class_label in zip(mask_labels, class_labels):
-                self.assertEqual(mask_label.shape[0], class_label.shape[0])
-                # this ensure padding has happened
-                self.assertEqual(mask_label.shape[1:], pixel_values.shape[2:])
+                # check the batch_size
+                for mask_label, class_label in zip(mask_labels, class_labels):
+                    self.assertEqual(mask_label.shape[0], class_label.shape[0])
+                    # this ensure padding has happened
+                    self.assertEqual(mask_label.shape[1:], pixel_values.shape[2:])
 
         common()
         common(is_instance_map=True)
@@ -286,31 +289,32 @@ def get_instance_segmentation_and_mapping(annotation):
         instance_seg2, inst2class2 = get_instance_segmentation_and_mapping(annotation2)
 
         # create a image processor
-        image_processing = MaskFormerImageProcessor(do_reduce_labels=True, ignore_index=255, size=(512, 512))
-
-        # prepare the images and annotations
-        inputs = image_processing(
-            [image1, image2],
-            [instance_seg1, instance_seg2],
-            instance_id_to_semantic_id=[inst2class1, inst2class2],
-            return_tensors="pt",
-        )
+        for image_processing_class in self.image_processor_list:
+            image_processing = image_processing_class(do_reduce_labels=True, ignore_index=255, size=(512, 512))
+
+            # prepare the images and annotations
+            inputs = image_processing(
+                [image1, image2],
+                [instance_seg1, instance_seg2],
+                instance_id_to_semantic_id=[inst2class1, inst2class2],
+                return_tensors="pt",
+            )
 
-        # verify the pixel values and pixel mask
-        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
-        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
+            # verify the pixel values and pixel mask
+            self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
+            self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
 
-        # verify the class labels
-        self.assertEqual(len(inputs["class_labels"]), 2)
-        torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([30, 55]))
-        torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([4, 4, 23, 55]))
+            # verify the class labels
+            self.assertEqual(len(inputs["class_labels"]), 2)
+            torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([30, 55]))
+            torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([4, 4, 23, 55]))
 
-        # verify the mask labels
-        self.assertEqual(len(inputs["mask_labels"]), 2)
-        self.assertEqual(inputs["mask_labels"][0].shape, (2, 512, 512))
-        self.assertEqual(inputs["mask_labels"][1].shape, (4, 512, 512))
-        self.assertEqual(inputs["mask_labels"][0].sum().item(), 41527.0)
-        self.assertEqual(inputs["mask_labels"][1].sum().item(), 26259.0)
+            # verify the mask labels
+            self.assertEqual(len(inputs["mask_labels"]), 2)
+            self.assertEqual(inputs["mask_labels"][0].shape, (2, 512, 512))
+            self.assertEqual(inputs["mask_labels"][1].shape, (4, 512, 512))
+            self.assertEqual(inputs["mask_labels"][0].sum().item(), 41527.0)
+            self.assertEqual(inputs["mask_labels"][1].sum().item(), 26259.0)
 
     def test_integration_semantic_segmentation(self):
         # load 2 images and corresponding semantic annotations from the hub
@@ -329,30 +333,31 @@ def test_integration_semantic_segmentation(self):
         )
 
         # create a image processor
-        image_processing = MaskFormerImageProcessor(do_reduce_labels=True, ignore_index=255, size=(512, 512))
+        for image_processing_class in self.image_processor_list:
+            image_processing = image_processing_class(do_reduce_labels=True, ignore_index=255, size=(512, 512))
 
-        # prepare the images and annotations
-        inputs = image_processing(
-            [image1, image2],
-            [annotation1, annotation2],
-            return_tensors="pt",
-        )
+            # prepare the images and annotations
+            inputs = image_processing(
+                [image1, image2],
+                [annotation1, annotation2],
+                return_tensors="pt",
+            )
 
-        # verify the pixel values and pixel mask
-        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
-        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
+            # verify the pixel values and pixel mask
+            self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 512))
+            self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 512))
 
-        # verify the class labels
-        self.assertEqual(len(inputs["class_labels"]), 2)
-        torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([2, 4, 60]))
-        torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([0, 3, 7, 8, 15, 28, 30, 143]))
+            # verify the class labels
+            self.assertEqual(len(inputs["class_labels"]), 2)
+            torch.testing.assert_close(inputs["class_labels"][0], torch.tensor([2, 4, 60]))
+            torch.testing.assert_close(inputs["class_labels"][1], torch.tensor([0, 3, 7, 8, 15, 28, 30, 143]))
 
-        # verify the mask labels
-        self.assertEqual(len(inputs["mask_labels"]), 2)
-        self.assertEqual(inputs["mask_labels"][0].shape, (3, 512, 512))
-        self.assertEqual(inputs["mask_labels"][1].shape, (8, 512, 512))
-        self.assertEqual(inputs["mask_labels"][0].sum().item(), 170200.0)
-        self.assertEqual(inputs["mask_labels"][1].sum().item(), 257036.0)
+            # verify the mask labels
+            self.assertEqual(len(inputs["mask_labels"]), 2)
+            self.assertEqual(inputs["mask_labels"][0].shape, (3, 512, 512))
+            self.assertEqual(inputs["mask_labels"][1].shape, (8, 512, 512))
+            self.assertEqual(inputs["mask_labels"][0].sum().item(), 170200.0)
+            self.assertEqual(inputs["mask_labels"][1].sum().item(), 257036.0)
 
     def test_integration_panoptic_segmentation(self):
         # load 2 images and corresponding panoptic annotations from the hub
@@ -386,34 +391,35 @@ def create_panoptic_map(annotation, segments_info):
         panoptic_map2, inst2class2 = create_panoptic_map(annotation2, segments_info2)
 
         # create a image processor
-        image_processing = MaskFormerImageProcessor(ignore_index=0, do_resize=False)
-
-        # prepare the images and annotations
-        pixel_values_list = [np.moveaxis(np.array(image1), -1, 0), np.moveaxis(np.array(image2), -1, 0)]
-        inputs = image_processing.encode_inputs(
-            pixel_values_list,
-            [panoptic_map1, panoptic_map2],
-            instance_id_to_semantic_id=[inst2class1, inst2class2],
-            return_tensors="pt",
-        )
+        for image_processing_class in self.image_processor_list:
+            image_processing = image_processing_class(ignore_index=0, do_resize=False)
+
+            # prepare the images and annotations
+            pixel_values_list = [np.moveaxis(np.array(image1), -1, 0), np.moveaxis(np.array(image2), -1, 0)]
+            inputs = image_processing(
+                pixel_values_list,
+                [panoptic_map1, panoptic_map2],
+                instance_id_to_semantic_id=[inst2class1, inst2class2],
+                return_tensors="pt",
+            )
 
-        # verify the pixel values and pixel mask
-        self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 711))
-        self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 711))
-
-        # verify the class labels
-        self.assertEqual(len(inputs["class_labels"]), 2)
-        expected_class_labels = torch.tensor([4, 17, 32, 42, 42, 42, 42, 42, 42, 42, 32, 12, 12, 12, 12, 12, 42, 42, 12, 12, 12, 42, 12, 12, 12, 12, 12, 3, 12, 12, 12, 12, 42, 42, 42, 12, 42, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 5, 12, 12, 12, 12, 12, 12, 12, 0, 43, 43, 43, 96, 43, 104, 43, 31, 125, 31, 125, 138, 87, 125, 149, 138, 125, 87, 87])  # fmt: skip
-        torch.testing.assert_close(inputs["class_labels"][0], torch.tensor(expected_class_labels))
-        expected_class_labels = torch.tensor([19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 67, 82, 19, 19, 17, 19, 19, 19, 19, 19, 19, 19, 19, 19, 12, 12, 42, 12, 12, 12, 12, 3, 14, 12, 12, 12, 12, 12, 12, 12, 12, 14, 5, 12, 12, 0, 115, 43, 43, 115, 43, 43, 43, 8, 8, 8, 138, 138, 125, 143])  # fmt: skip
-        torch.testing.assert_close(inputs["class_labels"][1], expected_class_labels)
-
-        # verify the mask labels
-        self.assertEqual(len(inputs["mask_labels"]), 2)
-        self.assertEqual(inputs["mask_labels"][0].shape, (79, 512, 711))
-        self.assertEqual(inputs["mask_labels"][1].shape, (61, 512, 711))
-        self.assertEqual(inputs["mask_labels"][0].sum().item(), 315193.0)
-        self.assertEqual(inputs["mask_labels"][1].sum().item(), 350747.0)
+            # verify the pixel values and pixel mask
+            self.assertEqual(inputs["pixel_values"].shape, (2, 3, 512, 711))
+            self.assertEqual(inputs["pixel_mask"].shape, (2, 512, 711))
+
+            # verify the class labels
+            self.assertEqual(len(inputs["class_labels"]), 2)
+            expected_class_labels = torch.tensor([4, 17, 32, 42, 42, 42, 42, 42, 42, 42, 32, 12, 12, 12, 12, 12, 42, 42, 12, 12, 12, 42, 12, 12, 12, 12, 12, 3, 12, 12, 12, 12, 42, 42, 42, 12, 42, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 5, 12, 12, 12, 12, 12, 12, 12, 0, 43, 43, 43, 96, 43, 104, 43, 31, 125, 31, 125, 138, 87, 125, 149, 138, 125, 87, 87])  # fmt: skip
+            torch.testing.assert_close(inputs["class_labels"][0], torch.tensor(expected_class_labels))
+            expected_class_labels = torch.tensor([19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 67, 82, 19, 19, 17, 19, 19, 19, 19, 19, 19, 19, 19, 19, 12, 12, 42, 12, 12, 12, 12, 3, 14, 12, 12, 12, 12, 12, 12, 12, 12, 14, 5, 12, 12, 0, 115, 43, 43, 115, 43, 43, 43, 8, 8, 8, 138, 138, 125, 143])  # fmt: skip
+            torch.testing.assert_close(inputs["class_labels"][1], expected_class_labels)
+
+            # verify the mask labels
+            self.assertEqual(len(inputs["mask_labels"]), 2)
+            self.assertEqual(inputs["mask_labels"][0].shape, (79, 512, 711))
+            self.assertEqual(inputs["mask_labels"][1].shape, (61, 512, 711))
+            self.assertEqual(inputs["mask_labels"][0].sum().item(), 315193.0)
+            self.assertEqual(inputs["mask_labels"][1].sum().item(), 350747.0)
 
     def test_binary_mask_to_rle(self):
         fake_binary_mask = np.zeros((20, 50))