diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ead547abb..9d8c401a66 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -68,6 +68,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Detection for LFW format () +- Export of masks with background class with id != 0 in the VOC format + () ### Security - TBD diff --git a/datumaro/components/annotation.py b/datumaro/components/annotation.py index 2424856db6..5e214ca34e 100644 --- a/datumaro/components/annotation.py +++ b/datumaro/components/annotation.py @@ -388,6 +388,7 @@ def from_instance_masks( instance_masks: Iterable[Mask], instance_ids: Optional[Iterable[int]] = None, instance_labels: Optional[Iterable[int]] = None, + background_label_id: int = 0, ) -> CompiledMask: """ Joins instance masks into a single mask. Masks are sorted by @@ -398,6 +399,9 @@ def from_instance_masks( By default, mask positions are used. instance_labels: Instance label id values for the produced class mask. By default, mask labels are used. + background_label_id: The background label index. Masks with label None or + with this label are mapped to the same instance id 0. + By default, the background label is 0. """ from datumaro.util.mask_tools import make_index_mask @@ -427,28 +431,31 @@ def from_instance_masks( it = iter(masks) + # Generate an index mask + index_mask = None instance_map = [0] - class_map = [0] - - m, idx, instance_id, class_id = next(it) - if not class_id: - idx = 0 - index_mask = make_index_mask(m.image, idx, dtype=index_dtype) - instance_map.append(instance_id) - class_map.append(class_id) - + class_map = [background_label_id] for m, idx, instance_id, class_id in it: - if not class_id: + if class_id in [background_label_id, None]: + # Optimization A: map all background masks to the same idx 0 idx = 0 - index_mask = np.where(m.image, idx, index_mask) + + if index_mask is not None: + index_mask = np.where(m.image, idx, index_mask) + else: + index_mask = make_index_mask(m.image, idx, dtype=index_dtype) + instance_map.append(instance_id) class_map.append(class_id) # Generate compiled masks + # Map the index mask to segmentation masks if np.array_equal(instance_map, range(max_index)): + # Optimization B: can reuse the index mask generated in the Optimization A merged_instance_mask = index_mask else: + # TODO: squash spaces in the instance indices? merged_instance_mask = np.array(instance_map, dtype=np.min_scalar_type(instance_map))[ index_mask ] diff --git a/datumaro/plugins/voc_format/converter.py b/datumaro/plugins/voc_format/converter.py index 045722ac67..3e05039737 100644 --- a/datumaro/plugins/voc_format/converter.py +++ b/datumaro/plugins/voc_format/converter.py @@ -1,11 +1,12 @@ # Copyright (C) 2020-2022 Intel Corporation +# Copyright (C) 2022 CVAT.ai Corporation # # SPDX-License-Identifier: MIT import logging as log import os import os.path as osp -from collections import OrderedDict, defaultdict +from collections import defaultdict from enum import Enum, auto from itertools import chain from typing import Dict, Optional, Set @@ -28,23 +29,13 @@ from datumaro.components.errors import MediaTypeError from datumaro.components.extractor import DatasetItem from datumaro.components.media import Image -from datumaro.util import find, str_to_bool +from datumaro.util import str_to_bool from datumaro.util.annotation_util import make_label_id_mapping from datumaro.util.image import save_image -from datumaro.util.mask_tools import paint_mask, remap_mask +from datumaro.util.mask_tools import paint_mask from datumaro.util.meta_file_util import has_meta_file -from .format import ( - VocInstColormap, - VocPath, - VocTask, - make_voc_categories, - make_voc_label_map, - parse_label_map, - parse_meta_file, - write_label_map, - write_meta_file, -) +from .format import VocInstColormap, VocLabelMap, VocPath, VocTask, make_voc_categories def _convert_attr(name, attributes, type_conv, default=None): @@ -111,7 +102,7 @@ def build_cmdline_parser(cls, **kwargs): "--apply-colormap", type=str_to_bool, default=True, - help="Use colormap for class and instance masks " "(default: %(default)s)", + help="Use colormap for class and instance masks (default: %(default)s)", ) parser.add_argument( "--label-map", @@ -129,7 +120,7 @@ def build_cmdline_parser(cls, **kwargs): "--keep-empty", type=str_to_bool, default=False, - help="Write subset lists even if they are empty " "(default: %(default)s)", + help="Write subset lists even if they are empty (default: %(default)s)", ) parser.add_argument( "--tasks", @@ -212,7 +203,7 @@ def make_dirs(self): self._inst_dir = inst_dir self._images_dir = images_dir - def get_label(self, label_id): + def get_label_name(self, label_id: int) -> str: return self._extractor.categories()[AnnotationType.label].items[label_id].name def save_subsets(self): @@ -294,7 +285,7 @@ def _export_annotations(self, item: DatasetItem, *, image_filename: str, lists: main_bboxes = [] layout_bboxes = [] for bbox in bboxes: - label = self.get_label(bbox.label) + label = self.get_label_name(bbox.label) if self._is_part(label): layout_bboxes.append(bbox) elif self._is_label(label): @@ -305,7 +296,7 @@ def _export_annotations(self, item: DatasetItem, *, image_filename: str, lists: obj_elem = ET.SubElement(root_elem, "object") - obj_label = self.get_label(obj.label) + obj_label = self.get_label_name(obj.label) ET.SubElement(obj_elem, "name").text = obj_label if "pose" in attr: @@ -329,7 +320,7 @@ def _export_annotations(self, item: DatasetItem, *, image_filename: str, lists: lambda x: obj.group and obj.group == x.group, layout_bboxes ): part_elem = ET.SubElement(obj_elem, "part") - ET.SubElement(part_elem, "name").text = self.get_label(part_bbox.label) + ET.SubElement(part_elem, "name").text = self.get_label_name(part_bbox.label) _write_xml_bbox(part_bbox.get_bbox(), part_elem) objects_with_parts.append(new_obj_id) @@ -379,7 +370,7 @@ def _export_annotations(self, item: DatasetItem, *, image_filename: str, lists: lists.action_list[item.id] = objects_with_actions for label_ann in labels: - label = self.get_label(label_ann.label) + label = self.get_label_name(label_ann.label) if not self._is_label(label): continue class_list = lists.class_lists.get(item.id, set()) @@ -390,7 +381,9 @@ def _export_annotations(self, item: DatasetItem, *, image_filename: str, lists: if masks and VocTask.segmentation in self._tasks: compiled_mask = CompiledMask.from_instance_masks( - masks, instance_labels=[self._label_id_mapping(m.label) for m in masks] + masks, + instance_labels=[self._label_id_mapping(m.label) for m in masks], + background_label_id=self._label_id_mapping(None), ) self.save_segm( @@ -486,7 +479,7 @@ def save_class_lists(self, subset_name, class_lists): def _write_item(f, item, item_labels): if not item_labels: return - item_labels = [self.get_label(l) for l in item_labels] + item_labels = [self.get_label_name(l) for l in item_labels] presented = label in item_labels f.write("%s % d\n" % (item, 1 if presented else -1)) @@ -574,101 +567,74 @@ def _write_item(f, item, item_layouts): def save_segm(self, path, mask, colormap=None): if self._apply_colormap: if colormap is None: - colormap = self._categories[AnnotationType.mask].colormap + colormap = self._output_categories[AnnotationType.mask].colormap mask = paint_mask(mask, colormap) save_image(path, mask, create_dir=True) def save_label_map(self): if self._save_dataset_meta: - write_meta_file(self._save_dir, self._label_map) + self._label_map.dump_to_meta_file(self._save_dir) else: - path = osp.join(self._save_dir, VocPath.LABELMAP_FILE) - write_label_map(path, self._label_map) + self._label_map.dump_to_file(osp.join(self._save_dir, VocPath.LABELMAP_FILE)) def _load_categories(self, label_map_source): if label_map_source == LabelmapType.voc.name: # use the default VOC colormap - label_map = make_voc_label_map() - - elif ( - label_map_source == LabelmapType.source.name - and AnnotationType.mask not in self._extractor.categories() - ): - # generate colormap for input labels - labels = self._extractor.categories().get(AnnotationType.label, LabelCategories()) - label_map = OrderedDict((item.name, [None, [], []]) for item in labels.items) - - elif ( - label_map_source == LabelmapType.source.name - and AnnotationType.mask in self._extractor.categories() - ): - # use source colormap - labels = self._extractor.categories()[AnnotationType.label] - colors = self._extractor.categories()[AnnotationType.mask] - label_map = OrderedDict() - for idx, item in enumerate(labels.items): - color = colors.colormap.get(idx) - if color is not None: - label_map[item.name] = [color, [], []] + label_map = VocLabelMap.make_default() + + elif label_map_source == LabelmapType.source.name: + label_map = VocLabelMap.from_categories(self._extractor.categories()) elif isinstance(label_map_source, dict): - label_map = OrderedDict(sorted(label_map_source.items(), key=lambda e: e[0])) + # TODO: move sorting to CVAT + label_map = VocLabelMap(sorted(label_map_source.items(), key=lambda e: e[0])) elif isinstance(label_map_source, str) and osp.isfile(label_map_source): if has_meta_file(label_map_source): - label_map = parse_meta_file(label_map_source) + label_map = VocLabelMap.parse_from_meta_file(label_map_source) else: - label_map = parse_label_map(label_map_source) + label_map = VocLabelMap.parse_from_file(label_map_source) else: - raise Exception( - "Wrong labelmap specified: '%s', " - "expected one of %s or a file path" + raise ValueError( + "Wrong labelmap specified: '%s', expected one of %s or a file path" % (label_map_source, ", ".join(t.name for t in LabelmapType)) ) - bg_label = find(label_map.items(), lambda x: x[1][0] == (0, 0, 0)) - if bg_label is None: - bg_label = "background" - if bg_label not in label_map: - has_colors = any(v[0] is not None for v in label_map.values()) - color = (0, 0, 0) if has_colors else None - label_map[bg_label] = [color, [], []] - label_map.move_to_end(bg_label, last=False) - - self._categories = make_voc_categories(label_map) + bg_label = label_map.find_or_create_background_label() + output_categories = make_voc_categories(label_map) # Update colors with assigned values - colormap = self._categories[AnnotationType.mask].colormap + colormap = output_categories[AnnotationType.mask].colormap for label_id, color in colormap.items(): - label_desc = label_map[self._categories[AnnotationType.label].items[label_id].name] - label_desc[0] = color + label_desc = label_map[output_categories[AnnotationType.label].items[label_id].name] + label_desc.color = color + self._output_categories = output_categories self._label_map = label_map + self._bg_label = bg_label self._label_id_mapping = self._make_label_id_map() def _is_label(self, s): - return self._label_map.get(s) is not None + return self._label_map.is_label(s) def _is_part(self, s): - for label_desc in self._label_map.values(): - if s in label_desc[1]: - return True - return False + return self._label_map.is_part(s) def _is_action(self, label, s): - return s in self._get_actions(label) + return self._label_map.is_action(label, s) def _get_actions(self, label): - label_desc = self._label_map.get(label) - if not label_desc: - return [] - return label_desc[2] + return self._label_map.get_actions(label) def _make_label_id_map(self): + src_cat: LabelCategories = self._extractor.categories().get(AnnotationType.label) + dst_cat: LabelCategories = self._output_categories[AnnotationType.label] + bg_label_id = dst_cat.find(self._bg_label)[0] map_id, id_mapping, src_labels, dst_labels = make_label_id_mapping( - self._extractor.categories().get(AnnotationType.label), - self._categories[AnnotationType.label], + src_cat, + dst_cat, + fallback=bg_label_id, ) void_labels = [ @@ -687,7 +653,9 @@ def _make_label_id_map(self): src_id, src_label, id_mapping[src_id], - self._categories[AnnotationType.label].items[id_mapping[src_id]].name, + self._output_categories[AnnotationType.label] + .items[id_mapping[src_id]] + .name, ) for src_id, src_label in src_labels.items() ] @@ -696,9 +664,6 @@ def _make_label_id_map(self): return map_id - def _remap_mask(self, mask): - return remap_mask(mask, self._label_id_mapping) - @classmethod def patch(cls, dataset, patch, save_dir, **kwargs): conv = cls(patch.as_dataset(dataset), save_dir=save_dir, **kwargs) diff --git a/datumaro/plugins/voc_format/extractor.py b/datumaro/plugins/voc_format/extractor.py index b8a9ba20b2..b5b86561ab 100644 --- a/datumaro/plugins/voc_format/extractor.py +++ b/datumaro/plugins/voc_format/extractor.py @@ -30,14 +30,7 @@ from datumaro.util.mask_tools import invert_colormap, lazy_mask from datumaro.util.meta_file_util import has_meta_file -from .format import ( - VocInstColormap, - VocPath, - VocTask, - make_voc_categories, - parse_label_map, - parse_meta_file, -) +from .format import VocInstColormap, VocLabelMap, VocPath, VocTask, make_voc_categories _inverse_inst_colormap = invert_colormap(VocInstColormap) @@ -81,11 +74,11 @@ def _get_label_id(self, label: str) -> int: def _load_categories(self, dataset_path): label_map = None if has_meta_file(dataset_path): - label_map = parse_meta_file(dataset_path) + label_map = VocLabelMap.parse_from_meta_file(dataset_path) else: label_map_path = osp.join(dataset_path, VocPath.LABELMAP_FILE) if osp.isfile(label_map_path): - label_map = parse_label_map(label_map_path) + label_map = VocLabelMap.parse_from_file(label_map_path) return make_voc_categories(label_map) diff --git a/datumaro/plugins/voc_format/format.py b/datumaro/plugins/voc_format/format.py index 5cbed30057..80768b9f7a 100644 --- a/datumaro/plugins/voc_format/format.py +++ b/datumaro/plugins/voc_format/format.py @@ -1,14 +1,19 @@ # Copyright (C) 2019-2022 Intel Corporation +# Copyright (C) 2022 CVAT.ai Corporation # # SPDX-License-Identifier: MIT +from __future__ import annotations + import os.path as osp +import sys from collections import OrderedDict from enum import Enum, auto from itertools import chain -from typing import Dict, List, Optional, Tuple +from typing import Dict, Optional, Set, Union import numpy as np +from attr import define, field from datumaro.components.annotation import ( AnnotationType, @@ -124,162 +129,258 @@ class VocPath: } -LabelMapConfig = Dict[str, Tuple[Optional[RgbColor], List[str], List[str]]] -"""A type representing a label map config""" -# Not totally type-correct, tuple elements are supposed to support modification. -# Therefore, the tuple is typically a list -# TODO: refactor, make type annotations conform with actual usage +@define +class VocLabelInfo: + color: Optional[RgbColor] = None + parts: Set[str] = field(factory=set, converter=set) + actions: Set[str] = field(factory=set, converter=set) -def make_voc_label_map() -> LabelMapConfig: - labels = sorted(VocLabel, key=lambda l: l.value) - label_map = OrderedDict((label.name, [VocColormap[label.value], [], []]) for label in labels) - label_map[VocLabel.person.name][1] = [p.name for p in VocBodyPart] - label_map[VocLabel.person.name][2] = [a.name for a in VocAction] - return label_map +if sys.version_info < (3, 9): + _voc_label_map_base = OrderedDict +else: + _voc_label_map_base = OrderedDict[str, VocLabelInfo] -def parse_label_map(path: str) -> LabelMapConfig: +class VocLabelMap(_voc_label_map_base): + """ + Provides VOC-specific info about labels. + A mapping of type: str -> VocLabelInfo """ - Parses a label map file in the format: - 'name : color (r, g, b) : parts (hand, feet, ...) : actions (siting, standing, ...)' - Parameters: - path: File path + DEFAULT_BACKGROUND_LABEL = "background" + DEFAULT_BACKGROUND_COLOR: RgbColor = (0, 0, 0) - Returns: - A dictionary: label -> (color, parts, actions) - """ + def __init__(self, *args, **kwargs) -> None: + d = OrderedDict(*args, **kwargs) + + super().__init__( + (k, v if isinstance(v, VocLabelInfo) else VocLabelInfo(*v)) for k, v in d.items() + ) + + def is_label(self, label: str) -> bool: + return self.get(label) is not None + + def is_part(self, part: str) -> bool: + for label_desc in self.values(): + if part in label_desc.parts: + return True + return False + + def is_action(self, label: str, action: str) -> bool: + return action in self.get_actions(label) - label_map = OrderedDict() - with open(path, "r", encoding="utf-8") as f: - for line in f: - # skip empty and commented lines - line = line.strip() - if not line or line[0] == "#": - continue - - # name : color : parts : actions - label_desc = line.strip().split(":") - if len(label_desc) != 4: - raise InvalidAnnotationError( - f"Label description has wrong number of fields '{len(label_desc)}'. " - "Expected 4 ':'-separated fields." - ) - name = label_desc[0] - - if name in label_map: - raise InvalidAnnotationError(f"Label '{name}' is already defined in the label map") - - if 1 < len(label_desc) and label_desc[1]: - color = label_desc[1].split(",") - if len(color) != 3: + def get_actions(self, label: str) -> Set[str]: + label_desc = self.get(label) + if not label_desc: + return set() + return label_desc.actions + + def has_colors(self) -> bool: + return any(v.color is not None for v in self.values()) + + def find_background_label( + self, + *, + name: str = DEFAULT_BACKGROUND_LABEL, + color: RgbColor = DEFAULT_BACKGROUND_COLOR, + ) -> Optional[str]: + bg_label = find(self.items(), lambda x: x[1].color == color) + if bg_label is not None: + return bg_label[0] + + if name in self: + return name + + return None + + def find_or_create_background_label( + self, + *, + name: str = DEFAULT_BACKGROUND_LABEL, + color: RgbColor = DEFAULT_BACKGROUND_COLOR, + ) -> str: + bg_label = self.find_background_label(color=color, name=name) + + if bg_label is None: + bg_label = name + color = color if self.has_colors() else None + self[bg_label] = VocLabelInfo(color) + self.move_to_end(bg_label, last=False) + + return bg_label + + @staticmethod + def make_default() -> VocLabelMap: + labels = sorted(VocLabel, key=lambda l: l.value) + label_map = VocLabelMap( + (label.name, VocLabelInfo(VocColormap[label.value])) for label in labels + ) + label_map[VocLabel.person.name].parts = [p.name for p in VocBodyPart] + label_map[VocLabel.person.name].actions = [a.name for a in VocAction] + return label_map + + @staticmethod + def from_categories(categories: CategoriesInfo) -> VocLabelMap: + if AnnotationType.mask not in categories: + # generate a new colormap for the input labels + labels = categories.get(AnnotationType.label, LabelCategories()) + label_map = VocLabelMap((item.name, VocLabelInfo()) for item in labels.items) + + else: + # use the source colormap + labels: LabelCategories = categories[AnnotationType.label] + colors: MaskCategories = categories[AnnotationType.mask] + label_map = VocLabelMap() + for idx, item in enumerate(labels.items): + color = colors.colormap.get(idx) + if color is not None: + label_map[item.name] = VocLabelInfo(color) + + return label_map + + @staticmethod + def parse_from_file(path: str) -> VocLabelMap: + """ + Parses a label map file in the format: + 'name : color (r, g, b) : parts (hand, feet, ...) : actions (siting, standing, ...)' + + Parameters: + path: File path + + Returns: + A dictionary: label -> (color, parts, actions) + """ + + label_map = OrderedDict() + with open(path, "r", encoding="utf-8") as f: + for line in f: + # skip empty and commented lines + line = line.strip() + if not line or line[0] == "#": + continue + + # name : color : parts : actions + label_desc = line.strip().split(":") + if len(label_desc) != 4: raise InvalidAnnotationError( - f"Label '{name}' has wrong color '{color}'. Expected an 'r,g,b' triplet." + f"Label description has wrong number of fields '{len(label_desc)}'. " + "Expected 4 ':'-separated fields." ) - color = tuple(int(c) for c in color) - else: - color = None - - if 2 < len(label_desc) and label_desc[2]: - parts = [s.strip() for s in label_desc[2].split(",")] - else: - parts = [] + name = label_desc[0] - if 3 < len(label_desc) and label_desc[3]: - actions = [s.strip() for s in label_desc[3].split(",")] - else: - actions = [] + if name in label_map: + raise InvalidAnnotationError( + f"Label '{name}' is already defined in the label map" + ) - label_map[name] = [color, parts, actions] - return label_map + if 1 < len(label_desc) and label_desc[1]: + color = label_desc[1].split(",") + if len(color) != 3: + raise InvalidAnnotationError( + f"Label '{name}' has wrong color '{color}'. Expected an 'r,g,b' triplet." + ) + color = tuple(int(c) for c in color) + else: + color = None + if 2 < len(label_desc) and label_desc[2]: + parts = [s.strip() for s in label_desc[2].split(",")] + else: + parts = [] -def parse_meta_file(path: str) -> LabelMapConfig: - # Uses custom format with extra fields - meta_file = path - if osp.isdir(path): - meta_file = get_meta_file(path) + if 3 < len(label_desc) and label_desc[3]: + actions = [s.strip() for s in label_desc[3].split(",")] + else: + actions = [] - dataset_meta = parse_json_file(meta_file) + label_map[name] = [color, parts, actions] + return VocLabelMap(label_map) - label_map = OrderedDict() - parts = dataset_meta.get("parts", {}) - actions = dataset_meta.get("actions", {}) + @staticmethod + def parse_from_meta_file(path: str) -> VocLabelMap: + # Uses a custom format with extra fields + meta_file = path + if osp.isdir(path): + meta_file = get_meta_file(path) - for i, label in enumerate(dataset_meta.get("labels", [])): - label_map[label] = [None, parts.get(str(i), []), actions.get(str(i), [])] + dataset_meta = parse_json_file(meta_file) - colors = dataset_meta.get("segmentation_colors", []) + label_map = OrderedDict() + parts = dataset_meta.get("parts", {}) + actions = dataset_meta.get("actions", {}) - for i, label in enumerate(dataset_meta.get("label_map", {}).values()): - if label not in label_map: - label_map[label] = [None, [], []] + for i, label in enumerate(dataset_meta.get("labels", [])): + label_map[label] = [None, parts.get(str(i), []), actions.get(str(i), [])] - if any(colors) and colors[i] is not None: - label_map[label][0] = tuple(colors[i]) + colors = dataset_meta.get("segmentation_colors", []) - return label_map + for i, label in enumerate(dataset_meta.get("label_map", {}).values()): + if label not in label_map: + label_map[label] = [None, [], []] + if any(colors) and colors[i] is not None: + label_map[label][0] = tuple(colors[i]) -def write_label_map(path: str, label_map: LabelMapConfig): - with open(path, "w", encoding="utf-8") as f: - f.write("# label:color_rgb:parts:actions\n") - for label_name, label_desc in label_map.items(): - if label_desc[0]: - color_rgb = ",".join(str(c) for c in label_desc[0]) - else: - color_rgb = "" + return VocLabelMap(label_map) - parts = ",".join(str(p) for p in label_desc[1]) - actions = ",".join(str(a) for a in label_desc[2]) + def dump_to_file(self, path: str): + with open(path, "w", encoding="utf-8") as f: + f.write("# label:color_rgb:parts:actions\n") + for label_name, label_desc in self.items(): + if label_desc.color: + color_rgb = ",".join(str(c) for c in label_desc.color) + else: + color_rgb = "" - f.write("%s\n" % ":".join([label_name, color_rgb, parts, actions])) + parts = ",".join(str(p) for p in label_desc.parts) + actions = ",".join(str(a) for a in label_desc.actions) + f.write("%s\n" % ":".join([label_name, color_rgb, parts, actions])) -def write_meta_file(path: str, label_map: LabelMapConfig): - # Uses custom format with extra fields - dataset_meta = {} + def dump_to_meta_file(self, output_dir: str): + # Uses custom format with extra fields + dataset_meta = {} - labels = [] - labels_dict = {} - segmentation_colors = [] - parts = {} - actions = {} + labels = [] + labels_dict = {} + segmentation_colors = [] + parts = {} + actions = {} - for i, (label_name, label_desc) in enumerate(label_map.items()): - labels.append(label_name) - if label_desc[0]: - labels_dict[str(i)] = label_name - segmentation_colors.append( - [int(label_desc[0][0]), int(label_desc[0][1]), int(label_desc[0][2])] - ) + for i, (label_name, label_desc) in enumerate(self.items()): + labels.append(label_name) + if label_desc.color: + labels_dict[str(i)] = label_name + segmentation_colors.append(list(map(int, label_desc.color))) - parts[str(i)] = label_desc[1] - actions[str(i)] = label_desc[2] + parts[str(i)] = list(label_desc.parts) + actions[str(i)] = list(label_desc.actions) - dataset_meta["labels"] = labels + dataset_meta["labels"] = labels - if any(segmentation_colors): - dataset_meta["label_map"] = labels_dict - dataset_meta["segmentation_colors"] = segmentation_colors + if any(segmentation_colors): + dataset_meta["label_map"] = labels_dict + dataset_meta["segmentation_colors"] = segmentation_colors - bg_label = find(label_map.items(), lambda x: x[1] == (0, 0, 0)) - if bg_label is not None: - dataset_meta["background_label"] = str(bg_label[0]) + bg_label = self.find_background_label() + if bg_label: + dataset_meta["background_label"] = bg_label - if any(parts): - dataset_meta["parts"] = parts + if any(parts): + dataset_meta["parts"] = parts - if any(actions): - dataset_meta["actions"] = actions + if any(actions): + dataset_meta["actions"] = actions - dump_json_file(get_meta_file(path), dataset_meta) + dump_json_file(get_meta_file(output_dir), dataset_meta) -def make_voc_categories(label_map: Optional[LabelMapConfig] = None) -> CategoriesInfo: +def make_voc_categories(label_map: Optional[Union[VocLabelMap, Dict]] = None) -> CategoriesInfo: if label_map is None: - label_map = make_voc_label_map() + label_map = VocLabelMap.make_default() + elif not isinstance(label_map, VocLabelMap): + label_map = VocLabelMap(label_map) categories = {} @@ -287,18 +388,17 @@ def make_voc_categories(label_map: Optional[LabelMapConfig] = None) -> Categorie label_categories.attributes.update(["difficult", "truncated", "occluded"]) for label, desc in label_map.items(): - label_categories.add(label, attributes=desc[2]) - for part in OrderedDict((k, None) for k in chain(*(desc[1] for desc in label_map.values()))): + label_categories.add(label, attributes=desc.actions) + for part in sorted(set(chain(*(desc.parts for desc in label_map.values())))): label_categories.add(part) categories[AnnotationType.label] = label_categories - has_colors = any(v[0] is not None for v in label_map.values()) - if not has_colors: # generate new colors + if not label_map.has_colors(): # generate new colors colormap = generate_colormap(len(label_map)) else: # only copy defined colors label_id = lambda label: label_categories.find(label)[0] colormap = { - label_id(name): desc[0] for name, desc in label_map.items() if desc[0] is not None + label_id(name): desc.color for name, desc in label_map.items() if desc.color is not None } mask_categories = MaskCategories(colormap) mask_categories.inverse_colormap # pylint: disable=pointless-statement diff --git a/datumaro/util/__init__.py b/datumaro/util/__init__.py index d4643932a6..a7181ac1c4 100644 --- a/datumaro/util/__init__.py +++ b/datumaro/util/__init__.py @@ -1,11 +1,12 @@ # Copyright (C) 2019-2022 Intel Corporation +# Copyright (C) 2022 CVAT.ai Corporation # # SPDX-License-Identifier: MIT from functools import wraps from inspect import isclass from itertools import islice -from typing import Any, Iterable, Optional, Tuple, Union +from typing import Any, Callable, Iterable, Tuple, TypeVar, Union import attrs import orjson @@ -14,8 +15,13 @@ str_to_bool = attrs.converters.to_bool +T = TypeVar("T") +U = TypeVar("U") -def find(iterable, pred=lambda x: True, default=None): + +def find( + iterable: Iterable[T], pred: Callable[[T], bool] = lambda x: True, *, default: U = None +) -> Union[T, U]: return next((x for x in iterable if pred(x)), default) @@ -28,7 +34,7 @@ def cast(value, type_conv, default=None): return default -def to_snake_case(s): +def to_snake_case(s: str) -> str: if not s: return "" @@ -46,7 +52,7 @@ def to_snake_case(s): return "".join(name) -def pairs(iterable): +def pairs(iterable: Iterable[T]) -> Iterable[Tuple[T, T]]: a = iter(iterable) return zip(a, a) diff --git a/tests/cli/test_voc_format.py b/tests/cli/test_voc_format.py index 586b606f1f..a1a2770bfd 100644 --- a/tests/cli/test_voc_format.py +++ b/tests/cli/test_voc_format.py @@ -375,7 +375,7 @@ def test_can_save_and_load_voc_dataset(self): "pose": "Unspecified", }, ), - Bbox(5.5, 6.0, 2.0, 2.0, label=22, id=0, group=1), + Bbox(5.5, 6.0, 2.0, 2.0, label=24, id=0, group=1), Mask(image=np.ones([10, 20]), label=2, group=1), ], ), @@ -414,7 +414,7 @@ def test_can_save_and_load_voc_layout_dataset(self): **{a.name: a.value % 2 == 1 for a in VOC.VocAction}, }, ), - Bbox(5.5, 6.0, 2.0, 2.0, label=22, id=0, group=1), + Bbox(5.5, 6.0, 2.0, 2.0, label=24, id=0, group=1), ], ), DatasetItem( diff --git a/tests/cli/test_yolo_format.py b/tests/cli/test_yolo_format.py index dfe459786f..58d1e8d3f2 100644 --- a/tests/cli/test_yolo_format.py +++ b/tests/cli/test_yolo_format.py @@ -89,7 +89,7 @@ def test_can_convert_voc_to_yolo(self): annotations=[ Bbox(1.0, 2.0, 2.0, 2.0, label=8), Bbox(4.0, 5.0, 2.0, 2.0, label=15), - Bbox(5.5, 6, 2, 2, label=22), + Bbox(5.5, 6, 2, 2, label=24), ], ), DatasetItem( diff --git a/tests/test_extractor_tfds.py b/tests/test_extractor_tfds.py index 397e1d5b33..7d84ecb1fc 100644 --- a/tests/test_extractor_tfds.py +++ b/tests/test_extractor_tfds.py @@ -18,6 +18,7 @@ import tensorflow_datasets as tfds +@skipIf(not TFDS_EXTRACTOR_AVAILABLE, "TFDS is not installed") class TfdsDatasetsTest(TestCase): @mark_requirement(Requirements.DATUM_GENERAL_REQ) def test_metadata(self): diff --git a/tests/test_voc_format.py b/tests/test_voc_format.py index d700b9b7e5..0c738f4c32 100644 --- a/tests/test_voc_format.py +++ b/tests/test_voc_format.py @@ -84,27 +84,27 @@ def test_colormap_generator(self): @mark_requirement(Requirements.DATUM_GENERAL_REQ) def test_can_write_and_parse_labelmap(self): - src_label_map = VOC.make_voc_label_map() - src_label_map["qq"] = [None, ["part1", "part2"], ["act1", "act2"]] - src_label_map["ww"] = [(10, 20, 30), [], ["act3"]] + src_label_map = VOC.VocLabelMap.make_default() + src_label_map["qq"] = VOC.VocLabelInfo(None, ["part1", "part2"], ["act1", "act2"]) + src_label_map["ww"] = VOC.VocLabelInfo((10, 20, 30), [], ["act3"]) with TestDir() as test_dir: file_path = osp.join(test_dir, "test.txt") - VOC.write_label_map(file_path, src_label_map) - dst_label_map = VOC.parse_label_map(file_path) + src_label_map.dump_to_file(file_path) + dst_label_map = VOC.VocLabelMap.parse_from_file(file_path) self.assertEqual(src_label_map, dst_label_map) @mark_requirement(Requirements.DATUM_GENERAL_REQ) def test_can_write_and_parse_dataset_meta_file(self): - src_label_map = VOC.make_voc_label_map() - src_label_map["qq"] = [None, ["part1", "part2"], ["act1", "act2"]] - src_label_map["ww"] = [(10, 20, 30), [], ["act3"]] + src_label_map = VOC.VocLabelMap.make_default() + src_label_map["qq"] = VOC.VocLabelInfo(None, ["part1", "part2"], ["act1", "act2"]) + src_label_map["ww"] = VOC.VocLabelInfo((10, 20, 30), [], ["act3"]) with TestDir() as test_dir: - VOC.write_meta_file(test_dir, src_label_map) - dst_label_map = VOC.parse_meta_file(test_dir) + src_label_map.dump_to_meta_file(test_dir) + dst_label_map = VOC.VocLabelMap.parse_from_meta_file(test_dir) self.assertEqual(src_label_map, dst_label_map) @@ -116,7 +116,7 @@ def test_can_report_invalid_line_in_labelmap(self): f.write("a\n") with self.assertRaisesRegex(InvalidAnnotationError, "Expected 4 ':'-separated fields"): - VOC.parse_label_map(path) + VOC.VocLabelMap.parse_from_file(path) @mark_requirement(Requirements.DATUM_GENERAL_REQ) def test_can_report_repeated_label_in_labelmap(self): @@ -127,7 +127,7 @@ def test_can_report_repeated_label_in_labelmap(self): f.write("a:::\n") with self.assertRaisesRegex(InvalidAnnotationError, "already defined"): - VOC.parse_label_map(path) + VOC.VocLabelMap.parse_from_file(path) @mark_requirement(Requirements.DATUM_GENERAL_REQ) def test_can_report_invalid_color_in_labelmap(self): @@ -137,7 +137,7 @@ def test_can_report_invalid_color_in_labelmap(self): f.write("a:10,20::\n") with self.assertRaisesRegex(InvalidAnnotationError, "Expected an 'r,g,b' triplet"): - VOC.parse_label_map(path) + VOC.VocLabelMap.parse_from_file(path) class TestExtractorBase(Extractor): @@ -287,7 +287,7 @@ def test_can_import_voc_layout_dataset(self): **{a.name: a.value % 2 == 1 for a in VOC.VocAction}, }, ), - Bbox(5.5, 6.0, 2.0, 2.0, label=22, group=2), + Bbox(5.5, 6.0, 2.0, 2.0, label=24, group=2), ], ), DatasetItem( @@ -1666,6 +1666,34 @@ def test_background_masks_dont_introduce_instances_but_cover_others(self): self.assertTrue(np.array_equal([0, 1], np.unique(cls_mask))) self.assertTrue(np.array_equal([0, 1], np.unique(inst_mask))) + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_export_masks_with_non_0_background_color(self): + dataset = Dataset.from_iterable( + [ + DatasetItem( + 1, + media=Image(data=np.zeros((4, 1, 1))), + annotations=[ + Mask([[1, 1, 0, 0]], label=0, attributes={"z_order": 1}), + ], + ) + ], + categories=["fg"], + ) + + label_map = OrderedDict( + fg=[(20, 20, 20), [], []], + background=[(0, 0, 0), [], []], + ) + + with TestDir() as test_dir: + VocConverter.convert(dataset, test_dir, apply_colormap=False, label_map=label_map) + + cls_mask = load_mask(osp.join(test_dir, "SegmentationClass", "1.png")) + inst_mask = load_mask(osp.join(test_dir, "SegmentationObject", "1.png")) + self.assertTrue(np.array_equal([[1, 1, 0, 0]], cls_mask)) + self.assertTrue(np.array_equal([[1, 1, 0, 0]], inst_mask)) + @mark_requirement(Requirements.DATUM_GENERAL_REQ) def test_can_save_dataset_with_image_info(self): class TestExtractor(TestExtractorBase):