Optimize mask conversions (#1097)

datumaro/datumaro/components/converters/voc.py

@@ -20,7 +20,7 @@ from datumaro.components.formats.voc import (VocTask, VocPath,
     parse_label_map, make_voc_label_map, make_voc_categories, write_label_map
 )
 from datumaro.util.image import save_image
-from datumaro.util.mask_tools import apply_colormap, remap_mask
+from datumaro.util.mask_tools import paint_mask, remap_mask
 
 
 def _convert_attr(name, attributes, type_conv, default=None, warn=True):
@@ -367,7 +367,7 @@ class _Converter:
             if colormap is None:
                 colormap = self._categories[AnnotationType.mask].colormap
             data = self._remap_mask(data)
-            data = apply_colormap(data, colormap)
+            data = paint_mask(data, colormap)
         save_image(path, data)
 
     def save_label_map(self):

datumaro/datumaro/util/mask_tools.py

@@ -30,65 +30,79 @@ def invert_colormap(colormap):
         tuple(a): index for index, a in colormap.items()
     }
 
+def check_is_mask(mask):
+    assert len(mask.shape) in {2, 3}
+    if len(mask.shape) == 3:
+        assert mask.shape[2] == 1
+
 _default_colormap = generate_colormap()
 _default_unpaint_colormap = invert_colormap(_default_colormap)
 
-def _default_unpaint_colormap_fn(r, g, b):
-    return _default_unpaint_colormap[(r, g, b)]
+def unpaint_mask(painted_mask, inverse_colormap=None):
+    # Covert color mask to index mask
 
-def unpaint_mask(painted_mask, colormap=None):
-    # expect HWC BGR [0; 255] image
-    # expect RGB->index colormap
+    # mask: HWC BGR [0; 255]
+    # colormap: (R, G, B) -> index
     assert len(painted_mask.shape) == 3
-    if colormap is None:
-        colormap = _default_unpaint_colormap_fn
-    if callable(colormap):
-        map_fn = lambda a: colormap(int(a[2]), int(a[1]), int(a[0]))
+    if inverse_colormap is None:
+        inverse_colormap = _default_unpaint_colormap
+
+    if callable(inverse_colormap):
+        map_fn = lambda a: inverse_colormap(
+                (a >> 16) & 255, (a >> 8) & 255, a & 255
+            )
     else:
-        map_fn = lambda a: colormap[(int(a[2]), int(a[1]), int(a[0]))]
+        map_fn = lambda a: inverse_colormap[(
+                (a >> 16) & 255, (a >> 8) & 255, a & 255
+            )]
 
-    unpainted_mask = np.apply_along_axis(map_fn,
-        1, np.reshape(painted_mask, (-1, 3)))
-    unpainted_mask = np.reshape(unpainted_mask, (painted_mask.shape[:2]))
-    return unpainted_mask.astype(int)
+    painted_mask = painted_mask.astype(int)
+    painted_mask = painted_mask[:, :, 0] + \
+                   (painted_mask[:, :, 1] << 8) + \
+                   (painted_mask[:, :, 2] << 16)
+    uvals, unpainted_mask = np.unique(painted_mask, return_inverse=True)
+    palette = np.array([map_fn(v) for v in uvals], dtype=np.float32)
+    unpainted_mask = palette[unpainted_mask].reshape(painted_mask.shape[:2])
 
+    return unpainted_mask
 
-def apply_colormap(mask, colormap=None):
-    # expect HW [0; max_index] mask
-    # expect index->RGB colormap
-    assert len(mask.shape) == 2
+def paint_mask(mask, colormap=None):
+    # Applies colormap to index mask
+
+    # mask: HW(C) [0; max_index] mask
+    # colormap: index -> (R, G, B)
+    check_is_mask(mask)
 
     if colormap is None:
         colormap = _default_colormap
     if callable(colormap):
-        map_fn = lambda p: colormap(int(p[0]))[::-1]
+        map_fn = colormap
     else:
-        map_fn = lambda p: colormap[int(p[0])][::-1]
-    painted_mask = np.apply_along_axis(map_fn, 1, np.reshape(mask, (-1, 1)))
+        map_fn = lambda c: colormap.get(c, (-1, -1, -1))
+    palette = np.array([map_fn(c)[::-1] for c in range(256)], dtype=np.float32)
 
-    painted_mask = np.reshape(painted_mask, (*mask.shape, 3))
-    return painted_mask.astype(np.float32)
+    mask = mask.astype(np.uint8)
+    painted_mask = palette[mask].reshape((*mask.shape[:2], 3))
+    return painted_mask
 
 def remap_mask(mask, map_fn):
     # Changes mask elements from one colormap to another
-    assert len(mask.shape) == 2
 
-    shape = mask.shape
-    mask = np.reshape(mask, (-1, 1))
-    mask = np.apply_along_axis(map_fn, 1, mask)
-    mask = np.reshape(mask, shape)
-    return mask
+    # mask: HW(C) [0; max_index] mask
+    check_is_mask(mask)
+
+    return np.array([map_fn(c) for c in range(256)], dtype=np.uint8)[mask]
 
 
-def load_mask(path, colormap=None):
+def load_mask(path, inverse_colormap=None):
     mask = load_image(path)
-    if colormap is not None:
+    if inverse_colormap is not None:
         if len(mask.shape) == 3 and mask.shape[2] != 1:
-            mask = unpaint_mask(mask, colormap=colormap)
+            mask = unpaint_mask(mask, inverse_colormap)
     return mask
 
-def lazy_mask(path, colormap=None):
-    return lazy_image(path, lambda path: load_mask(path, colormap))
+def lazy_mask(path, inverse_colormap=None):
+    return lazy_image(path, lambda path: load_mask(path, inverse_colormap))
 
 
 def mask_to_rle(binary_mask):

datumaro/tests/test_masks.py

@@ -66,4 +66,59 @@ class PolygonConversionsTest(TestCase):
         self.assertEqual(len(initial), len(computed))
         for i, (e_mask, c_mask) in enumerate(zip(expected, computed)):
             self.assertTrue(np.array_equal(e_mask, c_mask),
-                '#%s: %s\n%s\n' % (i, e_mask, c_mask))
+                '#%s: %s\n%s\n' % (i, e_mask, c_mask))
+
+class ColormapOperationsTest(TestCase):
+    def test_can_paint_mask(self):
+        mask = np.zeros((1, 3), dtype=np.uint8)
+        mask[:, 0] = 0
+        mask[:, 1] = 1
+        mask[:, 2] = 2
+
+        colormap = mask_tools.generate_colormap(3)
+
+        expected = np.zeros((*mask.shape, 3), dtype=np.uint8)
+        expected[:, 0] = colormap[0][::-1]
+        expected[:, 1] = colormap[1][::-1]
+        expected[:, 2] = colormap[2][::-1]
+
+        actual = mask_tools.paint_mask(mask, colormap)
+
+        self.assertTrue(np.array_equal(expected, actual),
+            '%s\nvs.\n%s' % (expected, actual))
+
+    def test_can_unpaint_mask(self):
+        colormap = mask_tools.generate_colormap(3)
+        inverse_colormap = mask_tools.invert_colormap(colormap)
+
+        mask = np.zeros((1, 3, 3), dtype=np.uint8)
+        mask[:, 0] = colormap[0][::-1]
+        mask[:, 1] = colormap[1][::-1]
+        mask[:, 2] = colormap[2][::-1]
+
+        expected = np.zeros((1, 3), dtype=np.uint8)
+        expected[:, 0] = 0
+        expected[:, 1] = 1
+        expected[:, 2] = 2
+
+        actual = mask_tools.unpaint_mask(mask, inverse_colormap)
+
+        self.assertTrue(np.array_equal(expected, actual),
+            '%s\nvs.\n%s' % (expected, actual))
+
+    def test_can_remap_mask(self):
+        class_count = 10
+        remap_fn = lambda c: class_count - c
+
+        src = np.empty((class_count, class_count), dtype=np.uint8)
+        for c in range(class_count):
+            src[c:, c:] = c
+
+        expected = np.empty_like(src)
+        for c in range(class_count):
+            expected[c:, c:] = remap_fn(c)
+
+        actual = mask_tools.remap_mask(src, remap_fn)
+
+        self.assertTrue(np.array_equal(expected, actual),
+            '%s\nvs.\n%s' % (expected, actual))