diff --git a/mmdet/structures/mask/structures.py b/mmdet/structures/mask/structures.py index 8619d91dd2f..132dc45ece8 100644 --- a/mmdet/structures/mask/structures.py +++ b/mmdet/structures/mask/structures.py @@ -422,7 +422,6 @@ def translate(self, >>> border_value = 0 >>> interpolation = 'bilinear' >>> # Note, There seem to be issues when: - >>> # * out_shape is different than self's shape >>> # * the mask dtype is not supported by cv2.AffineWarp >>> new = self.translate(out_shape, offset, direction, >>> border_value, interpolation) @@ -432,8 +431,16 @@ def translate(self, if len(self.masks) == 0: translated_masks = np.empty((0, *out_shape), dtype=np.uint8) else: + masks = self.masks + if masks.shape[-2:] != out_shape: + empty_masks = np.zeros((masks.shape[0], *out_shape), + dtype=masks.dtype) + min_h = min(out_shape[0], masks.shape[1]) + min_w = min(out_shape[1], masks.shape[2]) + empty_masks[:, :min_h, :min_w] = masks[:, :min_h, :min_w] + masks = empty_masks translated_masks = mmcv.imtranslate( - self.masks.transpose((1, 2, 0)), + masks.transpose((1, 2, 0)), offset, direction, border_value=border_value, diff --git a/tests/test_structures/test_mask/test_mask_structures.py b/tests/test_structures/test_mask/test_mask_structures.py new file mode 100644 index 00000000000..d3596a512de --- /dev/null +++ b/tests/test_structures/test_mask/test_mask_structures.py @@ -0,0 +1,43 @@ +from unittest import TestCase + +import numpy as np +from mmengine.testing import assert_allclose + +from mmdet.structures.mask import BitmapMasks + + +class TestMaskStructures(TestCase): + + def test_bitmap_translate_same_size(self): + mask_array = np.zeros((5, 10, 10), dtype=np.uint8) + mask_array[:, 0:5, 0:5] = 1 + mask_target = np.zeros((5, 10, 10), dtype=np.uint8) + mask_target[:, 0:5, 5:10] = 1 + + mask = BitmapMasks(mask_array, 10, 10) + mask = mask.translate((10, 10), 5) + assert mask.masks.shape == (5, 10, 10) + assert_allclose(mask_target, mask.masks) + + def test_bitmap_translate_diff_size(self): + # test out shape larger + mask_array = np.zeros((5, 10, 10), dtype=np.uint8) + mask_array[:, 0:5, 0:5] = 1 + + mask_target = np.zeros((5, 20, 20), dtype=np.uint8) + mask_target[:, 0:5, 5:10] = 1 + mask = BitmapMasks(mask_array, 10, 10) + mask = mask.translate((20, 20), 5) + assert mask.masks.shape == (5, 20, 20) + assert_allclose(mask_target, mask.masks) + + # test out shape smaller + mask_array = np.zeros((5, 10, 10), dtype=np.uint8) + mask_array[:, 0:5, 0:5] = 1 + + mask_target = np.zeros((5, 20, 8), dtype=np.uint8) + mask_target[:, 0:5, 5:] = 1 + mask = BitmapMasks(mask_array, 10, 10) + mask = mask.translate((20, 8), 5) + assert mask.masks.shape == (5, 20, 8) + assert_allclose(mask_target, mask.masks)