diff --git a/datumaro/README.md b/datumaro/README.md index 5139976db05..66a43e52809 100644 --- a/datumaro/README.md +++ b/datumaro/README.md @@ -177,6 +177,24 @@ project = Project.load('directory') datum project diff mymodel_inference/ --format tensorboard --output-dir diff ``` +- Change colors in PASCAL VOC-like `.png` masks: + ```bash + datum project import --format voc --input-path + + # Create a color map file with desired colors: + # + # label : color_rgb : parts : actions + # cat:0,0,255:: + # dog:255,0,0:: + # + # Save as mycolormap.txt + + datum project export --format voc_segmentation -- --label-map mycolormap.txt + # add "--apply-colormap=0" to save grayscale (indexed) masks + # check "--help" option for more info + # use "datum --loglevel debug" for extra conversion info + ``` + diff --git a/datumaro/datumaro/plugins/voc_format/converter.py b/datumaro/datumaro/plugins/voc_format/converter.py index 4e4fefd157e..15b4086553c 100644 --- a/datumaro/datumaro/plugins/voc_format/converter.py +++ b/datumaro/datumaro/plugins/voc_format/converter.py @@ -16,6 +16,7 @@ from datumaro.components.extractor import (DEFAULT_SUBSET_NAME, AnnotationType, LabelCategories, CompiledMask, ) +from datumaro.util import str_to_bool, find from datumaro.util.image import save_image from datumaro.util.mask_tools import paint_mask, remap_mask @@ -466,6 +467,12 @@ def _load_categories(self, label_map_source=None): elif isinstance(label_map_source, str) and osp.isfile(label_map_source): label_map = parse_label_map(label_map_source) + has_black = find(label_map.items(), + lambda e: e[0] == 'background' or e[1][0] == (0, 0, 0)) + if not has_black and 'background' not in label_map: + label_map['background'] = [(0, 0, 0), [], []] + label_map.move_to_end('background', last=False) + else: raise Exception("Wrong labelmap specified, " "expected one of %s or a file path" % \ @@ -559,7 +566,7 @@ def build_cmdline_parser(cls, **kwargs): parser.add_argument('--save-images', action='store_true', help="Save images (default: %(default)s)") - parser.add_argument('--apply-colormap', type=bool, default=True, + parser.add_argument('--apply-colormap', type=str_to_bool, default=True, help="Use colormap for class and instance masks " "(default: %(default)s)") parser.add_argument('--label-map', type=cls._get_labelmap, default=None, diff --git a/datumaro/datumaro/util/__init__.py b/datumaro/datumaro/util/__init__.py index 19edb82de8b..1db427b1797 100644 --- a/datumaro/datumaro/util/__init__.py +++ b/datumaro/datumaro/util/__init__.py @@ -74,4 +74,13 @@ def take_by(iterable, count): if len(batch) == 0: break - yield batch \ No newline at end of file + yield batch + +def str_to_bool(s): + t = s.lower() + if t in {'true', '1', 'ok', 'yes', 'y'}: + return True + elif t in {'false', '0', 'no', 'n'}: + return False + else: + raise ValueError("Can't convert value '%s' to bool" % s)