diff --git a/datumaro/datumaro/cli/project/__init__.py b/datumaro/datumaro/cli/project/__init__.py index c8fec0ded871..f3ed7589c6ad 100644 --- a/datumaro/datumaro/cli/project/__init__.py +++ b/datumaro/datumaro/cli/project/__init__.py @@ -54,14 +54,14 @@ def build_import_parser(parser): help="Source project format (options: %s)" % (', '.join(importers_list))) parser.add_argument('-d', '--dest', default='.', dest='dst_dir', help="Directory to save the new project to (default: current dir)") - parser.add_argument('extra_args', nargs=argparse.REMAINDER, - help="Additional arguments for importer") parser.add_argument('-n', '--name', default=None, help="Name of the new project (default: same as project dir)") parser.add_argument('--overwrite', action='store_true', help="Overwrite existing files in the save directory") parser.add_argument('--copy', action='store_true', help="Make a deep copy instead of saving source links") + # parser.add_argument('extra_args', nargs=argparse.REMAINDER, + # help="Additional arguments for importer (pass '-- -h' for help)") return parser def import_command(args): @@ -111,8 +111,8 @@ def build_export_parser(parser): help="Output format") parser.add_argument('-p', '--project', dest='project_dir', default='.', help="Directory of the project to operate on (default: current dir)") - parser.add_argument('--save-images', action='store_true', - help="Save images") + parser.add_argument('extra_args', nargs=argparse.REMAINDER, default=None, + help="Additional arguments for converter (pass '-- -h' for help)") return parser def export_command(args): @@ -125,7 +125,7 @@ def export_command(args): save_dir=dst_dir, output_format=args.output_format, filter_expr=args.filter, - save_images=args.save_images) + cmdline_args=args.extra_args) log.info("Project exported to '%s' as '%s'" % \ (dst_dir, args.output_format)) diff --git a/datumaro/datumaro/components/converter.py b/datumaro/datumaro/components/converter.py index 7d2aece0e349..9ea404d962de 100644 --- a/datumaro/datumaro/components/converter.py +++ b/datumaro/datumaro/components/converter.py @@ -4,5 +4,16 @@ # SPDX-License-Identifier: MIT class Converter: + def __init__(self, cmdline_args=None): + pass + def __call__(self, extractor, save_dir): raise NotImplementedError() + + def _parse_cmdline(self, cmdline): + parser = self.build_cmdline_parser() + + if len(cmdline) != 0 and cmdline[0] == '--': + cmdline = cmdline[1:] + args = parser.parse_args(cmdline) + return vars(args) \ No newline at end of file diff --git a/datumaro/datumaro/components/converters/datumaro.py b/datumaro/datumaro/components/converters/datumaro.py index e3d4fcd66980..4e80b5855aa4 100644 --- a/datumaro/datumaro/components/converters/datumaro.py +++ b/datumaro/datumaro/components/converters/datumaro.py @@ -19,7 +19,6 @@ ) from datumaro.components.formats.datumaro import DatumaroPath from datumaro.util.image import save_image -from datumaro.util.mask_tools import apply_colormap def _cast(value, type_conv, default=None): @@ -118,13 +117,6 @@ def _save_mask(self, mask): if mask is None: return mask_id - if self._converter._apply_colormap: - categories = self._converter._extractor.categories() - categories = categories[AnnotationType.mask] - colormap = categories.colormap - - mask = apply_colormap(mask, colormap) - mask_id = self._next_mask_id self._next_mask_id += 1 @@ -232,12 +224,10 @@ def _convert_points_categories(self, obj): return converted class _Converter: - def __init__(self, extractor, save_dir, - save_images=False, apply_colormap=False): + def __init__(self, extractor, save_dir, save_images=False,): self._extractor = extractor self._save_dir = save_dir self._save_images = save_images - self._apply_colormap = apply_colormap def convert(self): os.makedirs(self._save_dir, exist_ok=True) @@ -282,13 +272,27 @@ def _save_image(self, item): save_image(image_path, image) class DatumaroConverter(Converter): - def __init__(self, save_images=False, apply_colormap=False): + def __init__(self, save_images=False, cmdline_args=None): super().__init__() - self._save_images = save_images - self._apply_colormap = apply_colormap + + self._options = { + 'save_images': save_images, + } + + if cmdline_args is not None: + self._options.update(self._parse_cmdline(cmdline_args)) + + @classmethod + def build_cmdline_parser(cls, parser=None): + import argparse + if not parser: + parser = argparse.ArgumentParser() + + parser.add_argument('--save-images', action='store_true', + help="Save images (default: %(default)s)") + + return parser def __call__(self, extractor, save_dir): - converter = _Converter(extractor, save_dir, - apply_colormap=self._apply_colormap, - save_images=self._save_images) + converter = _Converter(extractor, save_dir, **self._options) converter.convert() \ No newline at end of file diff --git a/datumaro/datumaro/components/converters/ms_coco.py b/datumaro/datumaro/components/converters/ms_coco.py index b91662511df6..3354f2f09f3d 100644 --- a/datumaro/datumaro/components/converters/ms_coco.py +++ b/datumaro/datumaro/components/converters/ms_coco.py @@ -14,7 +14,7 @@ from datumaro.components.extractor import ( DEFAULT_SUBSET_NAME, AnnotationType, PointsObject, BboxObject ) -from datumaro.components.formats.ms_coco import CocoAnnotationType, CocoPath +from datumaro.components.formats.ms_coco import CocoTask, CocoPath from datumaro.util import find from datumaro.util.image import save_image import datumaro.util.mask_tools as mask_tools @@ -29,8 +29,9 @@ def _cast(value, type_conv, default=None): return default class _TaskConverter: - def __init__(self): + def __init__(self, context): self._min_ann_id = 1 + self._context = context data = { 'licenses': [], @@ -191,6 +192,13 @@ def save_annotations(self, item): rle = mask_utils.merge(rles) area = mask_utils.area(rle) + if self._context._merge_polygons: + binary_mask = mask_utils.decode(rle).astype(np.bool) + binary_mask = np.asfortranarray(binary_mask, dtype=np.uint8) + segmentation = mask_tools.convert_mask_to_rle(binary_mask) + is_crowd = True + bbox = [int(i) for i in mask_utils.toBbox(rle)] + if ann.group is not None: # Mark the group as visited to prevent repeats for a in annotations[:]: @@ -201,6 +209,18 @@ def save_annotations(self, item): is_crowd = False segmentation = [ann.get_polygon()] area = ann.area() + + if self._context._merge_polygons: + h, w, _ = item.image.shape + rles = mask_utils.frPyObjects(segmentation, h, w) + rle = mask_utils.merge(rles) + area = mask_utils.area(rle) + binary_mask = mask_utils.decode(rle).astype(np.bool) + binary_mask = np.asfortranarray(binary_mask, dtype=np.uint8) + segmentation = mask_tools.convert_mask_to_rle(binary_mask) + is_crowd = True + bbox = [int(i) for i in mask_utils.toBbox(rle)] + if bbox is None: bbox = ann.get_bbox() @@ -340,22 +360,30 @@ def save_annotations(self, item): class _Converter: _TASK_CONVERTER = { - CocoAnnotationType.image_info: _ImageInfoConverter, - CocoAnnotationType.instances: _InstancesConverter, - CocoAnnotationType.person_keypoints: _KeypointsConverter, - CocoAnnotationType.captions: _CaptionsConverter, - CocoAnnotationType.labels: _LabelsConverter, + CocoTask.image_info: _ImageInfoConverter, + CocoTask.instances: _InstancesConverter, + CocoTask.person_keypoints: _KeypointsConverter, + CocoTask.captions: _CaptionsConverter, + CocoTask.labels: _LabelsConverter, } - def __init__(self, extractor, save_dir, save_images=False, task=None): - if not task: - task = list(self._TASK_CONVERTER.keys()) - elif task in CocoAnnotationType: - task = [task] - self._task = task + def __init__(self, extractor, save_dir, + tasks=None, save_images=False, merge_polygons=False): + assert tasks is None or isinstance(tasks, (CocoTask, list)) + if tasks is None: + tasks = list(self._TASK_CONVERTER) + elif isinstance(tasks, CocoTask): + tasks = [tasks] + else: + for t in tasks: + assert t in CocoTask + self._tasks = tasks + self._extractor = extractor self._save_dir = save_dir + self._save_images = save_images + self._merge_polygons = merge_polygons def make_dirs(self): self._images_dir = osp.join(self._save_dir, CocoPath.IMAGES_DIR) @@ -365,11 +393,13 @@ def make_dirs(self): os.makedirs(self._ann_dir, exist_ok=True) def make_task_converter(self, task): - return self._TASK_CONVERTER[task]() + if task not in self._TASK_CONVERTER: + raise NotImplementedError() + return self._TASK_CONVERTER[task](self) def make_task_converters(self): return { - task: self.make_task_converter(task) for task in self._task + task: self.make_task_converter(task) for task in self._tasks } def save_image(self, item, filename): @@ -411,32 +441,56 @@ def convert(self): '%s_%s.json' % (task.name, subset_name))) class CocoConverter(Converter): - def __init__(self, task=None, save_images=False): + def __init__(self, + tasks=None, save_images=False, merge_polygons=False, + cmdline_args=None): super().__init__() - self._task = task - self._save_images = save_images + + self._options = { + 'tasks': tasks, + 'save_images': save_images, + 'merge_polygons': merge_polygons, + } + + if cmdline_args is not None: + self._options.update(self._parse_cmdline(cmdline_args)) + + @staticmethod + def _split_tasks_string(s): + return [CocoTask[i.strip()] for i in s.split(',')] + + @classmethod + def build_cmdline_parser(cls, parser=None): + import argparse + if not parser: + parser = argparse.ArgumentParser() + + parser.add_argument('--save-images', action='store_true', + help="Save images (default: %(default)s)") + parser.add_argument('--merge-polygons', action='store_true', + help="Merge instance polygons into a mask (default: %(default)s)") + parser.add_argument('--tasks', type=cls._split_tasks_string, + default=None, + help="COCO task filter, comma-separated list of {%s} " + "(default: all)" % ', '.join([t.name for t in CocoTask])) + + return parser def __call__(self, extractor, save_dir): - converter = _Converter(extractor, save_dir, - save_images=self._save_images, task=self._task) + converter = _Converter(extractor, save_dir, **self._options) converter.convert() -def CocoInstancesConverter(save_images=False): - return CocoConverter(CocoAnnotationType.instances, - save_images=save_images) +def CocoInstancesConverter(**kwargs): + return CocoConverter(CocoTask.instances, **kwargs) -def CocoImageInfoConverter(save_images=False): - return CocoConverter(CocoAnnotationType.image_info, - save_images=save_images) +def CocoImageInfoConverter(**kwargs): + return CocoConverter(CocoTask.image_info, **kwargs) -def CocoPersonKeypointsConverter(save_images=False): - return CocoConverter(CocoAnnotationType.person_keypoints, - save_images=save_images) +def CocoPersonKeypointsConverter(**kwargs): + return CocoConverter(CocoTask.person_keypoints, **kwargs) -def CocoCaptionsConverter(save_images=False): - return CocoConverter(CocoAnnotationType.captions, - save_images=save_images) +def CocoCaptionsConverter(**kwargs): + return CocoConverter(CocoTask.captions, **kwargs) -def CocoLabelsConverter(save_images=False): - return CocoConverter(CocoAnnotationType.labels, - save_images=save_images) \ No newline at end of file +def CocoLabelsConverter(**kwargs): + return CocoConverter(CocoTask.labels, **kwargs) \ No newline at end of file diff --git a/datumaro/datumaro/components/converters/voc.py b/datumaro/datumaro/components/converters/voc.py index 034749bb7839..2c76be81cca3 100644 --- a/datumaro/datumaro/components/converters/voc.py +++ b/datumaro/datumaro/components/converters/voc.py @@ -31,11 +31,18 @@ class _Converter: _BODY_PARTS = set([entry.name for entry in VocBodyPart]) _ACTIONS = set([entry.name for entry in VocAction]) - def __init__(self, task, extractor, save_dir, - apply_colormap=True, save_images=False): + def __init__(self, extractor, save_dir, + tasks=None, apply_colormap=True, save_images=False): + assert tasks is None or isinstance(tasks, (VocTask, list)) + if tasks is None: + tasks = list(VocTask) + elif isinstance(tasks, VocTask): + tasks = [tasks] + else: + for t in tasks: + assert t in VocTask + self._tasks = tasks - assert not task or task in VocTask - self._task = task self._extractor = extractor self._save_dir = save_dir self._apply_colormap = apply_colormap @@ -205,10 +212,10 @@ def save_subsets(self): objects_with_actions[new_obj_id][action] = presented - if self._task in [None, + if set(self._tasks) & set([None, VocTask.detection, VocTask.person_layout, - VocTask.action_classification]: + VocTask.action_classification]): with open(osp.join(self._ann_dir, item_id + '.xml'), 'w') as f: f.write(ET.tostring(root_elem, encoding='unicode', pretty_print=True)) @@ -245,19 +252,19 @@ def save_subsets(self): action_list[item_id] = None segm_list[item_id] = None - if self._task in [None, + if set(self._tasks) & set([None, VocTask.classification, VocTask.detection, VocTask.action_classification, - VocTask.person_layout]: + VocTask.person_layout]): self.save_clsdet_lists(subset_name, clsdet_list) - if self._task in [None, VocTask.classification]: + if set(self._tasks) & set([None, VocTask.classification]): self.save_class_lists(subset_name, class_lists) - if self._task in [None, VocTask.action_classification]: + if set(self._tasks) & set([None, VocTask.action_classification]): self.save_action_lists(subset_name, action_list) - if self._task in [None, VocTask.person_layout]: + if set(self._tasks) & set([None, VocTask.person_layout]): self.save_layout_lists(subset_name, layout_list) - if self._task in [None, VocTask.segmentation]: + if set(self._tasks) & set([None, VocTask.segmentation]): self.save_segm_lists(subset_name, segm_list) def save_action_lists(self, subset_name, action_list): @@ -337,34 +344,57 @@ def save_segm(self, path, annotation, colormap): save_image(path, data) class VocConverter(Converter): - def __init__(self, task=None, save_images=False, apply_colormap=False): + def __init__(self, + tasks=None, save_images=False, apply_colormap=False, + cmdline_args=None): super().__init__() - self._task = task - self._save_images = save_images - self._apply_colormap = apply_colormap + + self._options = { + 'tasks': tasks, + 'save_images': save_images, + 'apply_colormap': apply_colormap, + } + + if cmdline_args is not None: + self._options.update(self._parse_cmdline(cmdline_args)) + + @staticmethod + def _split_tasks_string(s): + return [VocTask[i.strip()] for i in s.split(',')] + + @classmethod + def build_cmdline_parser(cls, parser=None): + import argparse + if not parser: + parser = argparse.ArgumentParser() + + parser.add_argument('--save-images', action='store_true', + help="Save images (default: %(default)s)") + parser.add_argument('--apply-colormap', type=bool, default=True, + help="Use colormap for class and instance masks " + "(default: %(default)s)") + parser.add_argument('--tasks', type=cls._split_tasks_string, + default=None, + help="VOC task filter, comma-separated list of {%s} " + "(default: all)" % ', '.join([t.name for t in VocTask])) + + return parser def __call__(self, extractor, save_dir): - converter = _Converter(self._task, extractor, save_dir, - apply_colormap=self._apply_colormap, - save_images=self._save_images) + converter = _Converter(extractor, save_dir, **self._options) converter.convert() -def VocClassificationConverter(save_images=False): - return VocConverter(VocTask.classification, - save_images=save_images) +def VocClassificationConverter(**kwargs): + return VocConverter(VocTask.classification, **kwargs) -def VocDetectionConverter(save_images=False): - return VocConverter(VocTask.detection, - save_images=save_images) +def VocDetectionConverter(**kwargs): + return VocConverter(VocTask.detection, **kwargs) -def VocLayoutConverter(save_images=False): - return VocConverter(VocTask.person_layout, - save_images=save_images) +def VocLayoutConverter(**kwargs): + return VocConverter(VocTask.person_layout, **kwargs) -def VocActionConverter(save_images=False): - return VocConverter(VocTask.action_classification, - save_images=save_images) +def VocActionConverter(**kwargs): + return VocConverter(VocTask.action_classification, **kwargs) -def VocSegmentationConverter(save_images=False, apply_colormap=True): - return VocConverter(VocTask.segmentation, - save_images=save_images, apply_colormap=apply_colormap) +def VocSegmentationConverter(**kwargs): + return VocConverter(VocTask.segmentation, **kwargs) diff --git a/datumaro/datumaro/components/converters/yolo.py b/datumaro/datumaro/components/converters/yolo.py index cdf6195363b6..4bf746939d01 100644 --- a/datumaro/datumaro/components/converters/yolo.py +++ b/datumaro/datumaro/components/converters/yolo.py @@ -27,11 +27,26 @@ def _make_yolo_bbox(img_size, box): class YoloConverter(Converter): # https://github.com/AlexeyAB/darknet#how-to-train-to-detect-your-custom-objects - def __init__(self, task=None, save_images=False, apply_colormap=False): + def __init__(self, save_images=False, cmdline_args=None): super().__init__() - self._task = task self._save_images = save_images - self._apply_colormap = apply_colormap + + if cmdline_args is not None: + options = self._parse_cmdline(cmdline_args) + for k, v in options.items(): + if hasattr(self, '_' + str(k)): + setattr(self, '_' + str(k), v) + + @classmethod + def build_cmdline_parser(cls, parser=None): + import argparse + if not parser: + parser = argparse.ArgumentParser() + + parser.add_argument('--save-images', action='store_true', + help="Save images (default: %(default)s)") + + return parser def __call__(self, extractor, save_dir): os.makedirs(save_dir, exist_ok=True) diff --git a/datumaro/datumaro/components/extractors/datumaro.py b/datumaro/datumaro/components/extractors/datumaro.py index 429cb38e094d..6bb336533114 100644 --- a/datumaro/datumaro/components/extractors/datumaro.py +++ b/datumaro/datumaro/components/extractors/datumaro.py @@ -5,6 +5,7 @@ from collections import defaultdict import json +import logging as log import os.path as osp from datumaro.components.extractor import (Extractor, DatasetItem, @@ -140,11 +141,10 @@ def _load_annotations(self, item): mask = None if osp.isfile(mask_path): - mask_cat = self._categories.get(AnnotationType.mask) - if mask_cat is not None: - mask = lazy_mask(mask_path, mask_cat.inverse_colormap) - else: - mask = lazy_image(mask_path) + mask = lazy_mask(mask_path) + else: + log.warn("Not found mask image file '%s', skipped." % \ + mask_path) loaded.append(MaskObject(label=label_id, image=mask, id=ann_id, attributes=attributes, group=group)) diff --git a/datumaro/datumaro/components/extractors/ms_coco.py b/datumaro/datumaro/components/extractors/ms_coco.py index 55d876557c9b..537a297b2d1a 100644 --- a/datumaro/datumaro/components/extractors/ms_coco.py +++ b/datumaro/datumaro/components/extractors/ms_coco.py @@ -16,7 +16,7 @@ BboxObject, CaptionObject, LabelCategories, PointsCategories ) -from datumaro.components.formats.ms_coco import CocoAnnotationType, CocoPath +from datumaro.components.formats.ms_coco import CocoTask, CocoPath from datumaro.util.image import lazy_image @@ -103,9 +103,9 @@ def _load_categories(self): self._categories = {} - label_loader = loaders.get(CocoAnnotationType.labels) - instances_loader = loaders.get(CocoAnnotationType.instances) - person_kp_loader = loaders.get(CocoAnnotationType.person_keypoints) + label_loader = loaders.get(CocoTask.labels) + instances_loader = loaders.get(CocoTask.instances) + person_kp_loader = loaders.get(CocoTask.person_keypoints) if label_loader is None and instances_loader is not None: label_loader = instances_loader @@ -209,7 +209,7 @@ def _parse_annotation(self, ann, ann_type, parsed_annotations, if 'score' in ann: attributes['score'] = ann['score'] - if ann_type is CocoAnnotationType.instances: + if ann_type is CocoTask.instances: x, y, w, h = ann['bbox'] label_id = self._parse_label(ann) group = None @@ -253,13 +253,13 @@ def _parse_annotation(self, ann, ann_type, parsed_annotations, BboxObject(x, y, w, h, label=label_id, id=ann_id, attributes=attributes, group=group) ) - elif ann_type is CocoAnnotationType.labels: + elif ann_type is CocoTask.labels: label_id = self._parse_label(ann) parsed_annotations.append( LabelObject(label=label_id, id=ann_id, attributes=attributes) ) - elif ann_type is CocoAnnotationType.person_keypoints: + elif ann_type is CocoTask.person_keypoints: keypoints = ann['keypoints'] points = [p for i, p in enumerate(keypoints) if i % 3 != 2] visibility = keypoints[2::3] @@ -276,7 +276,7 @@ def _parse_annotation(self, ann, ann_type, parsed_annotations, parsed_annotations.append( BboxObject(*bbox, label=label_id, group=group) ) - elif ann_type is CocoAnnotationType.captions: + elif ann_type is CocoTask.captions: caption = ann['caption'] parsed_annotations.append( CaptionObject(caption, @@ -289,21 +289,21 @@ def _parse_annotation(self, ann, ann_type, parsed_annotations, class CocoImageInfoExtractor(CocoExtractor): def __init__(self, path, **kwargs): - super().__init__(path, task=CocoAnnotationType.image_info, **kwargs) + super().__init__(path, task=CocoTask.image_info, **kwargs) class CocoCaptionsExtractor(CocoExtractor): def __init__(self, path, **kwargs): - super().__init__(path, task=CocoAnnotationType.captions, **kwargs) + super().__init__(path, task=CocoTask.captions, **kwargs) class CocoInstancesExtractor(CocoExtractor): def __init__(self, path, **kwargs): - super().__init__(path, task=CocoAnnotationType.instances, **kwargs) + super().__init__(path, task=CocoTask.instances, **kwargs) class CocoPersonKeypointsExtractor(CocoExtractor): def __init__(self, path, **kwargs): - super().__init__(path, task=CocoAnnotationType.person_keypoints, + super().__init__(path, task=CocoTask.person_keypoints, **kwargs) class CocoLabelsExtractor(CocoExtractor): def __init__(self, path, **kwargs): - super().__init__(path, task=CocoAnnotationType.labels, **kwargs) \ No newline at end of file + super().__init__(path, task=CocoTask.labels, **kwargs) \ No newline at end of file diff --git a/datumaro/datumaro/components/formats/ms_coco.py b/datumaro/datumaro/components/formats/ms_coco.py index bb3bb82bafdb..2a9cddc2c1be 100644 --- a/datumaro/datumaro/components/formats/ms_coco.py +++ b/datumaro/datumaro/components/formats/ms_coco.py @@ -6,11 +6,11 @@ from enum import Enum -CocoAnnotationType = Enum('CocoAnnotationType', [ +CocoTask = Enum('CocoTask', [ 'instances', 'person_keypoints', 'captions', - 'labels', # extension, does not exist in original COCO format + 'labels', # extension, does not exist in the original COCO format 'image_info', 'panoptic', 'stuff', diff --git a/datumaro/datumaro/components/importers/ms_coco.py b/datumaro/datumaro/components/importers/ms_coco.py index 2119cfbdfe09..30d959b0024e 100644 --- a/datumaro/datumaro/components/importers/ms_coco.py +++ b/datumaro/datumaro/components/importers/ms_coco.py @@ -7,16 +7,16 @@ import os import os.path as osp -from datumaro.components.formats.ms_coco import CocoAnnotationType, CocoPath +from datumaro.components.formats.ms_coco import CocoTask, CocoPath class CocoImporter: _COCO_EXTRACTORS = { - CocoAnnotationType.instances: 'coco_instances', - CocoAnnotationType.person_keypoints: 'coco_person_kp', - CocoAnnotationType.captions: 'coco_captions', - CocoAnnotationType.labels: 'coco_labels', - CocoAnnotationType.image_info: 'coco_images', + CocoTask.instances: 'coco_instances', + CocoTask.person_keypoints: 'coco_person_kp', + CocoTask.captions: 'coco_captions', + CocoTask.labels: 'coco_labels', + CocoTask.image_info: 'coco_images', } def __init__(self, task_filter=None): @@ -58,12 +58,12 @@ def find_subsets(dataset_dir): name_parts = osp.splitext(ann_file)[0].rsplit('_', maxsplit=1) ann_type = name_parts[0] try: - ann_type = CocoAnnotationType[ann_type] + ann_type = CocoTask[ann_type] except KeyError: raise Exception( 'Unknown subset type %s, only known are: %s' % \ (ann_type, - ', '.join([e.name for e in CocoAnnotationType]) + ', '.join([e.name for e in CocoTask]) )) subset_name = name_parts[1] subsets[subset_name][ann_type] = subset_path diff --git a/datumaro/datumaro/components/project.py b/datumaro/datumaro/components/project.py index 8bd0ef6427d0..0bce665ee0b4 100644 --- a/datumaro/datumaro/components/project.py +++ b/datumaro/datumaro/components/project.py @@ -542,7 +542,7 @@ def update(self, items): return self def save(self, save_dir=None, merge=False, recursive=True, - save_images=False, apply_colormap=True): + save_images=False): if save_dir is None: assert self.config.project_dir save_dir = self.config.project_dir @@ -562,7 +562,6 @@ def save(self, save_dir=None, merge=False, recursive=True, converter_kwargs = { 'save_images': save_images, - 'apply_colormap': apply_colormap, } if merge: diff --git a/datumaro/datumaro/util/image.py b/datumaro/datumaro/util/image.py index 784a1218772d..02364de90eaf 100644 --- a/datumaro/datumaro/util/image.py +++ b/datumaro/datumaro/util/image.py @@ -28,7 +28,7 @@ def load_image(path): if _IMAGE_BACKEND == _IMAGE_BACKENDS.cv2: import cv2 - image = cv2.imread(path) + image = cv2.imread(path, cv2.IMREAD_UNCHANGED) image = image.astype(np.float32) elif _IMAGE_BACKEND == _IMAGE_BACKENDS.PIL: from PIL import Image @@ -39,13 +39,15 @@ def load_image(path): else: raise NotImplementedError() - assert len(image.shape) == 3 - assert image.shape[2] in [1, 3, 4] + assert len(image.shape) in [2, 3] + if len(image.shape) == 3: + assert image.shape[2] in [3, 4] return image def save_image(path, image, params=None): if _IMAGE_BACKEND == _IMAGE_BACKENDS.cv2: import cv2 + image = image.astype(np.uint8) cv2.imwrite(path, image, params=params) elif _IMAGE_BACKEND == _IMAGE_BACKENDS.PIL: from PIL import Image @@ -109,8 +111,9 @@ def decode_image(image_bytes): else: raise NotImplementedError() - assert len(image.shape) == 3 - assert image.shape[2] in [1, 3, 4] + assert len(image.shape) in [2, 3] + if len(image.shape) == 3: + assert image.shape[2] in [3, 4] return image diff --git a/datumaro/tests/test_coco_format.py b/datumaro/tests/test_coco_format.py index 9bea8a1f2bca..17d7155752f8 100644 --- a/datumaro/tests/test_coco_format.py +++ b/datumaro/tests/test_coco_format.py @@ -138,9 +138,8 @@ def test_can_import(self): self.assertFalse(ann_2_mask is None) class CocoConverterTest(TestCase): - def _test_save_and_load(self, source_dataset, converter_type, test_dir, - importer_params=None): - converter = converter_type() + def _test_save_and_load(self, source_dataset, converter, test_dir, + importer_params=None, target_dataset=None): converter(source_dataset, test_dir.path) if not importer_params: @@ -149,6 +148,8 @@ def _test_save_and_load(self, source_dataset, converter_type, test_dir, **importer_params) parsed_dataset = project.make_dataset() + if target_dataset is not None: + source_dataset = target_dataset source_subsets = [s if s else DEFAULT_SUBSET_NAME for s in source_dataset.subsets()] self.assertListEqual( @@ -195,7 +196,7 @@ def subsets(self): with TestDir() as test_dir: self._test_save_and_load(TestExtractor(), - CocoCaptionsConverter, test_dir) + CocoCaptionsConverter(), test_dir) def test_can_save_and_load_instances(self): class TestExtractor(Extractor): @@ -249,7 +250,7 @@ def categories(self): with TestDir() as test_dir: self._test_save_and_load(TestExtractor(), - CocoInstancesConverter, test_dir) + CocoInstancesConverter(), test_dir) def test_can_save_and_load_instances_with_mask_conversion(self): class TestExtractor(Extractor): @@ -291,9 +292,64 @@ def categories(self): with TestDir() as test_dir: self._test_save_and_load(TestExtractor(), - CocoInstancesConverter, test_dir, + CocoInstancesConverter(), test_dir, {'merge_instance_polygons': True}) + def test_can_merge_instance_polygons_to_mask_in_coverter(self): + label_categories = LabelCategories() + for i in range(10): + label_categories.add(str(i)) + + class SrcTestExtractor(Extractor): + def __iter__(self): + items = [ + DatasetItem(id=0, image=np.zeros((5, 10, 3)), + annotations=[ + PolygonObject([0, 0, 4, 0, 4, 4], + label=3, id=4, group=4, + attributes={ 'is_crowd': False }), + PolygonObject([5, 0, 9, 0, 5, 5], + label=3, id=4, group=4, + attributes={ 'is_crowd': False }), + ] + ), + ] + return iter(items) + + def categories(self): + return { AnnotationType.label: label_categories } + + class DstTestExtractor(Extractor): + def __iter__(self): + items = [ + DatasetItem(id=0, image=np.zeros((5, 10, 3)), + annotations=[ + BboxObject(1, 0, 8, 4, label=3, id=4, group=4, + attributes={ 'is_crowd': True }), + MaskObject(np.array([ + [0, 1, 1, 1, 0, 1, 1, 1, 1, 0], + [0, 0, 1, 1, 0, 1, 1, 1, 0, 0], + [0, 0, 0, 1, 0, 1, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], + # only internal fragment (without the border), + # but not everywhere... + dtype=np.bool), + attributes={ 'is_crowd': True }, + label=3, id=4, group=4), + ] + ), + ] + return iter(items) + + def categories(self): + return { AnnotationType.label: label_categories } + + with TestDir() as test_dir: + self._test_save_and_load(SrcTestExtractor(), + CocoInstancesConverter(merge_polygons=True), test_dir, + target_dataset=DstTestExtractor()) + def test_can_save_and_load_images(self): class TestExtractor(Extractor): def __iter__(self): @@ -314,7 +370,7 @@ def subsets(self): with TestDir() as test_dir: self._test_save_and_load(TestExtractor(), - CocoImageInfoConverter, test_dir) + CocoImageInfoConverter(), test_dir) def test_can_save_and_load_labels(self): class TestExtractor(Extractor): @@ -350,7 +406,7 @@ def categories(self): with TestDir() as test_dir: self._test_save_and_load(TestExtractor(), - CocoLabelsConverter, test_dir) + CocoLabelsConverter(), test_dir) def test_can_save_and_load_keypoints(self): class TestExtractor(Extractor): @@ -397,7 +453,7 @@ def categories(self): with TestDir() as test_dir: self._test_save_and_load(TestExtractor(), - CocoPersonKeypointsConverter, test_dir) + CocoPersonKeypointsConverter(), test_dir) def test_can_save_dataset_with_no_subsets(self): class TestExtractor(Extractor): @@ -429,4 +485,4 @@ def categories(self): with TestDir() as test_dir: self._test_save_and_load(TestExtractor(), - CocoConverter, test_dir) \ No newline at end of file + CocoConverter(), test_dir) \ No newline at end of file diff --git a/datumaro/tests/test_datumaro_format.py b/datumaro/tests/test_datumaro_format.py index 63907be6f449..3402ccba22f3 100644 --- a/datumaro/tests/test_datumaro_format.py +++ b/datumaro/tests/test_datumaro_format.py @@ -18,7 +18,7 @@ class DatumaroConverterTest(TestCase): class TestExtractor(Extractor): def __iter__(self): items = [ - DatasetItem(id=100, subset='train', + DatasetItem(id=100, subset='train', image=np.ones((10, 6, 3)), annotations=[ CaptionObject('hello', id=1), CaptionObject('world', id=2, group=5), @@ -75,8 +75,7 @@ def test_can_save_and_load(self): with TestDir() as test_dir: source_dataset = self.TestExtractor() - converter = DatumaroConverter( - save_images=True, apply_colormap=True) + converter = DatumaroConverter(save_images=True) converter(source_dataset, test_dir.path) project = Project.import_from(test_dir.path, 'datumaro') diff --git a/datumaro/tests/test_image.py b/datumaro/tests/test_image.py index 143d6c4e41e4..424fd9c88dd1 100644 --- a/datumaro/tests/test_image.py +++ b/datumaro/tests/test_image.py @@ -17,9 +17,12 @@ def tearDown(self): def test_save_and_load_backends(self): backends = image_module._IMAGE_BACKENDS - for save_backend, load_backend in product(backends, backends): + for save_backend, load_backend, c in product(backends, backends, [1, 3]): with TestDir() as test_dir: - src_image = np.random.randint(0, 255 + 1, (2, 4, 3)) + if c == 1: + src_image = np.random.randint(0, 255 + 1, (2, 4)) + else: + src_image = np.random.randint(0, 255 + 1, (2, 4, c)) path = osp.join(test_dir.path, 'img.png') # lossless image_module._IMAGE_BACKEND = save_backend @@ -33,8 +36,11 @@ def test_save_and_load_backends(self): def test_encode_and_decode_backends(self): backends = image_module._IMAGE_BACKENDS - for save_backend, load_backend in product(backends, backends): - src_image = np.random.randint(0, 255 + 1, (2, 4, 3)) + for save_backend, load_backend, c in product(backends, backends, [1, 3]): + if c == 1: + src_image = np.random.randint(0, 255 + 1, (2, 4)) + else: + src_image = np.random.randint(0, 255 + 1, (2, 4, c)) image_module._IMAGE_BACKEND = save_backend buffer = image_module.encode_image(src_image, '.png') # lossless