Add ignore_index argument to Mask.as_class_mask() and Mask.as_instanc…

…e_mask() (#1409)

### Summary

- Ticket no. 137625
- Same as title.

### How to test
Added an unit test for this change.

### Checklist
<!-- Put an 'x' in all the boxes that apply -->
- [x] I have added unit tests to cover my changes.​
- [ ] I have added integration tests to cover my changes.​
- [x] I have added the description of my changes into
[CHANGELOG](https://github.com/openvinotoolkit/datumaro/blob/develop/CHANGELOG.md).​
- [x] I have updated the
[documentation](https://github.com/openvinotoolkit/datumaro/tree/develop/docs)
accordingly

### License

- [x] I submit _my code changes_ under the same [MIT
License](https://github.com/openvinotoolkit/datumaro/blob/develop/LICENSE)
that covers the project.
  Feel free to contact the maintainers if that's a concern.
- [x] I have updated the license header for each file (see an example
below).

```python
# Copyright (C) 2024 Intel Corporation
#
# SPDX-License-Identifier: MIT
```

---------

Signed-off-by: Kim, Vinnam <[email protected]>

CHANGELOG.md

@@ -5,6 +5,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Apr. 2024 Release 2.1.0
+### New features
+
+### Enhancements
+- Add ignore_index argument to Mask.as_class_mask() and Mask.as_instance_mask()
+  (<https://github.com/openvinotoolkit/datumaro/pull/1409>)
+
+### Bug fixes
+
 ## Apr. 2024 Release 2.0.0
 ### New features
 - Changed supported Python version range (>=3.9, <=3.11)

src/datumaro/components/annotation.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2022 Intel Corporation
+# Copyright (C) 2021-2024 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
 
@@ -370,23 +370,58 @@ def image(self) -> BinaryMaskImage:
             image = image()
         return image
 
-    def as_class_mask(self, label_id: Optional[int] = None) -> IndexMaskImage:
-        """
-        Produces a class index mask. Mask label id can be changed.
+    def as_class_mask(
+        self,
+        label_id: Optional[int] = None,
+        ignore_index: int = 0,
+        dtype: Optional[np.dtype] = None,
+    ) -> IndexMaskImage:
+        """Produces a class index mask based on the binary mask.
+
+        Args:
+            label_id: Scalar value to represent the class index of the mask.
+                If not specified, `self.label` will be used. Defaults to None.
+            ignore_index: Scalar value to fill in the zeros in the binary mask.
+                Defaults to 0.
+            dtype: Data type for the resulting mask. If not specified,
+                it will be inferred from the provided `label_id` to hold its value.
+                For example, if `label_id=255`, the inferred dtype will be `np.uint8`.
+                Defaults to None.
+
+        Returns:
+            IndexMaskImage: Class index mask generated from the binary mask.
         """
         if label_id is None:
             label_id = self.label
         from datumaro.util.mask_tools import make_index_mask
 
-        return make_index_mask(self.image, label_id)
+        return make_index_mask(self.image, index=label_id, ignore_index=ignore_index, dtype=dtype)
 
-    def as_instance_mask(self, instance_id: int) -> IndexMaskImage:
-        """
-        Produces a instance index mask.
+    def as_instance_mask(
+        self,
+        instance_id: int,
+        ignore_index: int = 0,
+        dtype: Optional[np.dtype] = None,
+    ) -> IndexMaskImage:
+        """Produces an instance index mask based on the binary mask.
+
+        Args:
+            instance_id: Scalar value to represent the instance id.
+            ignore_index: Scalar value to fill in the zeros in the binary mask.
+                Defaults to 0.
+            dtype: Data type for the resulting mask. If not specified,
+                it will be inferred from the provided `label_id` to hold its value.
+                For example, if `label_id=255`, the inferred dtype will be `np.uint8`.
+                Defaults to None.
+
+        Returns:
+            IndexMaskImage: Instance index mask generated from the binary mask.
         """
         from datumaro.util.mask_tools import make_index_mask
 
-        return make_index_mask(self.image, instance_id)
+        return make_index_mask(
+            self.image, index=instance_id, ignore_index=ignore_index, dtype=dtype
+        )
 
     def get_area(self) -> int:
         return np.count_nonzero(self.image)

src/datumaro/util/mask_tools.py

@@ -1,10 +1,10 @@
-# Copyright (C) 2019-2021 Intel Corporation
+# Copyright (C) 2019-2024 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
-
+import logging as log
 from functools import partial
 from itertools import chain
-from typing import Dict, Tuple
+from typing import Dict, Optional, Tuple
 
 import numpy as np
 from pycocotools import mask as pycocotools_mask
@@ -124,8 +124,72 @@ def remap_mask(mask, map_fn):
     return np.array([map_fn(c) for c in range(256)], dtype=np.uint8)[mask]
 
 
-def make_index_mask(binary_mask, index, dtype=None):
-    return binary_mask * np.array([index], dtype=dtype or np.min_scalar_type(index))
+def make_index_mask(
+    binary_mask: np.ndarray,
+    index: int,
+    ignore_index: int = 0,
+    dtype: Optional[np.dtype] = None,
+):
+    """Create an index mask from a binary mask by filling a given index value.
+
+    Args:
+        binary_mask: Binary mask to create an index mask.
+        index: Scalar value to fill the ones in the binary mask.
+        ignore_index: Scalar value to fill in the zeros in the binary mask.
+            Defaults to 0.
+        dtype: Data type for the resulting mask. If not specified,
+                it will be inferred from the provided `index` to hold its value.
+                For example, if `index=255`, the inferred dtype will be `np.uint8`.
+                Defaults to None.
+
+    Returns:
+        np.ndarray: Index mask created from the binary mask.
+
+    Raises:
+        ValueError: If dtype is not specified and incompatible scalar types are used for index
+            and ignore_index.
+
+    Examples:
+        >>> binary_mask = np.eye(2, dtype=np.bool_)
+        >>> index_mask = make_index_mask(binary_mask, index=10, ignore_index=255, dtype=np.uint8)
+        >>> print(index_mask)
+        array([[ 10, 255],
+               [255,  10]], dtype=uint8)
+    """
+    if dtype is None:
+        dtype = np.min_scalar_type(index)
+        if dtype != np.min_scalar_type(ignore_index):
+            msg = (
+                "Given dtype is None, "
+                "but inferred dtypes from the given index and ignore_index are different each other. "
+                "Please mannually set dtype"
+            )
+            raise ValueError(msg, index, ignore_index)
+
+    flipped_zero_np_scalar = ~np.full(tuple(), fill_value=0, dtype=dtype)
+
+    # NOTE: This dispatching rule is required for a performance boost
+    if ignore_index == flipped_zero_np_scalar:
+        flipped_index = ~np.full(tuple(), fill_value=index, dtype=dtype)
+        return ~(binary_mask * flipped_index)
+    elif index < ignore_index:
+        diff = ignore_index - index
+        mask = ~binary_mask * np.full(tuple(), fill_value=diff, dtype=dtype)
+        mask += index
+        return mask
+    elif index > ignore_index:
+        diff = index - ignore_index
+        mask = binary_mask * np.full(tuple(), fill_value=diff, dtype=dtype)
+        mask += ignore_index
+        return mask
+
+    # index == ignore_index
+    msg = (
+        "index == ignore_index. "
+        f"It will create an index mask filling with a single value, index={index}"
+    )
+    log.warning(msg)
+    return np.full_like(binary_mask, fill_value=index, dtype=dtype)
 
 
 def make_binary_mask(mask):

tests/unit/test_masks.py

@@ -1,3 +1,7 @@
+# Copyright (C) 2019-2024 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
 from unittest import TestCase
 
 import numpy as np
@@ -223,3 +227,43 @@ def test_can_decode_compiled_mask(self):
         labels = compiled_mask.get_instance_labels()
 
         self.assertEqual({instance_idx: class_idx}, labels)
+
+
+class MaskToolsTest:
+    """New test implementation based on PyTest framework.
+
+    The other tests in this file should be also migrated into this test class.
+    """
+
+    def test_make_index_mask(self):
+        binary_mask = np.eye(2, dtype=np.bool_)
+
+        def _test(expected, actual):
+            assert np.allclose(expected, actual) and actual.dtype == expected.dtype
+
+        _test(
+            np.array([[10, 0], [0, 10]], dtype=np.uint8),
+            mask_tools.make_index_mask(binary_mask=binary_mask, index=10, ignore_index=0),
+        )
+
+        _test(
+            np.array([[10, 255], [255, 10]], dtype=np.uint8),
+            mask_tools.make_index_mask(binary_mask=binary_mask, index=10, ignore_index=255),
+        )
+
+        _test(
+            np.array([[10, 100], [100, 10]], dtype=np.uint8),
+            mask_tools.make_index_mask(binary_mask=binary_mask, index=10, ignore_index=100),
+        )
+
+        _test(
+            np.array([[200, 100], [100, 200]], dtype=np.uint8),
+            mask_tools.make_index_mask(binary_mask=binary_mask, index=200, ignore_index=100),
+        )
+
+        _test(
+            np.array([[10, 65535], [65535, 10]], dtype=np.uint16),
+            mask_tools.make_index_mask(
+                binary_mask=binary_mask, index=10, ignore_index=65535, dtype=np.uint16
+            ),
+        )