Optimize mask merging and improve memory use

CHANGELOG.md

@@ -6,6 +6,25 @@ 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).
 
 
+## [Unreleased]
+### Added
+-
+
+### Changed
+-
+
+### Deprecated
+-
+
+### Removed
+-
+
+### Fixed
+- High memory consumption and low performance of mask import/export, #53 (<https://github.com/openvinotoolkit/datumaro/pull/101>)
+
+### Security
+-
+
 ## 01/23/2021 - Release v0.1.5
 ### Added
 - `WiderFace` dataset format (<https://github.com/openvinotoolkit/datumaro/pull/65>, <https://github.com/openvinotoolkit/datumaro/pull/90>)

datumaro/components/extractor.py

@@ -167,10 +167,12 @@ def image(self):
     def as_class_mask(self, label_id=None):
         if label_id is None:
             label_id = self.label
-        return self.image * label_id
+        from datumaro.util.mask_tools import make_index_mask
+        return make_index_mask(self.image, label_id)
 
     def as_instance_mask(self, instance_id):
-        return self.image * instance_id
+        from datumaro.util.mask_tools import make_index_mask
+        return make_index_mask(self.image, instance_id)
 
     def get_area(self):
         return np.count_nonzero(self.image)
@@ -203,7 +205,7 @@ class RleMask(Mask):
     @staticmethod
     def _lazy_decode(rle):
         from pycocotools import mask as mask_utils
-        return lambda: mask_utils.decode(rle).astype(np.bool)
+        return lambda: mask_utils.decode(rle)
 
     def get_area(self):
         from pycocotools import mask as mask_utils
@@ -222,7 +224,7 @@ class CompiledMask:
     @staticmethod
     def from_instance_masks(instance_masks,
             instance_ids=None, instance_labels=None):
-        from datumaro.util.mask_tools import merge_masks
+        from datumaro.util.mask_tools import make_index_mask
 
         if instance_ids is not None:
             assert len(instance_ids) == len(instance_masks)
@@ -234,17 +236,27 @@ def from_instance_masks(instance_masks,
         else:
             instance_labels = [None] * len(instance_masks)
 
-        instance_masks = sorted(
-            zip(instance_masks, instance_ids, instance_labels),
-            key=lambda m: m[0].z_order)
+        instance_masks = sorted(enumerate(instance_masks),
+            key=lambda m: m[1].z_order)
+        instance_masks = ((m.image,
+                instance_ids[i] if instance_ids[i] is not None else 1 + j,
+                instance_labels[i] if instance_labels[i] is not None else m.label
+            ) for j, (i, m) in enumerate(instance_masks))
 
-        instance_mask = [m.as_instance_mask(id if id is not None else 1 + idx)
-            for idx, (m, id, _) in enumerate(instance_masks)]
-        instance_mask = merge_masks(instance_mask)
+        # 1. Avoid memory explosion on materialization of all masks
+        # 2. Optimize materialization calls
+        it = iter(instance_masks)
 
-        cls_mask = [m.as_class_mask(c) for m, _, c in instance_masks]
-        cls_mask = merge_masks(cls_mask)
-        return __class__(class_mask=cls_mask, instance_mask=instance_mask)
+        m, instance_id, class_id = next(it)
+        merged_instance_mask = make_index_mask(m, instance_id)
+        merged_class_mask = make_index_mask(m, class_id)
+
+        for m, instance_id, class_id in it:
+            merged_instance_mask = np.where(m, instance_id, merged_instance_mask)
+            merged_class_mask = np.where(m, class_id, merged_class_mask)
+
+        return __class__(class_mask=merged_class_mask,
+            instance_mask=merged_instance_mask)
 
     def __init__(self, class_mask=None, instance_mask=None):
         self._class_mask = class_mask

datumaro/plugins/coco_format/converter.py

@@ -8,7 +8,7 @@
 import os
 import os.path as osp
 from enum import Enum
-from itertools import groupby
+from itertools import chain, groupby
 
 import pycocotools.mask as mask_utils
 
@@ -224,16 +224,17 @@ def find_instance_parts(self, group, img_width, img_height):
                 mask = mask_tools.rles_to_mask(polygons, img_width, img_height)
 
             if masks:
+                masks = (m.image for m in masks)
                 if mask is not None:
-                    masks += [mask]
-                mask = mask_tools.merge_masks([m.image for m in masks])
+                    masks += chain(masks, [mask])
+                mask = mask_tools.merge_masks(masks)
 
             if mask is not None:
                 mask = mask_tools.mask_to_rle(mask)
             polygons = []
         else:
             if masks:
-                mask = mask_tools.merge_masks([m.image for m in masks])
+                mask = mask_tools.merge_masks(m.image for m in masks)
                 polygons += mask_tools.mask_to_polygons(mask)
             mask = None
 

datumaro/plugins/mots_format.py

@@ -138,8 +138,8 @@ def _save_annotations(self, item, anno_dir):
 
         instance_ids = [int(a.attributes['track_id']) for a in masks]
         masks = sorted(zip(masks, instance_ids), key=lambda e: e[0].z_order)
-        mask = merge_masks([
-            m.image * (MotsPath.MAX_INSTANCES * (1 + m.label) + id)
-            for m, id in masks])
+        mask = merge_masks(
+            (m.image, MotsPath.MAX_INSTANCES * (1 + m.label) + id)
+            for m, id in masks)
         save_image(osp.join(anno_dir, item.id + '.png'), mask,
             create_dir=True, dtype=np.uint16)

datumaro/plugins/tf_detection_api_format/converter.py

@@ -105,7 +105,7 @@ def _find_instance_parts(self, group, img_width, img_height):
 
         mask = None
         if self._save_masks:
-            mask = merge_masks([m.image for m in masks])
+            mask = merge_masks(m.image for m in masks)
 
         return [leader, mask, bbox]
 

datumaro/plugins/transforms.py

@@ -4,6 +4,7 @@
 
 from collections import Counter
 from enum import Enum
+from itertools import chain
 import logging as log
 import os.path as osp
 import random
@@ -129,6 +130,8 @@ def merge_segments(cls, instance, img_width, img_height,
         masks = [a for a in instance if a.type == AnnotationType.mask]
         if not polygons and not masks:
             return []
+        if not polygons and len(masks) == 1:
+            return masks
 
         leader = find_group_leader(polygons + masks)
         instance = []
@@ -143,9 +146,9 @@ def merge_segments(cls, instance, img_width, img_height,
             instance += polygons # keep unused polygons
 
         if masks:
-            masks = [m.image for m in masks]
+            masks = (m.image for m in masks)
             if mask is not None:
-                masks += [mask]
+                masks = chain(masks, [mask])
             mask = mask_tools.merge_masks(masks)
 
         if mask is None:

datumaro/util/mask_tools.py

@@ -66,16 +66,19 @@ def unpaint_mask(painted_mask, inverse_colormap=None):
                    (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)
+    palette = np.array([map_fn(v) for v in uvals],
+        dtype=np.min_scalar_type(len(uvals)))
     unpainted_mask = palette[unpainted_mask].reshape(painted_mask.shape[:2])
 
     return unpainted_mask
 
 def paint_mask(mask, colormap=None):
-    # Applies colormap to index mask
+    """
+    Applies colormap to index mask
 
-    # mask: HW(C) [0; max_index] mask
-    # colormap: index -> (R, G, B)
+    mask: HW(C) [0; max_index] mask
+    colormap: index -> (R, G, B)
+    """
     check_is_mask(mask)
 
     if colormap is None:
@@ -84,25 +87,30 @@ def paint_mask(mask, colormap=None):
         map_fn = colormap
     else:
         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)
+    palette = np.array([map_fn(c)[::-1] for c in range(256)], dtype=np.uint8)
 
     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
+    """
+    Changes mask elements from one colormap to another
 
     # 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 make_index_mask(binary_mask, index):
-    return np.choose(binary_mask, np.array([0, index], dtype=np.uint8))
+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_binary_mask(mask):
-    return np.nonzero(mask)
+    if mask.dtype.kind == 'b':
+        return mask
+    return mask.astype(bool)
 
 
 def load_mask(path, inverse_colormap=None):
@@ -135,7 +143,7 @@ def mask_to_rle(binary_mask):
         'size': list(binary_mask.shape)
     }
 
-def mask_to_polygons(mask, tolerance=1.0, area_threshold=1):
+def mask_to_polygons(mask, area_threshold=1):
     """
     Convert an instance mask to polygons
 
@@ -149,25 +157,22 @@ def mask_to_polygons(mask, tolerance=1.0, area_threshold=1):
         A list of polygons like [[x1,y1, x2,y2 ...], [...]]
     """
     from pycocotools import mask as mask_utils
-    from skimage import measure
+    import cv2
 
     polygons = []
 
-    # pad mask with 0 around borders
-    padded_mask = np.pad(mask, pad_width=1, mode='constant', constant_values=0)
-    contours = measure.find_contours(padded_mask, 0.5)
-    # Fix coordinates after padding
-    contours = np.subtract(contours, 1)
+    contours, _ = cv2.findContours(mask.astype(np.uint8),
+        mode=cv2.RETR_TREE, method=cv2.CHAIN_APPROX_TC89_KCOS)
 
     for contour in contours:
-        if not np.array_equal(contour[0], contour[-1]):
-            contour = np.vstack((contour, contour[0])) # make polygon closed
-
-        contour = measure.approximate_polygon(contour, tolerance)
         if len(contour) <= 2:
             continue
 
-        contour = np.flip(contour, axis=1).flatten().clip(0) # [x0, y0, ...]
+        contour = contour.reshape((-1, 2))
+
+        if not np.array_equal(contour[0], contour[-1]):
+            contour = np.vstack((contour, contour[0])) # make polygon closed
+        contour = contour.flatten().clip(0) # [x0, y0, ...]
 
         # Check if the polygon is big enough
         rle = mask_utils.frPyObjects([contour], mask.shape[0], mask.shape[1])
@@ -277,12 +282,25 @@ def find_mask_bbox(mask):
 def merge_masks(masks):
     """
         Merges masks into one, mask order is responsible for z order.
+        To avoid memory explosion on mask materialization, consider passing
+        a generator.
+
+        Inputs: a sequence of index masks or (binary mask, index) pairs
+        Outputs: an index mask
     """
-    if not masks:
+    it = iter(masks)
+
+    try:
+        merged_mask = next(it)
+        if isinstance(merged_mask, tuple) and len(merged_mask) == 2:
+            merged_mask = merged_mask[0] * merged_mask[1]
+    except StopIteration:
         return None
 
-    merged_mask = masks[0]
-    for m in masks[1:]:
-        merged_mask = np.where(m != 0, m, merged_mask)
+    for m in it:
+        if isinstance(m, tuple) and len(m) == 2:
+            merged_mask = np.where(m[0], m[1], merged_mask)
+        else:
+            merged_mask = np.where(m, m, merged_mask)
 
     return merged_mask

tests/test_coco_format.py

@@ -390,11 +390,11 @@ def test_can_convert_masks_to_polygons(self):
             DatasetItem(id=1, image=np.zeros((5, 10, 3)),
                 annotations=[
                     Polygon(
-                        [3.0, 2.5, 1.0, 0.0, 3.5, 0.0, 3.0, 2.5],
+                        [1, 0, 3, 2, 3, 0, 1, 0],
                         label=3, id=4, group=4,
                         attributes={ 'is_crowd': False }),
                     Polygon(
-                        [5.0, 3.5, 4.5, 0.0, 8.0, 0.0, 5.0, 3.5],
+                        [5, 0, 5, 3, 8, 0, 5, 0],
                         label=3, id=4, group=4,
                         attributes={ 'is_crowd': False }),
                 ], attributes={'id': 1}

tests/test_transforms.py

@@ -45,8 +45,8 @@ def test_mask_to_polygons(self):
 
         expected = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 10, 3)), annotations=[
-                Polygon([3.0, 2.5, 1.0, 0.0, 3.5, 0.0, 3.0, 2.5]),
-                Polygon([5.0, 3.5, 4.5, 0.0, 8.0, 0.0, 5.0, 3.5]),
+                Polygon([1, 0, 3, 2, 3, 0, 1, 0]),
+                Polygon([5, 0, 5, 3, 8, 0, 5, 0]),
             ]),
         ])