From 2ebca5bb5a283433580a01335a378318f5dbb069 Mon Sep 17 00:00:00 2001 From: zhiltsov-max Date: Sat, 7 Mar 2020 20:57:45 +0300 Subject: [PATCH] [Datumaro] Dataset format auto detection (#1242) * Add dataset format detection * Add auto format detection for import * Split VOC extractor --- .../datumaro/cli/contexts/project/__init__.py | 61 +- datumaro/datumaro/components/extractor.py | 4 + .../datumaro/plugins/coco_format/importer.py | 8 +- .../datumaro/plugins/cvat_format/importer.py | 25 +- .../plugins/datumaro_format/importer.py | 26 +- .../tf_detection_api_format/importer.py | 17 +- .../datumaro/plugins/voc_format/converter.py | 18 +- .../datumaro/plugins/voc_format/extractor.py | 873 +++++------------- .../datumaro/plugins/voc_format/importer.py | 69 +- .../datumaro/plugins/yolo_format/importer.py | 22 +- datumaro/datumaro/util/log_utils.py | 16 + datumaro/tests/test_coco_format.py | 6 + datumaro/tests/test_cvat_format.py | 170 ++-- datumaro/tests/test_datumaro_format.py | 9 +- datumaro/tests/test_tfrecord_format.py | 29 + datumaro/tests/test_voc_format.py | 50 +- datumaro/tests/test_yolo_format.py | 26 + 17 files changed, 572 insertions(+), 857 deletions(-) create mode 100644 datumaro/datumaro/util/log_utils.py diff --git a/datumaro/datumaro/cli/contexts/project/__init__.py b/datumaro/datumaro/cli/contexts/project/__init__.py index 61bb77ad433d..742413cc0d33 100644 --- a/datumaro/datumaro/cli/contexts/project/__init__.py +++ b/datumaro/datumaro/cli/contexts/project/__init__.py @@ -132,8 +132,8 @@ def build_import_parser(parser_ctor=argparse.ArgumentParser): help="Overwrite existing files in the save directory") parser.add_argument('-i', '--input-path', required=True, dest='source', help="Path to import project from") - parser.add_argument('-f', '--format', required=True, - help="Source project format") + parser.add_argument('-f', '--format', + help="Source project format. Will try to detect, if not specified.") parser.add_argument('extra_args', nargs=argparse.REMAINDER, help="Additional arguments for importer (pass '-- -h' for help)") parser.set_defaults(command=import_command) @@ -164,22 +164,53 @@ def import_command(args): if project_name is None: project_name = osp.basename(project_dir) - try: - env = Environment() - importer = env.make_importer(args.format) - except KeyError: - raise CliException("Importer for format '%s' is not found" % \ - args.format) - - extra_args = {} - if hasattr(importer, 'from_cmdline'): - extra_args = importer.from_cmdline(args.extra_args) + env = Environment() + log.info("Importing project from '%s'" % args.source) + + if not args.format: + if args.extra_args: + raise CliException("Extra args can not be used without format") + + log.info("Trying to detect dataset format...") + + matches = [] + for format_name in env.importers.items: + log.debug("Checking '%s' format...", format_name) + importer = env.make_importer(format_name) + try: + match = importer.detect(args.source) + if match: + log.debug("format matched") + matches.append((format_name, importer)) + except NotImplementedError: + log.debug("Format '%s' does not support auto detection.", + format_name) + + if len(matches) == 0: + log.error("Failed to detect dataset format automatically. " + "Try to specify format with '-f/--format' parameter.") + return 1 + elif len(matches) != 1: + log.error("Multiple formats match the dataset: %s. " + "Try to specify format with '-f/--format' parameter.", + ', '.join(m[0] for m in matches)) + return 2 + + format_name, importer = matches[0] + args.format = format_name + else: + try: + importer = env.make_importer(args.format) + if hasattr(importer, 'from_cmdline'): + extra_args = importer.from_cmdline(args.extra_args) + except KeyError: + raise CliException("Importer for format '%s' is not found" % \ + args.format) - log.info("Importing project from '%s' as '%s'" % \ - (args.source, args.format)) + log.info("Importing project as '%s'" % args.format) source = osp.abspath(args.source) - project = importer(source, **extra_args) + project = importer(source, **locals().get('extra_args', {})) project.config.project_name = project_name project.config.project_dir = project_dir diff --git a/datumaro/datumaro/components/extractor.py b/datumaro/datumaro/components/extractor.py index bd5b39fa962f..247807539b9b 100644 --- a/datumaro/datumaro/components/extractor.py +++ b/datumaro/datumaro/components/extractor.py @@ -743,6 +743,10 @@ class SourceExtractor(Extractor): pass class Importer: + @classmethod + def detect(cls, path): + raise NotImplementedError() + def __call__(self, path, **extra_params): raise NotImplementedError() diff --git a/datumaro/datumaro/plugins/coco_format/importer.py b/datumaro/datumaro/plugins/coco_format/importer.py index bb129d7aed5a..42e5be9fc89d 100644 --- a/datumaro/datumaro/plugins/coco_format/importer.py +++ b/datumaro/datumaro/plugins/coco_format/importer.py @@ -9,6 +9,7 @@ import os.path as osp from datumaro.components.extractor import Importer +from datumaro.util.log_utils import logging_disabled from .format import CocoTask, CocoPath @@ -22,6 +23,11 @@ class CocoImporter(Importer): CocoTask.image_info: 'coco_image_info', } + @classmethod + def detect(cls, path): + with logging_disabled(log.WARN): + return len(cls.find_subsets(path)) != 0 + def __call__(self, path, **extra_params): from datumaro.components.project import Project # cyclic import project = Project() @@ -53,7 +59,7 @@ def find_subsets(path): if osp.basename(osp.normpath(path)) != CocoPath.ANNOTATIONS_DIR: path = osp.join(path, CocoPath.ANNOTATIONS_DIR) - subset_paths += glob(osp.join(path, '*_*.json')) + subset_paths += glob(osp.join(path, '*_*.json')) subsets = defaultdict(dict) for subset_path in subset_paths: diff --git a/datumaro/datumaro/plugins/cvat_format/importer.py b/datumaro/datumaro/plugins/cvat_format/importer.py index a81b5cb38cec..79be0c610524 100644 --- a/datumaro/datumaro/plugins/cvat_format/importer.py +++ b/datumaro/datumaro/plugins/cvat_format/importer.py @@ -15,18 +15,15 @@ class CvatImporter(Importer): EXTRACTOR_NAME = 'cvat' + @classmethod + def detect(cls, path): + return len(cls.find_subsets(path)) != 0 + def __call__(self, path, **extra_params): from datumaro.components.project import Project # cyclic import project = Project() - if path.endswith('.xml') and osp.isfile(path): - subset_paths = [path] - else: - subset_paths = glob(osp.join(path, '*.xml')) - - if osp.basename(osp.normpath(path)) != CvatPath.ANNOTATIONS_DIR: - path = osp.join(path, CvatPath.ANNOTATIONS_DIR) - subset_paths += glob(osp.join(path, '*.xml')) + subset_paths = self.find_subsets(path) if len(subset_paths) == 0: raise Exception("Failed to find 'cvat' dataset at '%s'" % path) @@ -46,3 +43,15 @@ def __call__(self, path, **extra_params): }) return project + + @staticmethod + def find_subsets(path): + if path.endswith('.xml') and osp.isfile(path): + subset_paths = [path] + else: + subset_paths = glob(osp.join(path, '*.xml')) + + if osp.basename(osp.normpath(path)) != CvatPath.ANNOTATIONS_DIR: + path = osp.join(path, CvatPath.ANNOTATIONS_DIR) + subset_paths += glob(osp.join(path, '*.xml')) + return subset_paths \ No newline at end of file diff --git a/datumaro/datumaro/plugins/datumaro_format/importer.py b/datumaro/datumaro/plugins/datumaro_format/importer.py index 0184ef9040f4..ed2f75275f0f 100644 --- a/datumaro/datumaro/plugins/datumaro_format/importer.py +++ b/datumaro/datumaro/plugins/datumaro_format/importer.py @@ -15,19 +15,15 @@ class DatumaroImporter(Importer): EXTRACTOR_NAME = 'datumaro' + @classmethod + def detect(cls, path): + return len(cls.find_subsets(path)) != 0 + def __call__(self, path, **extra_params): from datumaro.components.project import Project # cyclic import project = Project() - if path.endswith('.json') and osp.isfile(path): - subset_paths = [path] - else: - subset_paths = glob(osp.join(path, '*.json')) - - if osp.basename(osp.normpath(path)) != DatumaroPath.ANNOTATIONS_DIR: - path = osp.join(path, DatumaroPath.ANNOTATIONS_DIR) - subset_paths += glob(osp.join(path, '*.json')) - + subset_paths = self.find_subsets(path) if len(subset_paths) == 0: raise Exception("Failed to find 'datumaro' dataset at '%s'" % path) @@ -46,3 +42,15 @@ def __call__(self, path, **extra_params): }) return project + + @staticmethod + def find_subsets(path): + if path.endswith('.json') and osp.isfile(path): + subset_paths = [path] + else: + subset_paths = glob(osp.join(path, '*.json')) + + if osp.basename(osp.normpath(path)) != DatumaroPath.ANNOTATIONS_DIR: + path = osp.join(path, DatumaroPath.ANNOTATIONS_DIR) + subset_paths += glob(osp.join(path, '*.json')) + return subset_paths \ No newline at end of file diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/importer.py b/datumaro/datumaro/plugins/tf_detection_api_format/importer.py index 3000c8881635..9783a23cb15d 100644 --- a/datumaro/datumaro/plugins/tf_detection_api_format/importer.py +++ b/datumaro/datumaro/plugins/tf_detection_api_format/importer.py @@ -13,15 +13,15 @@ class TfDetectionApiImporter(Importer): EXTRACTOR_NAME = 'tf_detection_api' + @classmethod + def detect(cls, path): + return len(cls.find_subsets(path)) != 0 + def __call__(self, path, **extra_params): from datumaro.components.project import Project # cyclic import project = Project() - if path.endswith('.tfrecord') and osp.isfile(path): - subset_paths = [path] - else: - subset_paths = glob(osp.join(path, '*.tfrecord')) - + subset_paths = self.find_subsets(path) if len(subset_paths) == 0: raise Exception( "Failed to find 'tf_detection_api' dataset at '%s'" % path) @@ -42,3 +42,10 @@ def __call__(self, path, **extra_params): return project + @staticmethod + def find_subsets(path): + if path.endswith('.tfrecord') and osp.isfile(path): + subset_paths = [path] + else: + subset_paths = glob(osp.join(path, '*.tfrecord')) + return subset_paths \ No newline at end of file diff --git a/datumaro/datumaro/plugins/voc_format/converter.py b/datumaro/datumaro/plugins/voc_format/converter.py index 5467e52ffb85..729383d5d889 100644 --- a/datumaro/datumaro/plugins/voc_format/converter.py +++ b/datumaro/datumaro/plugins/voc_format/converter.py @@ -317,6 +317,9 @@ def save_subsets(self): self.save_segm_lists(subset_name, segm_list) def save_action_lists(self, subset_name, action_list): + if not action_list: + return + os.makedirs(self._action_subsets_dir, exist_ok=True) ann_file = osp.join(self._action_subsets_dir, subset_name + '.txt') @@ -342,11 +345,11 @@ def save_action_lists(self, subset_name, action_list): (item, 1 + obj_id, 1 if presented else -1)) def save_class_lists(self, subset_name, class_lists): - os.makedirs(self._cls_subsets_dir, exist_ok=True) - - if len(class_lists) == 0: + if not class_lists: return + os.makedirs(self._cls_subsets_dir, exist_ok=True) + for label in self._label_map: ann_file = osp.join(self._cls_subsets_dir, '%s_%s.txt' % (label, subset_name)) @@ -360,6 +363,9 @@ def save_class_lists(self, subset_name, class_lists): f.write('%s % d\n' % (item, 1 if presented else -1)) def save_clsdet_lists(self, subset_name, clsdet_list): + if not clsdet_list: + return + os.makedirs(self._cls_subsets_dir, exist_ok=True) ann_file = osp.join(self._cls_subsets_dir, subset_name + '.txt') @@ -368,6 +374,9 @@ def save_clsdet_lists(self, subset_name, clsdet_list): f.write('%s\n' % item) def save_segm_lists(self, subset_name, segm_list): + if not segm_list: + return + os.makedirs(self._segm_subsets_dir, exist_ok=True) ann_file = osp.join(self._segm_subsets_dir, subset_name + '.txt') @@ -376,6 +385,9 @@ def save_segm_lists(self, subset_name, segm_list): f.write('%s\n' % item) def save_layout_lists(self, subset_name, layout_list): + if not layout_list: + return + os.makedirs(self._layout_subsets_dir, exist_ok=True) ann_file = osp.join(self._layout_subsets_dir, subset_name + '.txt') diff --git a/datumaro/datumaro/plugins/voc_format/extractor.py b/datumaro/datumaro/plugins/voc_format/extractor.py index 31047e8532e3..8823426ef69d 100644 --- a/datumaro/datumaro/plugins/voc_format/extractor.py +++ b/datumaro/datumaro/plugins/voc_format/extractor.py @@ -5,16 +5,16 @@ from collections import defaultdict import logging as log -import os +import numpy as np import os.path as osp from xml.etree import ElementTree as ET -from datumaro.components.extractor import (SourceExtractor, Extractor, +from datumaro.components.extractor import (SourceExtractor, DEFAULT_SUBSET_NAME, DatasetItem, AnnotationType, Label, Mask, Bbox, CompiledMask ) from datumaro.util import dir_items -from datumaro.util.image import lazy_image, Image +from datumaro.util.image import Image from datumaro.util.mask_tools import lazy_mask, invert_colormap from .format import ( @@ -24,717 +24,276 @@ _inverse_inst_colormap = invert_colormap(VocInstColormap) -class VocExtractor(SourceExtractor): - class Subset(Extractor): - def __init__(self, name, parent): - super().__init__() - self._parent = parent - self._name = name - self.items = [] - - def __iter__(self): - for item_id in self.items: - yield self._parent._get(item_id, self._name) - - def __len__(self): - return len(self.items) - - def categories(self): - return self._parent.categories() - - def _load_subsets(self, subsets_dir): - dir_files = dir_items(subsets_dir, '.txt', truncate_ext=True) - subset_names = [s for s in dir_files if '_' not in s] - - subsets = {} - for subset_name in subset_names: - subset_file_name = subset_name - if subset_name == DEFAULT_SUBSET_NAME: - subset_name = None - subset = __class__.Subset(subset_name, self) - - subset.items = [] - with open(osp.join(subsets_dir, subset_file_name + '.txt'), 'r') as f: - for line in f: - line = line.split()[0].strip() - if line: - subset.items.append(line) - - subsets[subset_name] = subset - return subsets - - def _load_cls_annotations(self, subsets_dir, subset_names): - subset_file_names = [n if n else DEFAULT_SUBSET_NAME - for n in subset_names] - dir_files = dir_items(subsets_dir, '.txt', truncate_ext=True) - - label_annotations = defaultdict(list) - label_anno_files = [s for s in dir_files \ - if '_' in s and s[s.rfind('_') + 1:] in subset_file_names] - for ann_filename in label_anno_files: - with open(osp.join(subsets_dir, ann_filename + '.txt'), 'r') as f: - label = ann_filename[:ann_filename.rfind('_')] - label_id = self._get_label_id(label) - for line in f: - item, present = line.split() - if present == '1': - label_annotations[item].append(label_id) - - self._annotations[VocTask.classification] = dict(label_annotations) - - def _load_det_annotations(self): - det_anno_dir = osp.join(self._path, VocPath.ANNOTATIONS_DIR) - det_anno_items = dir_items(det_anno_dir, '.xml', truncate_ext=True) - det_annotations = dict() - for ann_item in det_anno_items: - with open(osp.join(det_anno_dir, ann_item + '.xml'), 'r') as f: - ann_file_data = f.read() - det_annotations[ann_item] = ann_file_data - - self._annotations[VocTask.detection] = det_annotations - - def _load_categories(self): - label_map = None - label_map_path = osp.join(self._path, VocPath.LABELMAP_FILE) - if osp.isfile(label_map_path): - label_map = parse_label_map(label_map_path) - self._categories = make_voc_categories(label_map) - - def __init__(self, path, task): +class _VocExtractor(SourceExtractor): + def __init__(self, path): super().__init__() + assert osp.isfile(path), path self._path = path - self._subsets = {} - self._categories = {} - self._annotations = {} - self._task = task + self._dataset_dir = osp.dirname(osp.dirname(osp.dirname(path))) + + subset = osp.splitext(osp.basename(path))[0] + if subset == DEFAULT_SUBSET_NAME: + subset = None + self._subset = subset - self._load_categories() + self._categories = self._load_categories(self._dataset_dir) + log.debug("Loaded labels: %s", ', '.join("'%s'" % l.name + for l in self._categories[AnnotationType.label].items)) + self._items = self._load_subset_list(path) + + def categories(self): + return self._categories def __len__(self): - length = 0 - for subset in self._subsets.values(): - length += len(subset) - return length + return len(self._items) def subsets(self): - return list(self._subsets) + if self._subset: + return [self._subset] + return None def get_subset(self, name): - return self._subsets[name] - - def categories(self): - return self._categories - - def __iter__(self): - for subset in self._subsets.values(): - for item in subset: - yield item - - def _get(self, item_id, subset_name): - image = osp.join(self._path, VocPath.IMAGES_DIR, - item_id + VocPath.IMAGE_EXT) - det_annotations = self._annotations.get(VocTask.detection) - if det_annotations is not None: - det_annotations = det_annotations.get(item_id) - if det_annotations is not None: - root_elem = ET.fromstring(det_annotations) - height = root_elem.find('size/height') - if height is not None: - height = int(height.text) - width = root_elem.find('size/width') - if width is not None: - width = int(width.text) - if height and width: - image = Image(path=image, size=(height, width)) - - annotations = self._get_annotations(item_id) - - return DatasetItem(annotations=annotations, - id=item_id, subset=subset_name, image=image) + if name != self._subset: + return None + return self def _get_label_id(self, label): label_id, _ = self._categories[AnnotationType.label].find(label) - if label_id is None: - log.debug("Unknown label '%s'. Loaded labels: %s", - label, - ', '.join("'%s'" % s.name - for s in self._categories[AnnotationType.label].items)) - raise Exception("Unknown label '%s'" % label) + assert label_id is not None, label return label_id @staticmethod - def _lazy_extract_mask(mask, c): - return lambda: mask == c + def _load_categories(dataset_path): + label_map = None + label_map_path = osp.join(dataset_path, VocPath.LABELMAP_FILE) + if osp.isfile(label_map_path): + label_map = parse_label_map(label_map_path) + return make_voc_categories(label_map) + + @staticmethod + def _load_subset_list(subset_path): + with open(subset_path) as f: + return [line.split()[0] for line in f] + +class VocClassificationExtractor(_VocExtractor): + def __iter__(self): + raw_anns = self._load_annotations() + for item_id in self._items: + image = osp.join(self._dataset_dir, VocPath.IMAGES_DIR, + item_id + VocPath.IMAGE_EXT) + anns = self._parse_annotations(raw_anns, item_id) + yield DatasetItem(id=item_id, subset=self._subset, + image=image, annotations=anns) + + def _load_annotations(self): + annotations = defaultdict(list) + task_dir = osp.dirname(self._path) + anno_files = [s for s in dir_items(task_dir, '.txt') + if s.endswith('_' + osp.basename(self._path))] + for ann_filename in anno_files: + with open(osp.join(task_dir, ann_filename)) as f: + label = ann_filename[:ann_filename.rfind('_')] + label_id = self._get_label_id(label) + for line in f: + item, present = line.split() + if present == '1': + annotations[item].append(label_id) + + return dict(annotations) + + @staticmethod + def _parse_annotations(raw_anns, item_id): + return [Label(label_id) for label_id in raw_anns.get(item_id, [])] - def _get_annotations(self, item_id): +class _VocXmlExtractor(_VocExtractor): + def __init__(self, path, task): + super().__init__(path) + self._task = task + + def __iter__(self): + anno_dir = osp.join(self._dataset_dir, VocPath.ANNOTATIONS_DIR) + + for item_id in self._items: + image = osp.join(self._dataset_dir, VocPath.IMAGES_DIR, + item_id + VocPath.IMAGE_EXT) + + anns = [] + ann_file = osp.join(anno_dir, item_id + '.xml') + if osp.isfile(ann_file): + root_elem = ET.parse(ann_file) + height = root_elem.find('size/height') + if height is not None: + height = int(height.text) + width = root_elem.find('size/width') + if width is not None: + width = int(width.text) + if height and width: + image = Image(path=image, size=(height, width)) + + anns = self._parse_annotations(root_elem) + + yield DatasetItem(id=item_id, subset=self._subset, + image=image, annotations=anns) + + def _parse_annotations(self, root_elem): item_annotations = [] - if self._task is VocTask.segmentation: - class_mask = None - segm_path = osp.join(self._path, VocPath.SEGMENTATION_DIR, - item_id + VocPath.SEGM_EXT) - if osp.isfile(segm_path): - inverse_cls_colormap = \ - self._categories[AnnotationType.mask].inverse_colormap - class_mask = lazy_mask(segm_path, inverse_cls_colormap) - - instances_mask = None - inst_path = osp.join(self._path, VocPath.INSTANCES_DIR, - item_id + VocPath.SEGM_EXT) - if osp.isfile(inst_path): - instances_mask = lazy_mask(inst_path, _inverse_inst_colormap) - - if instances_mask is not None: - compiled_mask = CompiledMask(class_mask, instances_mask) - - if class_mask is not None: - label_cat = self._categories[AnnotationType.label] - instance_labels = compiled_mask.get_instance_labels( - class_count=len(label_cat.items)) - else: - instance_labels = {i: None - for i in range(compiled_mask.instance_count)} - - for instance_id, label_id in instance_labels.items(): - image = compiled_mask.lazy_extract(instance_id) - - attributes = dict() - if label_id is not None: - actions = {a: False - for a in label_cat.items[label_id].attributes - } - attributes.update(actions) - - item_annotations.append(Mask( - image=image, label=label_id, - attributes=attributes, group=instance_id - )) - elif class_mask is not None: - log.warn("item '%s': has only class segmentation, " - "instance masks will not be available" % item_id) - classes = class_mask.image.unique() - for label_id in classes: - image = self._lazy_extract_mask(class_mask, label_id) - item_annotations.append(Mask(image=image, label=label_id)) - - cls_annotations = self._annotations.get(VocTask.classification) - if cls_annotations is not None and \ - self._task is VocTask.classification: - item_labels = cls_annotations.get(item_id) - if item_labels is not None: - for label_id in item_labels: - item_annotations.append(Label(label_id)) - - det_annotations = self._annotations.get(VocTask.detection) - if det_annotations is not None: - det_annotations = det_annotations.get(item_id) - if det_annotations is not None: - root_elem = ET.fromstring(det_annotations) - - for obj_id, object_elem in enumerate(root_elem.findall('object')): - obj_id += 1 - attributes = {} - group = obj_id + for obj_id, object_elem in enumerate(root_elem.findall('object')): + obj_id += 1 + attributes = {} + group = obj_id - obj_label_id = None - label_elem = object_elem.find('name') - if label_elem is not None: - obj_label_id = self._get_label_id(label_elem.text) + obj_label_id = None + label_elem = object_elem.find('name') + if label_elem is not None: + obj_label_id = self._get_label_id(label_elem.text) - obj_bbox = self._parse_bbox(object_elem) + obj_bbox = self._parse_bbox(object_elem) - if obj_label_id is None or obj_bbox is None: - continue + if obj_label_id is None or obj_bbox is None: + continue - difficult_elem = object_elem.find('difficult') - attributes['difficult'] = difficult_elem is not None and \ - difficult_elem.text == '1' - - truncated_elem = object_elem.find('truncated') - attributes['truncated'] = truncated_elem is not None and \ - truncated_elem.text == '1' - - occluded_elem = object_elem.find('occluded') - attributes['occluded'] = occluded_elem is not None and \ - occluded_elem.text == '1' - - pose_elem = object_elem.find('pose') - if pose_elem is not None: - attributes['pose'] = pose_elem.text - - point_elem = object_elem.find('point') - if point_elem is not None: - point_x = point_elem.find('x') - point_y = point_elem.find('y') - point = [float(point_x.text), float(point_y.text)] - attributes['point'] = point - - actions_elem = object_elem.find('actions') - actions = {a: False - for a in self._categories[AnnotationType.label] \ - .items[obj_label_id].attributes} - if actions_elem is not None: - for action_elem in actions_elem: - actions[action_elem.tag] = (action_elem.text == '1') - for action, present in actions.items(): - attributes[action] = present - - has_parts = False - for part_elem in object_elem.findall('part'): - part = part_elem.find('name').text - part_label_id = self._get_label_id(part) - part_bbox = self._parse_bbox(part_elem) - - if self._task is not VocTask.person_layout: - break - if part_bbox is None: - continue - has_parts = True - item_annotations.append(Bbox(*part_bbox, label=part_label_id, - group=group)) - - if self._task is VocTask.person_layout and not has_parts: - continue - if self._task is VocTask.action_classification and not actions: + difficult_elem = object_elem.find('difficult') + attributes['difficult'] = difficult_elem is not None and \ + difficult_elem.text == '1' + + truncated_elem = object_elem.find('truncated') + attributes['truncated'] = truncated_elem is not None and \ + truncated_elem.text == '1' + + occluded_elem = object_elem.find('occluded') + attributes['occluded'] = occluded_elem is not None and \ + occluded_elem.text == '1' + + pose_elem = object_elem.find('pose') + if pose_elem is not None: + attributes['pose'] = pose_elem.text + + point_elem = object_elem.find('point') + if point_elem is not None: + point_x = point_elem.find('x') + point_y = point_elem.find('y') + point = [float(point_x.text), float(point_y.text)] + attributes['point'] = point + + actions_elem = object_elem.find('actions') + actions = {a: False + for a in self._categories[AnnotationType.label] \ + .items[obj_label_id].attributes} + if actions_elem is not None: + for action_elem in actions_elem: + actions[action_elem.tag] = (action_elem.text == '1') + for action, present in actions.items(): + attributes[action] = present + + has_parts = False + for part_elem in object_elem.findall('part'): + part = part_elem.find('name').text + part_label_id = self._get_label_id(part) + part_bbox = self._parse_bbox(part_elem) + + if self._task is not VocTask.person_layout: + break + if part_bbox is None: continue + has_parts = True + item_annotations.append(Bbox(*part_bbox, label=part_label_id, + group=group)) + + if self._task is VocTask.person_layout and not has_parts: + continue + if self._task is VocTask.action_classification and not actions: + continue - item_annotations.append(Bbox(*obj_bbox, label=obj_label_id, - attributes=attributes, id=obj_id, group=group)) + item_annotations.append(Bbox(*obj_bbox, label=obj_label_id, + attributes=attributes, id=obj_id, group=group)) return item_annotations @staticmethod def _parse_bbox(object_elem): bbox_elem = object_elem.find('bndbox') - if bbox_elem is None: - return None - xmin = float(bbox_elem.find('xmin').text) xmax = float(bbox_elem.find('xmax').text) ymin = float(bbox_elem.find('ymin').text) ymax = float(bbox_elem.find('ymax').text) return [xmin, ymin, xmax - xmin, ymax - ymin] -class VocClassificationExtractor(VocExtractor): - def __init__(self, path): - super().__init__(path, task=VocTask.classification) - - subsets_dir = osp.join(path, VocPath.SUBSETS_DIR, 'Main') - subsets = self._load_subsets(subsets_dir) - self._subsets = subsets - - self._load_cls_annotations(subsets_dir, subsets) - -class VocDetectionExtractor(VocExtractor): +class VocDetectionExtractor(_VocXmlExtractor): def __init__(self, path): super().__init__(path, task=VocTask.detection) - subsets_dir = osp.join(path, VocPath.SUBSETS_DIR, 'Main') - subsets = self._load_subsets(subsets_dir) - self._subsets = subsets - - self._load_det_annotations() - -class VocSegmentationExtractor(VocExtractor): - def __init__(self, path): - super().__init__(path, task=VocTask.segmentation) - - subsets_dir = osp.join(path, VocPath.SUBSETS_DIR, 'Segmentation') - subsets = self._load_subsets(subsets_dir) - self._subsets = subsets - -class VocLayoutExtractor(VocExtractor): +class VocLayoutExtractor(_VocXmlExtractor): def __init__(self, path): super().__init__(path, task=VocTask.person_layout) - subsets_dir = osp.join(path, VocPath.SUBSETS_DIR, 'Layout') - subsets = self._load_subsets(subsets_dir) - self._subsets = subsets - - self._load_det_annotations() - -class VocActionExtractor(VocExtractor): +class VocActionExtractor(_VocXmlExtractor): def __init__(self, path): super().__init__(path, task=VocTask.action_classification) - subsets_dir = osp.join(path, VocPath.SUBSETS_DIR, 'Action') - subsets = self._load_subsets(subsets_dir) - self._subsets = subsets - - self._load_det_annotations() - - -class VocResultsExtractor(Extractor): - class Subset(Extractor): - def __init__(self, name, parent): - super().__init__() - self._parent = parent - self._name = name - self.items = [] - - def __iter__(self): - for item in self.items: - yield self._parent._get(item, self._name) - - def __len__(self): - return len(self.items) - - def categories(self): - return self._parent.categories() - - _SUPPORTED_TASKS = { - VocTask.classification: { - 'dir': 'Main', - 'mark': 'cls', - 'ext': '.txt', - 'path' : ['%(comp)s_cls_%(subset)s_%(label)s.txt'], - 'comp': ['comp1', 'comp2'], - }, - VocTask.detection: { - 'dir': 'Main', - 'mark': 'det', - 'ext': '.txt', - 'path': ['%(comp)s_det_%(subset)s_%(label)s.txt'], - 'comp': ['comp3', 'comp4'], - }, - VocTask.segmentation: { - 'dir': 'Segmentation', - 'mark': ['cls', 'inst'], - 'ext': '.png', - 'path': ['%(comp)s_%(subset)s_cls', '%(item)s.png'], - 'comp': ['comp5', 'comp6'], - }, - VocTask.person_layout: { - 'dir': 'Layout', - 'mark': 'layout', - 'ext': '.xml', - 'path': ['%(comp)s_layout_%(subset)s.xml'], - 'comp': ['comp7', 'comp8'], - }, - VocTask.action_classification: { - 'dir': 'Action', - 'mark': 'action', - 'ext': '.txt', - 'path': ['%(comp)s_action_%(subset)s_%(label)s.txt'], - 'comp': ['comp9', 'comp10'], - }, - } - - def _parse_txt_ann(self, path, subsets, annotations, task): - task_desc = self._SUPPORTED_TASKS[task] - task_dir = osp.join(path, task_desc['dir']) - ann_ext = task_desc['ext'] - if not osp.isdir(task_dir): - return - - ann_files = dir_items(task_dir, ann_ext, truncate_ext=True) - - for ann_file in ann_files: - ann_parts = filter(None, ann_file.strip().split('_')) - if len(ann_parts) != 4: - continue - _, mark, subset_name, label = ann_parts - if mark != task_desc['mark']: - continue - - label_id = self._get_label_id(label) - anns = defaultdict(list) - with open(osp.join(task_dir, ann_file + ann_ext), 'r') as f: - for line in f: - line_parts = line.split() - item = line_parts[0] - anns[item].append((label_id, *line_parts[1:])) - - subset = VocResultsExtractor.Subset(subset_name, self) - subset.items = list(anns) - - subsets[subset_name] = subset - annotations[subset_name] = dict(anns) - - def _parse_classification(self, path, subsets, annotations): - self._parse_txt_ann(path, subsets, annotations, - VocTask.classification) - - def _parse_detection(self, path, subsets, annotations): - self._parse_txt_ann(path, subsets, annotations, - VocTask.detection) - - def _parse_action(self, path, subsets, annotations): - self._parse_txt_ann(path, subsets, annotations, - VocTask.action_classification) - - def _load_categories(self): - label_map = None - label_map_path = osp.join(self._path, VocPath.LABELMAP_FILE) - if osp.isfile(label_map_path): - label_map = parse_label_map(label_map_path) - self._categories = make_voc_categories(label_map) - - def _get_label_id(self, label): - label_id = self._categories[AnnotationType.label].find(label) - assert label_id is not None - return label_id - - def __init__(self, path): - super().__init__() - - self._path = path - self._subsets = {} - self._annotations = {} - - self._load_categories() - - def __len__(self): - length = 0 - for subset in self._subsets.values(): - length += len(subset) - return length - - def subsets(self): - return list(self._subsets) - - def get_subset(self, name): - return self._subsets[name] - - def categories(self): - return self._categories - +class VocSegmentationExtractor(_VocExtractor): def __iter__(self): - for subset in self._subsets.values(): - for item in subset: - yield item - - def _get(self, item, subset_name): - image = None - image_path = osp.join(self._path, VocPath.IMAGES_DIR, - item + VocPath.IMAGE_EXT) - if osp.isfile(image_path): - image = lazy_image(image_path) - - annotations = self._get_annotations(item, subset_name) - - return DatasetItem(annotations=annotations, - id=item, subset=subset_name, image=image) - - def _get_annotations(self, item, subset_name): - raise NotImplementedError() - -class VocComp_1_2_Extractor(VocResultsExtractor): - def __init__(self, path): - super().__init__(path) - - subsets = {} - annotations = defaultdict(dict) - - self._parse_classification(path, subsets, annotations) - - self._subsets = subsets - self._annotations = dict(annotations) - - def _get_annotations(self, item, subset_name): - annotations = [] - - cls_ann = self._annotations[subset_name].get(item) - if cls_ann is not None: - for desc in cls_ann: - label_id, conf = desc - annotations.append(Label( - int(label_id), - attributes={ 'score': float(conf) } - )) - - return annotations - -class VocComp_3_4_Extractor(VocResultsExtractor): - def __init__(self, path): - super().__init__(path) - - subsets = {} - annotations = defaultdict(dict) - - self._parse_detection(path, subsets, annotations) - - self._subsets = subsets - self._annotations = dict(annotations) - - def _get_annotations(self, item, subset_name): - annotations = [] - - det_ann = self._annotations[subset_name].get(item) - if det_ann is not None: - for desc in det_ann: - label_id, conf, left, top, right, bottom = desc - annotations.append(Bbox( - x=float(left), y=float(top), - w=float(right) - float(left), h=float(bottom) - float(top), - label=int(label_id), - attributes={ 'score': float(conf) } - )) - - return annotations - -class VocComp_5_6_Extractor(VocResultsExtractor): - def __init__(self, path): - super().__init__(path) - - subsets = {} - annotations = defaultdict(dict) - - task_dir = osp.join(path, 'Segmentation') - if not osp.isdir(task_dir): - return - - ann_files = os.listdir(task_dir) - - for ann_dir in ann_files: - ann_parts = filter(None, ann_dir.strip().split('_')) - if len(ann_parts) != 4: - continue - _, subset_name, mark = ann_parts - if mark not in ['cls', 'inst']: - continue - - item_dir = osp.join(task_dir, ann_dir) - items = dir_items(item_dir, '.png', truncate_ext=True) - items = { name: osp.join(item_dir, item + '.png') \ - for name, item in items } - - subset = VocResultsExtractor.Subset(subset_name, self) - subset.items = list(items) - - subsets[subset_name] = subset - annotations[subset_name][mark] = items + for item_id in self._items: + image = osp.join(self._dataset_dir, VocPath.IMAGES_DIR, + item_id + VocPath.IMAGE_EXT) + anns = self._load_annotations(item_id) + yield DatasetItem(id=item_id, subset=self._subset, + image=image, annotations=anns) - self._subsets = subsets - self._annotations = dict(annotations) + @staticmethod + def _lazy_extract_mask(mask, c): + return lambda: mask == c - def _get_annotations(self, item, subset_name): - annotations = [] + def _load_annotations(self, item_id): + item_annotations = [] - segm_ann = self._annotations[subset_name] - cls_image_path = segm_ann.get(item) - if cls_image_path and osp.isfile(cls_image_path): + class_mask = None + segm_path = osp.join(self._dataset_dir, VocPath.SEGMENTATION_DIR, + item_id + VocPath.SEGM_EXT) + if osp.isfile(segm_path): inverse_cls_colormap = \ self._categories[AnnotationType.mask].inverse_colormap - annotations.append(Mask( - image=lazy_mask(cls_image_path, inverse_cls_colormap), - attributes={ 'class': True } - )) - - inst_ann = self._annotations[subset_name] - inst_image_path = inst_ann.get(item) - if inst_image_path and osp.isfile(inst_image_path): - annotations.append(Mask( - image=lazy_mask(inst_image_path, _inverse_inst_colormap), - attributes={ 'instances': True } - )) - - return annotations - -class VocComp_7_8_Extractor(VocResultsExtractor): - def __init__(self, path): - super().__init__(path) + class_mask = lazy_mask(segm_path, inverse_cls_colormap) - subsets = {} - annotations = defaultdict(dict) + instances_mask = None + inst_path = osp.join(self._dataset_dir, VocPath.INSTANCES_DIR, + item_id + VocPath.SEGM_EXT) + if osp.isfile(inst_path): + instances_mask = lazy_mask(inst_path, _inverse_inst_colormap) - task = VocTask.person_layout - task_desc = self._SUPPORTED_TASKS[task] - task_dir = osp.join(path, task_desc['dir']) - if not osp.isdir(task_dir): - return + if instances_mask is not None: + compiled_mask = CompiledMask(class_mask, instances_mask) - ann_ext = task_desc['ext'] - ann_files = dir_items(task_dir, ann_ext, truncate_ext=True) + if class_mask is not None: + label_cat = self._categories[AnnotationType.label] + instance_labels = compiled_mask.get_instance_labels( + class_count=len(label_cat.items)) + else: + instance_labels = {i: None + for i in range(compiled_mask.instance_count)} - for ann_file in ann_files: - ann_parts = filter(None, ann_file.strip().split('_')) - if len(ann_parts) != 4: - continue - _, mark, subset_name, _ = ann_parts - if mark != task_desc['mark']: - continue + for instance_id, label_id in instance_labels.items(): + image = compiled_mask.lazy_extract(instance_id) - layouts = {} - root = ET.parse(osp.join(task_dir, ann_file + ann_ext)) - root_elem = root.getroot() - for layout_elem in root_elem.findall('layout'): - item = layout_elem.find('image').text - obj_id = int(layout_elem.find('object').text) - conf = float(layout_elem.find('confidence').text) - parts = [] - for part_elem in layout_elem.findall('part'): - label_id = self._get_label_id(part_elem.find('class').text) - bbox_elem = part_elem.find('bndbox') - xmin = float(bbox_elem.find('xmin').text) - xmax = float(bbox_elem.find('xmax').text) - ymin = float(bbox_elem.find('ymin').text) - ymax = float(bbox_elem.find('ymax').text) - bbox = [xmin, ymin, xmax - xmin, ymax - ymin] - parts.append((label_id, bbox)) - layouts[item] = [obj_id, conf, parts] - - subset = VocResultsExtractor.Subset(subset_name, self) - subset.items = list(layouts) - - subsets[subset_name] = subset - annotations[subset_name] = layouts - - self._subsets = subsets - self._annotations = dict(annotations) - - def _get_annotations(self, item, subset_name): - annotations = [] - - layout_ann = self._annotations[subset_name].get(item) - if layout_ann is not None: - for desc in layout_ann: - obj_id, conf, parts = desc - attributes = { - 'score': conf, - 'object_id': obj_id, - } - - for part in parts: - label_id, bbox = part - annotations.append(Bbox( - *bbox, label=label_id, - attributes=attributes)) - - return annotations - -class VocComp_9_10_Extractor(VocResultsExtractor): - def __init__(self, path): - super().__init__(path) - - subsets = {} - annotations = defaultdict(dict) - - self._parse_action(path, subsets, annotations) - - self._subsets = subsets - self._annotations = dict(annotations) - - def _load_categories(self): - from collections import OrderedDict - from .format import VocAction - label_map = OrderedDict((a.name, [[], [], []]) for a in VocAction) - self._categories = make_voc_categories(label_map) - - def _get_annotations(self, item, subset_name): - annotations = [] - - action_ann = self._annotations[subset_name].get(item) - if action_ann is not None: - for desc in action_ann: - action_id, obj_id, conf = desc - annotations.append(Label( - action_id, - attributes={ - 'score': conf, - 'object_id': int(obj_id), + attributes = {} + if label_id is not None: + actions = {a: False + for a in label_cat.items[label_id].attributes } + attributes.update(actions) + + item_annotations.append(Mask( + image=image, label=label_id, + attributes=attributes, group=instance_id )) + elif class_mask is not None: + log.warn("item '%s': has only class segmentation, " + "instance masks will not be available" % item_id) + class_mask = class_mask() + classes = np.unique(class_mask) + for label_id in classes: + image = self._lazy_extract_mask(class_mask, label_id) + item_annotations.append(Mask(image=image, label=label_id)) - return annotations \ No newline at end of file + return item_annotations diff --git a/datumaro/datumaro/plugins/voc_format/importer.py b/datumaro/datumaro/plugins/voc_format/importer.py index f3d7c5ef3b97..78dc6cc9edea 100644 --- a/datumaro/datumaro/plugins/voc_format/importer.py +++ b/datumaro/datumaro/plugins/voc_format/importer.py @@ -3,11 +3,10 @@ # # SPDX-License-Identifier: MIT -import os +from glob import glob import os.path as osp from datumaro.components.extractor import Importer -from datumaro.util import find from .format import VocTask, VocPath @@ -21,61 +20,37 @@ class VocImporter(Importer): (VocTask.action_classification, 'voc_action', 'Action'), ] + @classmethod + def detect(cls, path): + return len(cls.find_subsets(path)) != 0 + def __call__(self, path, **extra_params): from datumaro.components.project import Project # cyclic import project = Project() - for task, extractor_type, task_dir in self._TASKS: - task_dir = osp.join(path, VocPath.SUBSETS_DIR, task_dir) - if not osp.isdir(task_dir): - continue + subset_paths = self.find_subsets(path) + if len(subset_paths) == 0: + raise Exception("Failed to find 'voc' dataset at '%s'" % path) - project.add_source(task.name, { - 'url': path, + for task, extractor_type, subset_path in subset_paths: + project.add_source('%s-%s' % + (task.name, osp.splitext(osp.basename(subset_path))[0]), + { + 'url': subset_path, 'format': extractor_type, 'options': dict(extra_params), }) - if len(project.config.sources) == 0: - raise Exception("Failed to find 'voc' dataset at '%s'" % path) - return project - -class VocResultsImporter: - _TASKS = [ - ('comp1', 'voc_comp_1_2', 'Main'), - ('comp2', 'voc_comp_1_2', 'Main'), - ('comp3', 'voc_comp_3_4', 'Main'), - ('comp4', 'voc_comp_3_4', 'Main'), - ('comp5', 'voc_comp_5_6', 'Segmentation'), - ('comp6', 'voc_comp_5_6', 'Segmentation'), - ('comp7', 'voc_comp_7_8', 'Layout'), - ('comp8', 'voc_comp_7_8', 'Layout'), - ('comp9', 'voc_comp_9_10', 'Action'), - ('comp10', 'voc_comp_9_10', 'Action'), - ] - - def __call__(self, path, **extra_params): - from datumaro.components.project import Project # cyclic import - project = Project() - - for task_name, extractor_type, task_dir in self._TASKS: - task_dir = osp.join(path, task_dir) + @staticmethod + def find_subsets(path): + subset_paths = [] + for task, extractor_type, task_dir in __class__._TASKS: + task_dir = osp.join(path, VocPath.SUBSETS_DIR, task_dir) if not osp.isdir(task_dir): continue - dir_items = os.listdir(task_dir) - if not find(dir_items, lambda x: x == task_name): - continue - - project.add_source(task_name, { - 'url': task_dir, - 'format': extractor_type, - 'options': dict(extra_params), - }) - - if len(project.config.sources) == 0: - raise Exception("Failed to find 'voc_results' dataset at '%s'" % \ - path) - - return project \ No newline at end of file + task_subsets = [p for p in glob(osp.join(task_dir, '*.txt')) + if '_' not in osp.basename(p)] + subset_paths += [(task, extractor_type, p) for p in task_subsets] + return subset_paths diff --git a/datumaro/datumaro/plugins/yolo_format/importer.py b/datumaro/datumaro/plugins/yolo_format/importer.py index fcee669dc6b9..4e14d0315aea 100644 --- a/datumaro/datumaro/plugins/yolo_format/importer.py +++ b/datumaro/datumaro/plugins/yolo_format/importer.py @@ -11,16 +11,16 @@ class YoloImporter(Importer): + @classmethod + def detect(cls, path): + return len(cls.find_configs(path)) != 0 + def __call__(self, path, **extra_params): from datumaro.components.project import Project # cyclic import project = Project() - if path.endswith('.data') and osp.isfile(path): - config_paths = [path] - else: - config_paths = glob(osp.join(path, '*.data')) - - if not osp.exists(path) or not config_paths: + config_paths = self.find_configs(path) + if len(config_paths) == 0: raise Exception("Failed to find 'yolo' dataset at '%s'" % path) for config_path in config_paths: @@ -35,4 +35,12 @@ def __call__(self, path, **extra_params): 'options': dict(extra_params), }) - return project \ No newline at end of file + return project + + @staticmethod + def find_configs(path): + if path.endswith('.data') and osp.isfile(path): + config_paths = [path] + else: + config_paths = glob(osp.join(path, '*.data')) + return config_paths \ No newline at end of file diff --git a/datumaro/datumaro/util/log_utils.py b/datumaro/datumaro/util/log_utils.py new file mode 100644 index 000000000000..6c8d8421e7e9 --- /dev/null +++ b/datumaro/datumaro/util/log_utils.py @@ -0,0 +1,16 @@ + +# Copyright (C) 2020 Intel Corporation +# +# SPDX-License-Identifier: MIT + +from contextlib import contextmanager +import logging + +@contextmanager +def logging_disabled(max_level=logging.CRITICAL): + previous_level = logging.root.manager.disable + logging.disable(max_level) + try: + yield + finally: + logging.disable(previous_level) \ No newline at end of file diff --git a/datumaro/tests/test_coco_format.py b/datumaro/tests/test_coco_format.py index 2caa03a7c09c..724fdc5a4afd 100644 --- a/datumaro/tests/test_coco_format.py +++ b/datumaro/tests/test_coco_format.py @@ -136,6 +136,12 @@ def categories(self): compare_datasets(self, DstExtractor(), dataset) + def test_can_detect(self): + with TestDir() as test_dir: + self.COCO_dataset_generate(test_dir) + + self.assertTrue(CocoImporter.detect(test_dir)) + class CocoConverterTest(TestCase): def _test_save_and_load(self, source_dataset, converter, test_dir, target_dataset=None, importer_args=None): diff --git a/datumaro/tests/test_cvat_format.py b/datumaro/tests/test_cvat_format.py index cc45bee921c1..4c9545d6f838 100644 --- a/datumaro/tests/test_cvat_format.py +++ b/datumaro/tests/test_cvat_format.py @@ -16,88 +16,94 @@ from datumaro.util.test_utils import TestDir, compare_datasets -class CvatExtractorTest(TestCase): - @staticmethod - def generate_dummy_cvat(path): - images_dir = osp.join(path, CvatPath.IMAGES_DIR) - anno_dir = osp.join(path, CvatPath.ANNOTATIONS_DIR) - - os.makedirs(images_dir) - os.makedirs(anno_dir) - - root_elem = ET.Element('annotations') - ET.SubElement(root_elem, 'version').text = '1.1' - - meta_elem = ET.SubElement(root_elem, 'meta') - task_elem = ET.SubElement(meta_elem, 'task') - ET.SubElement(task_elem, 'z_order').text = 'True' - ET.SubElement(task_elem, 'mode').text = 'interpolation' - - labels_elem = ET.SubElement(task_elem, 'labels') - - label1_elem = ET.SubElement(labels_elem, 'label') - ET.SubElement(label1_elem, 'name').text = 'label1' - label1_attrs_elem = ET.SubElement(label1_elem, 'attributes') - - label1_a1_elem = ET.SubElement(label1_attrs_elem, 'attribute') - ET.SubElement(label1_a1_elem, 'name').text = 'a1' - ET.SubElement(label1_a1_elem, 'input_type').text = 'checkbox' - ET.SubElement(label1_a1_elem, 'default_value').text = 'false' - ET.SubElement(label1_a1_elem, 'values').text = 'false\ntrue' - - label1_a2_elem = ET.SubElement(label1_attrs_elem, 'attribute') - ET.SubElement(label1_a2_elem, 'name').text = 'a2' - ET.SubElement(label1_a2_elem, 'input_type').text = 'radio' - ET.SubElement(label1_a2_elem, 'default_value').text = 'v1' - ET.SubElement(label1_a2_elem, 'values').text = 'v1\nv2\nv3' - - label2_elem = ET.SubElement(labels_elem, 'label') - ET.SubElement(label2_elem, 'name').text = 'label2' - - # item 1 - save_image(osp.join(images_dir, 'img0.jpg'), np.ones((8, 8, 3))) - item1_elem = ET.SubElement(root_elem, 'image') - item1_elem.attrib.update({ - 'id': '0', 'name': 'img0', 'width': '8', 'height': '8' - }) - - item1_ann1_elem = ET.SubElement(item1_elem, 'box') - item1_ann1_elem.attrib.update({ - 'label': 'label1', 'occluded': '1', 'z_order': '1', - 'xtl': '0', 'ytl': '2', 'xbr': '4', 'ybr': '4' - }) - item1_ann1_a1_elem = ET.SubElement(item1_ann1_elem, 'attribute') - item1_ann1_a1_elem.attrib['name'] = 'a1' - item1_ann1_a1_elem.text = 'true' - item1_ann1_a2_elem = ET.SubElement(item1_ann1_elem, 'attribute') - item1_ann1_a2_elem.attrib['name'] = 'a2' - item1_ann1_a2_elem.text = 'v3' - - item1_ann2_elem = ET.SubElement(item1_elem, 'polyline') - item1_ann2_elem.attrib.update({ - 'label': '', 'points': '1.0,2;3,4;5,6;7,8' - }) - - # item 2 - save_image(osp.join(images_dir, 'img1.jpg'), np.ones((10, 10, 3))) - item2_elem = ET.SubElement(root_elem, 'image') - item2_elem.attrib.update({ - 'id': '1', 'name': 'img1', 'width': '10', 'height': '10' - }) - - item2_ann1_elem = ET.SubElement(item2_elem, 'polygon') - item2_ann1_elem.attrib.update({ - 'label': '', 'points': '1,2;3,4;6,5', 'z_order': '1', - }) - - item2_ann2_elem = ET.SubElement(item2_elem, 'points') - item2_ann2_elem.attrib.update({ - 'label': 'label2', 'points': '1,2;3,4;5,6', 'z_order': '2', - }) - - with open(osp.join(anno_dir, 'train.xml'), 'w') as f: - f.write(ET.tostring(root_elem, encoding='unicode')) +def generate_dummy_cvat(path): + images_dir = osp.join(path, CvatPath.IMAGES_DIR) + anno_dir = osp.join(path, CvatPath.ANNOTATIONS_DIR) + + os.makedirs(images_dir) + os.makedirs(anno_dir) + + root_elem = ET.Element('annotations') + ET.SubElement(root_elem, 'version').text = '1.1' + + meta_elem = ET.SubElement(root_elem, 'meta') + task_elem = ET.SubElement(meta_elem, 'task') + ET.SubElement(task_elem, 'z_order').text = 'True' + ET.SubElement(task_elem, 'mode').text = 'interpolation' + + labels_elem = ET.SubElement(task_elem, 'labels') + + label1_elem = ET.SubElement(labels_elem, 'label') + ET.SubElement(label1_elem, 'name').text = 'label1' + label1_attrs_elem = ET.SubElement(label1_elem, 'attributes') + + label1_a1_elem = ET.SubElement(label1_attrs_elem, 'attribute') + ET.SubElement(label1_a1_elem, 'name').text = 'a1' + ET.SubElement(label1_a1_elem, 'input_type').text = 'checkbox' + ET.SubElement(label1_a1_elem, 'default_value').text = 'false' + ET.SubElement(label1_a1_elem, 'values').text = 'false\ntrue' + + label1_a2_elem = ET.SubElement(label1_attrs_elem, 'attribute') + ET.SubElement(label1_a2_elem, 'name').text = 'a2' + ET.SubElement(label1_a2_elem, 'input_type').text = 'radio' + ET.SubElement(label1_a2_elem, 'default_value').text = 'v1' + ET.SubElement(label1_a2_elem, 'values').text = 'v1\nv2\nv3' + + label2_elem = ET.SubElement(labels_elem, 'label') + ET.SubElement(label2_elem, 'name').text = 'label2' + + # item 1 + save_image(osp.join(images_dir, 'img0.jpg'), np.ones((8, 8, 3))) + item1_elem = ET.SubElement(root_elem, 'image') + item1_elem.attrib.update({ + 'id': '0', 'name': 'img0', 'width': '8', 'height': '8' + }) + + item1_ann1_elem = ET.SubElement(item1_elem, 'box') + item1_ann1_elem.attrib.update({ + 'label': 'label1', 'occluded': '1', 'z_order': '1', + 'xtl': '0', 'ytl': '2', 'xbr': '4', 'ybr': '4' + }) + item1_ann1_a1_elem = ET.SubElement(item1_ann1_elem, 'attribute') + item1_ann1_a1_elem.attrib['name'] = 'a1' + item1_ann1_a1_elem.text = 'true' + item1_ann1_a2_elem = ET.SubElement(item1_ann1_elem, 'attribute') + item1_ann1_a2_elem.attrib['name'] = 'a2' + item1_ann1_a2_elem.text = 'v3' + + item1_ann2_elem = ET.SubElement(item1_elem, 'polyline') + item1_ann2_elem.attrib.update({ + 'label': '', 'points': '1.0,2;3,4;5,6;7,8' + }) + + # item 2 + save_image(osp.join(images_dir, 'img1.jpg'), np.ones((10, 10, 3))) + item2_elem = ET.SubElement(root_elem, 'image') + item2_elem.attrib.update({ + 'id': '1', 'name': 'img1', 'width': '10', 'height': '10' + }) + + item2_ann1_elem = ET.SubElement(item2_elem, 'polygon') + item2_ann1_elem.attrib.update({ + 'label': '', 'points': '1,2;3,4;6,5', 'z_order': '1', + }) + + item2_ann2_elem = ET.SubElement(item2_elem, 'points') + item2_ann2_elem.attrib.update({ + 'label': 'label2', 'points': '1,2;3,4;5,6', 'z_order': '2', + }) + + with open(osp.join(anno_dir, 'train.xml'), 'w') as f: + f.write(ET.tostring(root_elem, encoding='unicode')) + +class CvatImporterTest(TestCase): + def test_can_detect(self): + with TestDir() as test_dir: + generate_dummy_cvat(test_dir) + self.assertTrue(CvatImporter.detect(test_dir)) + +class CvatExtractorTest(TestCase): def test_can_load(self): class TestExtractor(Extractor): def __iter__(self): @@ -130,7 +136,7 @@ def categories(self): } with TestDir() as test_dir: - self.generate_dummy_cvat(test_dir) + generate_dummy_cvat(test_dir) source_dataset = TestExtractor() parsed_dataset = CvatImporter()(test_dir).make_dataset() diff --git a/datumaro/tests/test_datumaro_format.py b/datumaro/tests/test_datumaro_format.py index e17da2a7b95a..4b71ddaed5fc 100644 --- a/datumaro/tests/test_datumaro_format.py +++ b/datumaro/tests/test_datumaro_format.py @@ -8,6 +8,7 @@ PolyLine, Bbox, Caption, LabelCategories, MaskCategories, PointsCategories ) +from datumaro.plugins.datumaro_format.importer import DatumaroImporter from datumaro.plugins.datumaro_format.converter import DatumaroConverter from datumaro.util.mask_tools import generate_colormap from datumaro.util.image import Image @@ -98,4 +99,10 @@ def test_can_save_and_load(self): self.assertEqual( source_dataset.categories(), - parsed_dataset.categories()) \ No newline at end of file + parsed_dataset.categories()) + + def test_can_detect(self): + with TestDir() as test_dir: + DatumaroConverter()(self.TestExtractor(), save_dir=test_dir) + + self.assertTrue(DatumaroImporter.detect(test_dir)) \ No newline at end of file diff --git a/datumaro/tests/test_tfrecord_format.py b/datumaro/tests/test_tfrecord_format.py index 0bd29ae41794..737ea6cf5f77 100644 --- a/datumaro/tests/test_tfrecord_format.py +++ b/datumaro/tests/test_tfrecord_format.py @@ -170,3 +170,32 @@ def test_labelmap_parsing(self): parsed = TfDetectionApiExtractor._parse_labelmap(text) self.assertEqual(expected, parsed) + +class TfrecordImporterTest(TestCase): + def test_can_detect(self): + class TestExtractor(Extractor): + def __iter__(self): + return iter([ + DatasetItem(id=1, subset='train', + image=np.ones((16, 16, 3)), + annotations=[ + Bbox(0, 4, 4, 8, label=2), + ] + ), + ]) + + def categories(self): + label_cat = LabelCategories() + for label in range(10): + label_cat.add('label_' + str(label)) + return { + AnnotationType.label: label_cat, + } + + def generate_dummy_tfrecord(path): + TfDetectionApiConverter()(TestExtractor(), save_dir=path) + + with TestDir() as test_dir: + generate_dummy_tfrecord(test_dir) + + self.assertTrue(TfDetectionApiImporter.detect(test_dir)) \ No newline at end of file diff --git a/datumaro/tests/test_voc_format.py b/datumaro/tests/test_voc_format.py index b91ee1a9325c..7abff4cc8515 100644 --- a/datumaro/tests/test_voc_format.py +++ b/datumaro/tests/test_voc_format.py @@ -189,14 +189,14 @@ def __iter__(self): for l in VOC.VocLabel if l.value % 2 == 1 ] ), - - DatasetItem(id='2007_000002', subset='test') ]) with TestDir() as test_dir: generate_dummy_voc(test_dir) - parsed_dataset = VocClassificationExtractor(test_dir) - compare_datasets(self, DstExtractor(), parsed_dataset) + + parsed_train = VocClassificationExtractor( + osp.join(test_dir, 'ImageSets', 'Main', 'train.txt')) + compare_datasets(self, DstExtractor(), parsed_train) def test_can_load_voc_det(self): class DstExtractor(TestExtractorBase): @@ -229,14 +229,13 @@ def __iter__(self): ), ] ), - - DatasetItem(id='2007_000002', subset='test') ]) with TestDir() as test_dir: generate_dummy_voc(test_dir) - parsed_dataset = VocDetectionExtractor(test_dir) - compare_datasets(self, DstExtractor(), parsed_dataset) + parsed_train = VocDetectionExtractor( + osp.join(test_dir, 'ImageSets', 'Main', 'train.txt')) + compare_datasets(self, DstExtractor(), parsed_train) def test_can_load_voc_segm(self): class DstExtractor(TestExtractorBase): @@ -250,14 +249,13 @@ def __iter__(self): ), ] ), - - DatasetItem(id='2007_000002', subset='test') ]) with TestDir() as test_dir: generate_dummy_voc(test_dir) - parsed_dataset = VocSegmentationExtractor(test_dir) - compare_datasets(self, DstExtractor(), parsed_dataset) + parsed_train = VocSegmentationExtractor( + osp.join(test_dir, 'ImageSets', 'Segmentation', 'train.txt')) + compare_datasets(self, DstExtractor(), parsed_train) def test_can_load_voc_layout(self): class DstExtractor(TestExtractorBase): @@ -285,14 +283,13 @@ def __iter__(self): ) ] ), - - DatasetItem(id='2007_000002', subset='test') ]) with TestDir() as test_dir: generate_dummy_voc(test_dir) - parsed_dataset = VocLayoutExtractor(test_dir) - compare_datasets(self, DstExtractor(), parsed_dataset) + parsed_train = VocLayoutExtractor( + osp.join(test_dir, 'ImageSets', 'Layout', 'train.txt')) + compare_datasets(self, DstExtractor(), parsed_train) def test_can_load_voc_action(self): class DstExtractor(TestExtractorBase): @@ -316,14 +313,13 @@ def __iter__(self): ), ] ), - - DatasetItem(id='2007_000002', subset='test') ]) with TestDir() as test_dir: generate_dummy_voc(test_dir) - parsed_dataset = VocActionExtractor(test_dir) - compare_datasets(self, DstExtractor(), parsed_dataset) + parsed_train = VocActionExtractor( + osp.join(test_dir, 'ImageSets', 'Action', 'train.txt')) + compare_datasets(self, DstExtractor(), parsed_train) class VocConverterTest(TestCase): def _test_save_and_load(self, source_dataset, converter, test_dir, @@ -757,16 +753,26 @@ def test_can_import(self): dataset = Project.import_from(test_dir, 'voc').make_dataset() - self.assertEqual(len(VOC.VocTask), len(dataset.sources)) + self.assertEqual(len(VOC.VocTask) * len(subsets), + len(dataset.sources)) self.assertEqual(set(subsets), set(dataset.subsets())) self.assertEqual( sum([len(s) for _, s in subsets.items()]), len(dataset)) + def test_can_detect_voc(self): + with TestDir() as test_dir: + generate_dummy_voc(test_dir) + + dataset_found = VocImporter.detect(test_dir) + + self.assertTrue(dataset_found) + class VocFormatTest(TestCase): 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']] with TestDir() as test_dir: file_path = osp.join(test_dir, 'test.txt') @@ -774,4 +780,4 @@ def test_can_write_and_parse_labelmap(self): VOC.write_label_map(file_path, src_label_map) dst_label_map = VOC.parse_label_map(file_path) - self.assertEqual(src_label_map, dst_label_map) \ No newline at end of file + self.assertEqual(src_label_map, dst_label_map) diff --git a/datumaro/tests/test_yolo_format.py b/datumaro/tests/test_yolo_format.py index e9a95108a9ee..4d29c349c24a 100644 --- a/datumaro/tests/test_yolo_format.py +++ b/datumaro/tests/test_yolo_format.py @@ -114,3 +114,29 @@ def categories(self): image_info={'1': (10, 15)}).make_dataset() compare_datasets(self, source_dataset, parsed_dataset) + +class YoloImporterTest(TestCase): + def test_can_detect(self): + class TestExtractor(Extractor): + def __iter__(self): + return iter([ + DatasetItem(id=1, subset='train', + image=Image(path='1.jpg', size=(10, 15)), + annotations=[ + Bbox(0, 2, 4, 2, label=2), + Bbox(3, 3, 2, 3, label=4), + ]), + ]) + + def categories(self): + label_categories = LabelCategories() + for i in range(10): + label_categories.add('label_' + str(i)) + return { + AnnotationType.label: label_categories, + } + + with TestDir() as test_dir: + YoloConverter()(TestExtractor(), save_dir=test_dir) + + self.assertTrue(YoloImporter.detect(test_dir)) \ No newline at end of file