Add convert_segmentation_map_to_binary_masks_sorted function for hand…

src/transformers/models/mask2former/image_processing_mask2former.py

@@ -343,6 +343,113 @@ def convert_segmentation_map_to_binary_masks(
     return binary_masks.astype(np.float32), labels.astype(np.int64)
 
 
+def convert_segmentation_map_to_binary_masks_sorted(
+    segmentation_map: "np.ndarray",
+    instance_id_to_semantic_id: Optional[dict[int, int]] = None,
+    ignore_index: Optional[int] = None,
+    do_reduce_labels: bool = False,
+    sort_by_area: bool = True,
+):
+    """
+    Converts a segmentation map into binary masks with optional area-based sorting.
+
+    This function extends the base `convert_segmentation_map_to_binary_masks` by adding
+    area-based sorting capability. When `sort_by_area=True`, masks are sorted by area
+    in descending order (largest first), which is particularly useful for handling
+    overlapping annotations in segmentation tasks.
+
+    Args:
+        segmentation_map (`np.ndarray`):
+            A 2D numpy array representing the segmentation map where each unique value
+            corresponds to a different object instance or class.
+        instance_id_to_semantic_id (`Dict[int, int]`, *optional*):
+            A dictionary mapping instance IDs to semantic class IDs. If provided,
+            the returned labels will be the semantic class IDs instead of instance IDs.
+        ignore_index (`int`, *optional*):
+            The index to ignore when creating binary masks. Pixels with this value
+            will not be included in any mask.
+        do_reduce_labels (`bool`, *optional*, defaults to `False`):
+            Whether to reduce labels by subtracting 1 from all non-zero labels.
+            If `True`, `ignore_index` must be provided.
+        sort_by_area (`bool`, *optional*, defaults to `True`):
+            Whether to sort the binary masks by area in descending order.
+    Returns:
+        `Tuple[np.ndarray, np.ndarray]`:
+            - **binary_masks** (`np.ndarray` of shape `(num_masks, height, width)`):
+              Binary masks for each object, optionally sorted by area.
+            - **labels** (`np.ndarray` of shape `(num_masks,)`):
+              Class or instance labels corresponding to each mask, in the same order as the masks.
+    Example:
+        ```python
+        import numpy as np
+        from transformers.models.mask2former.image_processing_mask2former import convert_segmentation_map_to_binary_masks_sorted
+
+        # Create a simple segmentation map with overlapping objects
+        segmentation_map = np.array([
+            [0, 1, 1, 2],
+            [0, 1, 2, 2],
+            [0, 0, 2, 2],
+        ])
+
+        # Convert to binary masks sorted by area
+        masks, labels = convert_segmentation_map_to_binary_masks_sorted(
+            segmentation_map,
+            ignore_index=0,
+            sort_by_area=True
+        )
+
+        # Object 2 has larger area (6 pixels) so it comes first
+        # Object 1 has smaller area (3 pixels) so it comes second
+        print(labels)  # [2, 1] - sorted by area (largest first)
+        ```
+    """
+    if do_reduce_labels and ignore_index is None:
+        raise ValueError("If `do_reduce_labels` is True, `ignore_index` must be provided.")
+
+    if do_reduce_labels:
+        segmentation_map = np.where(segmentation_map == 0, ignore_index, segmentation_map - 1)
+
+    # Get unique ids (class or instance ids based on input)
+    all_labels = np.unique(segmentation_map)
+
+    # Drop background label if applicable
+    if ignore_index is not None:
+        all_labels = all_labels[all_labels != ignore_index]
+
+    # Generate a binary mask for each object instance
+    binary_masks = [(segmentation_map == i) for i in all_labels]
+
+    # Stack the binary masks
+    if binary_masks:
+        binary_masks = np.stack(binary_masks, axis=0)
+    else:
+        binary_masks = np.zeros((0, *segmentation_map.shape))
+
+    # Convert instance ids to class ids
+    if instance_id_to_semantic_id is not None:
+        labels = np.zeros(all_labels.shape[0])
+
+        for label in all_labels:
+            class_id = instance_id_to_semantic_id[label + 1 if do_reduce_labels else label]
+            labels[all_labels == label] = class_id - 1 if do_reduce_labels else class_id
+    else:
+        labels = all_labels
+
+    # Sort by area if requested
+    if sort_by_area and len(binary_masks) > 0:
+        # Calculate areas for each mask
+        areas = np.sum(binary_masks, axis=(1, 2))
+
+        # Sort indices by area in descending order (largest first)
+        sort_indices = np.argsort(areas)[::-1]
+
+        # Reorder masks and labels
+        binary_masks = binary_masks[sort_indices]
+        labels = labels[sort_indices]
+
+    return binary_masks.astype(np.float32), labels.astype(np.int64)
+
+
 # Copied from transformers.models.maskformer.image_processing_maskformer.get_maskformer_resize_output_image_size with maskformer->mask2former
 def get_mask2former_resize_output_image_size(
     image: np.ndarray,

tests/models/mask2former/test_image_processing_mask2former.py

@@ -31,7 +31,10 @@
 
     if is_vision_available():
         from transformers import Mask2FormerImageProcessor
-        from transformers.models.mask2former.image_processing_mask2former import binary_mask_to_rle
+        from transformers.models.mask2former.image_processing_mask2former import (
+            binary_mask_to_rle,
+            convert_segmentation_map_to_binary_masks_sorted,
+        )
         from transformers.models.mask2former.modeling_mask2former import Mask2FormerForUniversalSegmentationOutput
 
         if is_torchvision_available():
@@ -603,3 +606,72 @@ def test_slow_fast_equivalence_batched(self):
             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())
+
+    def test_convert_segmentation_map_to_binary_masks_sorted(self):
+        segmentation_map = np.array(
+            [
+                [0, 1, 1, 2, 2, 2],
+                [0, 1, 2, 2, 2, 2],
+                [0, 0, 2, 2, 3, 3],
+            ]
+        )
+
+        masks, labels = convert_segmentation_map_to_binary_masks_sorted(
+            segmentation_map, sort_by_area=True, ignore_index=0
+        )
+
+        self.assertEqual(masks.shape[0], 3)
+        self.assertEqual(len(labels), 3)
+        self.assertEqual(labels[0], 2)
+
+        areas = np.sum(masks, axis=(1, 2))
+        self.assertTrue(all(areas[i] >= areas[i + 1] for i in range(len(areas) - 1)))
+
+    def test_convert_segmentation_map_to_binary_masks_sorted_edge_cases(self):
+        empty_map = np.zeros((3, 3), dtype=np.uint8)
+        masks, labels = convert_segmentation_map_to_binary_masks_sorted(empty_map, ignore_index=0)
+        self.assertEqual(masks.shape[0], 0)
+        self.assertEqual(len(labels), 0)
+
+        single_object_map = np.ones((3, 3), dtype=np.uint8)
+        masks, labels = convert_segmentation_map_to_binary_masks_sorted(single_object_map, ignore_index=0)
+        self.assertEqual(masks.shape[0], 1)
+        self.assertEqual(labels[0], 1)
+
+    def test_convert_segmentation_map_to_binary_masks_sorted_instance_mapping(self):
+        segmentation_map = np.array(
+            [
+                [0, 1, 1, 2],
+                [0, 1, 2, 2],
+                [0, 0, 2, 2],
+            ]
+        )
+
+        instance_id_to_semantic_id = {1: 10, 2: 20}
+
+        masks, labels = convert_segmentation_map_to_binary_masks_sorted(
+            segmentation_map, instance_id_to_semantic_id=instance_id_to_semantic_id, ignore_index=0, sort_by_area=True
+        )
+
+        self.assertEqual(len(labels), 2)
+        self.assertEqual(labels[0], 20)
+        self.assertEqual(labels[1], 10)
+
+    def test_convert_segmentation_map_to_binary_masks_sorted_area_ordering(self):
+        segmentation_map = np.array(
+            [
+                [1, 1, 2, 3, 3, 3],
+                [1, 1, 2, 3, 3, 3],
+                [0, 0, 2, 3, 3, 3],
+            ]
+        )
+
+        masks, labels = convert_segmentation_map_to_binary_masks_sorted(
+            segmentation_map, sort_by_area=True, ignore_index=0
+        )
+
+        areas = [np.sum(mask) for mask in masks]
+        expected_order = [3, 1, 2]
+
+        self.assertEqual(list(labels), expected_order)
+        self.assertTrue(areas[0] >= areas[1] >= areas[2])