diff --git a/CHANGELOG.md b/CHANGELOG.md index ab020d5cc5..751420238b 100644 --- a/CHANGELOG.md +++ b/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 () + +### Security +- + ## 01/23/2021 - Release v0.1.5 ### Added - `WiderFace` dataset format (, ) diff --git a/datumaro/components/extractor.py b/datumaro/components/extractor.py index 1c81462779..c695a6e69e 100644 --- a/datumaro/components/extractor.py +++ b/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 diff --git a/datumaro/plugins/coco_format/converter.py b/datumaro/plugins/coco_format/converter.py index 44d604ab09..6a8f1bc47d 100644 --- a/datumaro/plugins/coco_format/converter.py +++ b/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 diff --git a/datumaro/plugins/mots_format.py b/datumaro/plugins/mots_format.py index cd1fa3f609..57bcbec475 100644 --- a/datumaro/plugins/mots_format.py +++ b/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) diff --git a/datumaro/plugins/tf_detection_api_format/converter.py b/datumaro/plugins/tf_detection_api_format/converter.py index 382149dc34..f7a94aaf91 100644 --- a/datumaro/plugins/tf_detection_api_format/converter.py +++ b/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] diff --git a/datumaro/plugins/transforms.py b/datumaro/plugins/transforms.py index c5030251f2..3101b0c776 100644 --- a/datumaro/plugins/transforms.py +++ b/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: diff --git a/datumaro/util/mask_tools.py b/datumaro/util/mask_tools.py index 72224bccf7..b6c2bc9462 100644 --- a/datumaro/util/mask_tools.py +++ b/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 \ No newline at end of file diff --git a/tests/test_coco_format.py b/tests/test_coco_format.py index 9fcd26d30d..9dd8e82ceb 100644 --- a/tests/test_coco_format.py +++ b/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} diff --git a/tests/test_transforms.py b/tests/test_transforms.py index de3cd66943..5098d03634 100644 --- a/tests/test_transforms.py +++ b/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]), ]), ])