From 1bd2296c4b8fd7c987d99b8406b0c741719350fe Mon Sep 17 00:00:00 2001 From: Yusuke Niitani Date: Tue, 19 Feb 2019 15:21:42 +0900 Subject: [PATCH] add scale_mask --- chainercv/utils/__init__.py | 1 + chainercv/utils/mask/scale_mask.py | 70 +++++++++++++++++ docs/source/reference/utils.rst | 4 + .../utils_tests/mask_tests/test_scale_mask.py | 78 +++++++++++++++++++ 4 files changed, 153 insertions(+) create mode 100644 chainercv/utils/mask/scale_mask.py create mode 100644 tests/utils_tests/mask_tests/test_scale_mask.py diff --git a/chainercv/utils/__init__.py b/chainercv/utils/__init__.py index e13ac7f1c6..4573ec37c6 100644 --- a/chainercv/utils/__init__.py +++ b/chainercv/utils/__init__.py @@ -13,6 +13,7 @@ from chainercv.utils.link import prepare_pretrained_model # NOQA from chainercv.utils.mask.mask_iou import mask_iou # NOQA from chainercv.utils.mask.mask_to_bbox import mask_to_bbox # NOQA +from chainercv.utils.mask.scale_mask import scale_mask # NOQA from chainercv.utils.testing import assert_is_bbox # NOQA from chainercv.utils.testing import assert_is_bbox_dataset # NOQA from chainercv.utils.testing import assert_is_detection_link # NOQA diff --git a/chainercv/utils/mask/scale_mask.py b/chainercv/utils/mask/scale_mask.py new file mode 100644 index 0000000000..72e91b1604 --- /dev/null +++ b/chainercv/utils/mask/scale_mask.py @@ -0,0 +1,70 @@ +from __future__ import division + +import numpy as np +import PIL.Image + +import chainer +from chainercv import transforms + + +def scale_mask(mask, bbox, size): + """Scale instance segmentation mask while keeping the aspect ratio. + + This function exploits the sparsity of :obj:`mask` to speed up + resize operation. + + The input image will be resized so that + the shorter edge will be scaled to length :obj:`size` after + resizing. + + Args: + mask (array): An array whose shape is :math:`(R, H, W)`. + :math:`R` is the number of masks. + The dtype should be :obj:`numpy.bool`. + bbox (array): The bounding boxes around the masked region + of :obj:`mask`. This is expected to be the value + obtained by :obj:`bbox = chainercv.utils.mask_to_bbox(mask)`. + size (int): The length of the smaller edge. + + Returns: + array: + An array whose shape is :math:`(R, H, W)`. + :math:`R` is the number of masks. + The dtype should be :obj:`numpy.bool`. + + """ + xp = chainer.backends.cuda.get_array_module(mask) + mask = chainer.cuda.to_cpu(mask) + bbox = chainer.cuda.to_cpu(bbox) + + R, H, W = mask.shape + if H < W: + out_size = (size, int(size * W / H)) + scale = size / H + else: + out_size = (int(size * H / W), size) + scale = size / W + + bbox[:, :2] = np.floor(bbox[:, :2]) + bbox[:, 2:] = np.ceil(bbox[:, 2:]) + bbox = bbox.astype(np.int32) + scaled_bbox = bbox * scale + scaled_bbox[:, :2] = np.floor(scaled_bbox[:, :2]) + scaled_bbox[:, 2:] = np.ceil(scaled_bbox[:, 2:]) + scaled_bbox = scaled_bbox.astype(np.int32) + + out_mask = xp.zeros((R,) + out_size, dtype=np.bool) + for i, (m, bb, scaled_bb) in enumerate( + zip(mask, bbox, scaled_bbox)): + cropped_m = m[bb[0]:bb[2], bb[1]:bb[3]] + h = scaled_bb[2] - scaled_bb[0] + w = scaled_bb[3] - scaled_bb[1] + cropped_m = transforms.resize( + cropped_m[None].astype(np.float32), + (h, w), + interpolation=PIL.Image.NEAREST)[0] + if xp != np: + cropped_m = xp.array(cropped_m) + out_mask[i, scaled_bb[0]:scaled_bb[2], + scaled_bb[1]:scaled_bb[3]] = cropped_m + return out_mask diff --git a/docs/source/reference/utils.rst b/docs/source/reference/utils.rst index 9378a583d5..d90db325ec 100644 --- a/docs/source/reference/utils.rst +++ b/docs/source/reference/utils.rst @@ -87,6 +87,10 @@ mask_to_bbox ~~~~~~~~~~~~ .. autofunction:: mask_to_bbox +scale_mask +~~~~~~~~~~ +.. autofunction:: scale_mask + Testing Utilities ----------------- diff --git a/tests/utils_tests/mask_tests/test_scale_mask.py b/tests/utils_tests/mask_tests/test_scale_mask.py new file mode 100644 index 0000000000..b71a4405e7 --- /dev/null +++ b/tests/utils_tests/mask_tests/test_scale_mask.py @@ -0,0 +1,78 @@ +from __future__ import division + +import unittest + +import numpy as np +import PIL.Image + +from chainer.backends import cuda +from chainer import testing +from chainer.testing import attr + +from chainercv.transforms import resize +from chainercv.utils import generate_random_bbox +from chainercv.utils import mask_to_bbox +from chainercv.utils import scale_mask + + +@testing.parameterize( + {'mask': np.array( + [[[False, False], + [False, True]]]), + 'expected': np.array( + [[[False, False, False, False], + [False, False, False, False], + [False, False, True, True], + [False, False, True, True]]]) + } +) +class TestScaleMaskSimple(unittest.TestCase): + + def check(self, mask, expected): + in_type = type(mask) + bbox = mask_to_bbox(mask) + size = 4 + out_mask = scale_mask(mask, bbox, size) + + self.assertIsInstance(out_mask, in_type) + self.assertEqual(out_mask.dtype, np.bool) + + np.testing.assert_equal( + cuda.to_cpu(out_mask), + cuda.to_cpu(expected)) + + def test_scale_mask_simple_cpu(self): + self.check(self.mask, self.expected) + + @attr.gpu + def test_scale_mask_simple_gpu(self): + self.check(cuda.to_gpu(self.mask), cuda.to_gpu(self.expected)) + + +class TestScaleMaskCompareResize(unittest.TestCase): + + def test(self): + H = 80 + W = 90 + n_inst = 10 + + mask = np.zeros((n_inst, H, W), dtype=np.bool) + bbox = generate_random_bbox(n_inst, (H, W), 10, 30).astype(np.int32) + for i, bb in enumerate(bbox): + y_min, x_min, y_max, x_max = bb + m = np.random.randint(0, 2, size=(y_max - y_min, x_max - x_min)) + m[5, 5] = 1 # At least one element is one + mask[i, y_min:y_max, x_min:x_max] = m + bbox = mask_to_bbox(mask) + size = H * 2 + out_H = size + out_W = W * 2 + out_mask = scale_mask(mask, bbox, size) + + expected = resize( + mask.astype(np.float32), (out_H, out_W), + interpolation=PIL.Image.NEAREST).astype(np.bool) + np.testing.assert_equal(out_mask, expected) + + +testing.run_module(__name__, __file__)