Merge pull request #552 from knorth55/mask-iou

add mask_iou

chainercv/utils/__init__.py

@@ -9,6 +9,7 @@
 from chainercv.utils.iterator import apply_to_iterator  # NOQA
 from chainercv.utils.iterator import ProgressHook  # NOQA
 from chainercv.utils.iterator import unzip  # NOQA
+from chainercv.utils.mask.mask_iou import mask_iou  # 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

chainercv/utils/mask/mask_iou.py

@@ -0,0 +1,46 @@
+from __future__ import division
+
+
+from chainer import cuda
+
+
+def mask_iou(mask_a, mask_b):
+    """Calculate the Intersection of Unions (IoUs) between masks.
+
+    IoU is calculated as a ratio of area of the intersection
+    and area of the union.
+
+    This function accepts both :obj:`numpy.ndarray` and :obj:`cupy.ndarray` as
+    inputs. Please note that both :obj:`mask_a` and :obj:`mask_b` need to be
+    same type.
+    The output is same type as the type of the inputs.
+
+    Args:
+        mask_a (array): An array whose shape is :math:`(N, H, W)`.
+            :math:`N` is the number of masks.
+            The dtype should be :obj:`numpy.bool`.
+        mask_b (array): An array similar to :obj:`mask_a`,
+            whose shape is :math:`(K, H, W)`.
+            The dtype should be :obj:`numpy.bool`.
+
+    Returns:
+        array:
+        An array whose shape is :math:`(N, K)`. \
+        An element at index :math:`(n, k)` contains IoUs between \
+        :math:`n` th mask in :obj:`mask_a` and :math:`k` th mask \
+        in :obj:`mask_b`.
+
+    """
+    if mask_a.shape[1:] != mask_b.shape[1:]:
+        raise IndexError
+    xp = cuda.get_array_module(mask_a)
+
+    n_mask_a = len(mask_a)
+    n_mask_b = len(mask_b)
+    iou = xp.empty((n_mask_a, n_mask_b), dtype=xp.float32)
+    for n, m_a in enumerate(mask_a):
+        for k, m_b in enumerate(mask_b):
+            intersect = xp.bitwise_and(m_a, m_b).sum()
+            union = xp.bitwise_or(m_a, m_b).sum()
+            iou[n, k] = intersect / union
+    return iou

docs/source/reference/utils.rst

@@ -64,6 +64,13 @@ unzip
 .. autofunction:: unzip
 
 
+Mask Utilities
+--------------
+
+mask_iou
+~~~~~~~~
+.. autofunction:: mask_iou
+
 Testing Utilities
 -----------------
 

tests/utils_tests/mask_tests/test_mask_iou.py

@@ -0,0 +1,91 @@
+from __future__ import division
+
+import unittest
+
+import numpy as np
+
+from chainer import cuda
+from chainer import testing
+from chainer.testing import attr
+
+from chainercv.utils import mask_iou
+
+
+@testing.parameterize(
+    {'mask_a': np.array(
+        [[[False, False], [True, True]],
+         [[True, True], [False, False]]],
+        dtype=np.bool),
+     'mask_b': np.array(
+        [[[False, False], [True, True]],
+         [[True, True], [False, False]],
+         [[True, False], [True, True]],
+         [[True, True], [False, True]]],
+        dtype=np.bool),
+     'expected': np.array(
+        [[1., 0., 2 / 3, 1 / 4],
+         [0., 1., 1 / 4, 2 / 3]],
+        dtype=np.float32)
+     },
+    {'mask_a': np.array(
+        [[[False, False], [True, True]],
+         [[True, True], [False, False]],
+         [[True, True], [True, False]],
+         [[False, True], [True, True]]],
+        dtype=np.bool),
+     'mask_b': np.array(
+        [[[False, False], [True, True]],
+         [[True, True], [False, False]]],
+        dtype=np.bool),
+     'expected': np.array(
+        [[1., 0.], [0., 1.], [1 / 4, 2 / 3], [2 / 3, 1 / 4]],
+        dtype=np.float32)
+     },
+    {'mask_a': np.zeros((0, 2, 2), dtype=np.bool),
+     'mask_b': np.array([[[False, False], [False, False]]], dtype=np.bool),
+     'expected': np.zeros((0, 1), dtype=np.float32)
+     },
+)
+class TestMaskIou(unittest.TestCase):
+
+    def check(self, mask_a, mask_b, expected):
+        iou = mask_iou(mask_a, mask_b)
+
+        self.assertIsInstance(iou, type(expected))
+        np.testing.assert_equal(
+            cuda.to_cpu(iou),
+            cuda.to_cpu(expected))
+
+    def test_mask_iou_cpu(self):
+        self.check(self.mask_a, self.mask_b, self.expected)
+
+    @attr.gpu
+    def test_mask_iou_gpu(self):
+        self.check(
+            cuda.to_gpu(self.mask_a),
+            cuda.to_gpu(self.mask_b),
+            cuda.to_gpu(self.expected))
+
+
+@testing.parameterize(
+    {'mask_a': np.array([[[False], [True, True]]], dtype=np.bool),
+     'mask_b': np.array([[[False, False], [True, True]]], dtype=np.bool)
+     },
+    {'mask_a': np.array([[[False, False, True], [True, True]]], dtype=np.bool),
+     'mask_b': np.array([[[False, False], [True, True]]], dtype=np.bool)
+     },
+    {'mask_a': np.array([[[False, False], [True, True]]], dtype=np.bool),
+     'mask_b': np.array([[[False], [True, True]]], dtype=np.bool)
+     },
+    {'mask_a': np.array([[[False, False], [True, True]]], dtype=np.bool),
+     'mask_b': np.array([[[False, False, True], [True, True]]], dtype=np.bool)
+     },
+)
+class TestMaskIouInvalidShape(unittest.TestCase):
+
+    def test_mask_iou_invalid(self):
+        with self.assertRaises(IndexError):
+            mask_iou(self.mask_a, self.mask_b)
+
+
+testing.run_module(__name__, __file__)