diff --git a/datumaro/plugins/camvid_format.py b/datumaro/plugins/camvid_format.py index 8577a423ba..76de818196 100644 --- a/datumaro/plugins/camvid_format.py +++ b/datumaro/plugins/camvid_format.py @@ -135,11 +135,14 @@ def make_camvid_categories(label_map=None): class CamvidExtractor(SourceExtractor): - def __init__(self, path): + def __init__(self, path, subset=None): assert osp.isfile(path), path self._path = path self._dataset_dir = osp.dirname(path) - super().__init__(subset=osp.splitext(osp.basename(path))[0]) + + if not subset: + subset = osp.splitext(osp.basename(path))[0] + super().__init__(subset=subset) self._categories = self._load_categories(self._dataset_dir) self._items = list(self._load_items(path).values()) @@ -156,6 +159,11 @@ def _load_categories(self, path): def _load_items(self, path): items = {} + + labels = self._categories[AnnotationType.label]._indices + labels = { labels[label_name]: label_name + for label_name in labels } + with open(path, encoding='utf-8') as f: for line in f: line = line.strip() @@ -170,29 +178,27 @@ def _load_items(self, path): else: objects = line.split() image = objects[0] - item_id = ('/'.join(image.split('/')[2:]))[:-len(CamvidPath.IMAGE_EXT)] - image_path = osp.join(self._dataset_dir, - (image, image[1:])[image[0] == '/']) + item_id = osp.splitext(osp.join(*image.split('/')[2:]))[0] + image_path = osp.join(self._dataset_dir, image.lstrip('/')) + item_annotations = [] if 1 < len(objects): gt = objects[1] - gt_path = osp.join(self._dataset_dir, - (gt, gt[1:]) [gt[0] == '/']) - inverse_cls_colormap = \ - self._categories[AnnotationType.mask].inverse_colormap - mask = lazy_mask(gt_path, inverse_cls_colormap) - # loading mask through cache - mask = mask() + gt_path = osp.join(self._dataset_dir, gt.lstrip('/')) + mask = lazy_mask(gt_path, + self._categories[AnnotationType.mask].inverse_colormap) + mask = mask() # loading mask through cache + classes = np.unique(mask) - labels = self._categories[AnnotationType.label]._indices - labels = { labels[label_name]: label_name - for label_name in labels } for label_id in classes: if labels[label_id] in self._labels: image = self._lazy_extract_mask(mask, label_id) - item_annotations.append(Mask(image=image, label=label_id)) + item_annotations.append( + Mask(image=image, label=label_id)) + items[item_id] = DatasetItem(id=item_id, subset=self._subset, image=image_path, annotations=item_annotations) + return items @staticmethod diff --git a/datumaro/plugins/coco_format/extractor.py b/datumaro/plugins/coco_format/extractor.py index 4f94776b1d..29b97f7e27 100644 --- a/datumaro/plugins/coco_format/extractor.py +++ b/datumaro/plugins/coco_format/extractor.py @@ -21,11 +21,12 @@ class _CocoExtractor(SourceExtractor): - def __init__(self, path, task, merge_instance_polygons=False): + def __init__(self, path, task, merge_instance_polygons=False, subset=None): assert osp.isfile(path), path - subset = osp.splitext(osp.basename(path))[0].rsplit('_', maxsplit=1) - subset = subset[1] if len(subset) == 2 else None + if not subset: + subset = osp.splitext(osp.basename(path))[0].rsplit('_', maxsplit=1) + subset = subset[1] if len(subset) == 2 else None super().__init__(subset=subset) rootpath = '' diff --git a/datumaro/plugins/cvat_format/converter.py b/datumaro/plugins/cvat_format/converter.py index d327d7a20f..c0611e7dbb 100644 --- a/datumaro/plugins/cvat_format/converter.py +++ b/datumaro/plugins/cvat_format/converter.py @@ -11,7 +11,8 @@ from datumaro.components.converter import Converter from datumaro.components.dataset import ItemStatus -from datumaro.components.extractor import AnnotationType, DatasetItem +from datumaro.components.extractor import (AnnotationType, DatasetItem, + LabelCategories) from datumaro.util import cast, pairs from .format import CvatPath @@ -190,7 +191,8 @@ def _write_item(self, item, index): self._writer.close_image() def _write_meta(self): - label_cat = self._extractor.categories()[AnnotationType.label] + label_cat = self._extractor.categories().get( + AnnotationType.label, LabelCategories()) meta = OrderedDict([ ("task", OrderedDict([ ("id", ""), @@ -222,6 +224,8 @@ def _write_meta(self): self._writer.write_meta(meta) def _get_label(self, label_id): + if label_id is None: + return "" label_cat = self._extractor.categories()[AnnotationType.label] return label_cat.items[label_id] diff --git a/datumaro/plugins/cvat_format/extractor.py b/datumaro/plugins/cvat_format/extractor.py index 59a8f4ed94..466ab96a95 100644 --- a/datumaro/plugins/cvat_format/extractor.py +++ b/datumaro/plugins/cvat_format/extractor.py @@ -19,7 +19,7 @@ class CvatExtractor(SourceExtractor): _SUPPORTED_SHAPES = ('box', 'polygon', 'polyline', 'points') - def __init__(self, path): + def __init__(self, path, subset=None): assert osp.isfile(path), path rootpath = osp.dirname(path) images_dir = '' @@ -28,7 +28,9 @@ def __init__(self, path): self._images_dir = images_dir self._path = path - super().__init__(subset=osp.splitext(osp.basename(path))[0]) + if not subset: + subset = osp.splitext(osp.basename(path))[0] + super().__init__(subset=subset) items, categories = self._parse(path) self._items = list(self._load_items(items).values()) diff --git a/datumaro/plugins/icdar_format/converter.py b/datumaro/plugins/icdar_format/converter.py index 186d48d296..e7759b2eba 100644 --- a/datumaro/plugins/icdar_format/converter.py +++ b/datumaro/plugins/icdar_format/converter.py @@ -20,16 +20,18 @@ def apply(self): for subset_name, subset in self._extractor.subsets().items(): annotation = '' for item in subset: - if item.has_image and self._save_images: - self._save_image(item, osp.join(self._save_dir, subset_name, - IcdarPath.IMAGES_DIR, item.id + IcdarPath.IMAGE_EXT)) + image_filename = self._make_image_filename(item) + if self._save_images and item.has_image: + self._save_image(item, osp.join(self._save_dir, + subset_name, IcdarPath.IMAGES_DIR, image_filename)) - annotation += '%s, ' % (item.id + IcdarPath.IMAGE_EXT) + annotation += '%s, ' % image_filename for ann in item.annotations: if ann.type != AnnotationType.caption: continue annotation += '\"%s\"' % ann.caption annotation += '\n' + if len(annotation): anno_file = osp.join(self._save_dir, subset_name, 'gt.txt') os.makedirs(osp.dirname(anno_file), exist_ok=True) @@ -42,9 +44,9 @@ class IcdarTextLocalizationConverter(Converter): def apply(self): for subset_name, subset in self._extractor.subsets().items(): for item in subset: - if item.has_image and self._save_images: - self._save_image(item, osp.join(self._save_dir, subset_name, - IcdarPath.IMAGES_DIR, item.id + IcdarPath.IMAGE_EXT)) + if self._save_images and item.has_image: + self._save_image(item, + subdir=osp.join(subset_name, IcdarPath.IMAGES_DIR)) annotation = '' for ann in item.annotations: @@ -57,8 +59,9 @@ def apply(self): if ann.attributes and 'text' in ann.attributes: annotation += ',\"%s\"' % ann.attributes['text'] annotation += '\n' - anno_file = osp.join(self._save_dir, subset_name, osp.dirname(item.id), - 'gt_' + osp.basename(item.id) + '.txt') + + anno_file = osp.join(self._save_dir, subset_name, + osp.dirname(item.id), 'gt_' + osp.basename(item.id) + '.txt') os.makedirs(osp.dirname(anno_file), exist_ok=True) with open(anno_file, 'w', encoding='utf-8') as f: f.write(annotation) @@ -69,59 +72,57 @@ class IcdarTextSegmentationConverter(Converter): def apply(self): for subset_name, subset in self._extractor.subsets().items(): for item in subset: - if item.has_image and self._save_images: - self._save_image(item, osp.join(self._save_dir, subset_name, - IcdarPath.IMAGES_DIR, item.id + IcdarPath.IMAGE_EXT)) + self._save_item(subset_name, subset, item) - annotation = '' - colormap = [(255, 255, 255)] - anns = [a for a in item.annotations - if a.type == AnnotationType.mask] - if anns: - is_not_index = len([p for p in anns if 'index' not in p.attributes]) - if is_not_index: - raise Exception("Item %s: a mask must have" - "'index' attribute" % item.id) - anns = sorted(anns, key=lambda a: a.attributes['index']) - group = anns[0].group - for ann in anns: - if ann.group != group or (not ann.group and anns[0].group != 0): - annotation += '\n' - text = '' - if ann.attributes: - if 'text' in ann.attributes: - text = ann.attributes['text'] - if text == ' ': - annotation += '#' - if 'color' in ann.attributes and \ - len(ann.attributes['color'].split()) == 3: - color = ann.attributes['color'].split() - colormap.append( - (int(color[0]), int(color[1]), int(color[2]))) - annotation += ' '.join(p for p in color) - else: - raise Exception("Item %s: a mask must have " - "an RGB color attribute, e. g. '10 7 50'" % item.id) - if 'center' in ann.attributes: - annotation += ' %s' % ann.attributes['center'] - else: - annotation += ' - -' - bbox = ann.get_bbox() - annotation += ' %s %s %s %s' % (bbox[0], bbox[1], - bbox[0] + bbox[2], bbox[1] + bbox[3]) - annotation += ' \"%s\"' % text - annotation += '\n' - group = ann.group - - mask = CompiledMask.from_instance_masks(anns, - instance_labels=[m.attributes['index'] + 1 for m in anns]) - mask = paint_mask(mask.class_mask, - { i: colormap[i] for i in range(len(colormap)) }) - save_image(osp.join(self._save_dir, subset_name, - item.id + '_GT' + IcdarPath.GT_EXT), mask, create_dir=True) + def _save_item(self, subset_name, subset, item): + if self._save_images and item.has_image: + self._save_image(item, + subdir=osp.join(subset_name, IcdarPath.IMAGES_DIR)) - anno_file = osp.join(self._save_dir, subset_name, - item.id + '_GT' + '.txt') - os.makedirs(osp.dirname(anno_file), exist_ok=True) - with open(anno_file, 'w', encoding='utf-8') as f: - f.write(annotation) \ No newline at end of file + annotation = '' + colormap = [(255, 255, 255)] + anns = [a for a in item.annotations if a.type == AnnotationType.mask] + if anns: + anns = sorted(anns, key=lambda a: int(a.attributes.get('index', 0))) + group = anns[0].group + for ann in anns: + if ann.group != group or (not ann.group and anns[0].group != 0): + annotation += '\n' + text = '' + if ann.attributes: + if 'text' in ann.attributes: + text = ann.attributes['text'] + if text == ' ': + annotation += '#' + if 'color' in ann.attributes and \ + len(ann.attributes['color'].split()) == 3: + color = ann.attributes['color'].split() + colormap.append( + (int(color[0]), int(color[1]), int(color[2]))) + annotation += ' '.join(p for p in color) + else: + raise Exception("Item %s: a mask must have " + "an RGB color attribute, e.g. '10 7 50'" % item.id) + if 'center' in ann.attributes: + annotation += ' %s' % ann.attributes['center'] + else: + annotation += ' - -' + bbox = ann.get_bbox() + annotation += ' %s %s %s %s' % (bbox[0], bbox[1], + bbox[0] + bbox[2], bbox[1] + bbox[3]) + annotation += ' \"%s\"' % text + annotation += '\n' + group = ann.group + + mask = CompiledMask.from_instance_masks(anns, + instance_labels=[m.attributes['index'] + 1 for m in anns]) + mask = paint_mask(mask.class_mask, + { i: colormap[i] for i in range(len(colormap)) }) + save_image(osp.join(self._save_dir, subset_name, + item.id + '_GT' + IcdarPath.GT_EXT), mask, create_dir=True) + + anno_file = osp.join(self._save_dir, subset_name, + item.id + '_GT' + '.txt') + os.makedirs(osp.dirname(anno_file), exist_ok=True) + with open(anno_file, 'w', encoding='utf-8') as f: + f.write(annotation) \ No newline at end of file diff --git a/datumaro/plugins/icdar_format/extractor.py b/datumaro/plugins/icdar_format/extractor.py index 3b304a9670..e37df63f2e 100644 --- a/datumaro/plugins/icdar_format/extractor.py +++ b/datumaro/plugins/icdar_format/extractor.py @@ -2,20 +2,21 @@ # # SPDX-License-Identifier: MIT -from glob import glob +from glob import iglob import os.path as osp import numpy as np from datumaro.components.extractor import (Bbox, Caption, DatasetItem, Importer, Mask, MaskCategories, Polygon, SourceExtractor) +from datumaro.util.image import find_images from datumaro.util.mask_tools import lazy_mask from .format import IcdarPath, IcdarTask class _IcdarExtractor(SourceExtractor): - def __init__(self, path, task): + def __init__(self, path, task, subset=None): self._path = path self._task = task @@ -23,7 +24,11 @@ def __init__(self, path, task): if not osp.isfile(path): raise FileNotFoundError( "Can't read annotation file '%s'" % path) - super().__init__(subset=osp.basename(osp.dirname(path))) + + if not subset: + subset = osp.basename(osp.dirname(path)) + super().__init__(subset=subset) + self._dataset_dir = osp.dirname(osp.dirname(path)) self._items = list(self._load_recognition_items().values()) @@ -31,8 +36,13 @@ def __init__(self, path, task): if not osp.isdir(path): raise NotADirectoryError( "Can't open folder with annotation files '%s'" % path) - super().__init__(subset=osp.basename(path)) + + if not subset: + subset = osp.basename(path) + super().__init__(subset=subset) + self._dataset_dir = osp.dirname(path) + if task is IcdarTask.text_localization: self._items = list(self._load_localization_items().values()) else: @@ -40,6 +50,7 @@ def __init__(self, path, task): def _load_recognition_items(self): items = {} + with open(self._path, encoding='utf-8') as f: for line in f: line = line.strip() @@ -51,7 +62,7 @@ def _load_recognition_items(self): image = objects[0][:-1] captions = [] - item_id = image[:-len(IcdarPath.IMAGE_EXT)] + item_id = osp.splitext(image)[0] image_path = osp.join(osp.dirname(self._path), IcdarPath.IMAGES_DIR, image) if item_id not in items: @@ -63,21 +74,28 @@ def _load_recognition_items(self): if caption[0] == '\"' and caption[-1] == '\"': caption = caption[1:-1] annotations.append(Caption(caption)) + return items def _load_localization_items(self): items = {} - for path in glob(osp.join(self._path, '**', '*.txt'), recursive=True): + image_dir = osp.join(self._path, IcdarPath.IMAGES_DIR) + if osp.isdir(image_dir): + images = { osp.splitext(osp.relpath(p, image_dir))[0]: p + for p in find_images(image_dir, recursive=True) } + else: + images = {} + + for path in iglob(osp.join(self._path, '**', '*.txt'), recursive=True): item_id = osp.splitext(osp.relpath(path, self._path))[0] if osp.basename(item_id).startswith('gt_'): item_id = osp.join(osp.dirname(item_id), osp.basename(item_id)[3:]) item_id = item_id.replace('\\', '/') - image_path = osp.join(self._path, IcdarPath.IMAGES_DIR, - item_id + IcdarPath.IMAGE_EXT) + if item_id not in items: items[item_id] = DatasetItem(item_id, subset=self._subset, - image=image_path) + image=images.get(item_id)) annotations = items[item_id].annotations with open(path, encoding='utf-8') as f: @@ -89,12 +107,14 @@ def _load_localization_items(self): if 8 <= len(objects): points = [float(p) for p in objects[:8]] + attributes = {} if len(objects) == 9: text = objects[8] if text[0] == '\"' and text[-1] == '\"': text = text[1:-1] attributes['text'] = text + annotations.append( Polygon(points, attributes=attributes)) elif 4 <= len(objects): @@ -102,12 +122,14 @@ def _load_localization_items(self): y = float(objects[1]) w = float(objects[2]) - x h = float(objects[3]) - y + attributes = {} if len(objects) == 5: text = objects[4] if text[0] == '\"' and text[-1] == '\"': text = text[1:-1] attributes['text'] = text + annotations.append( Bbox(x, y, w, h, attributes=attributes)) return items @@ -115,16 +137,22 @@ def _load_localization_items(self): def _load_segmentation_items(self): items = {} - for path in glob(osp.join(self._path, '**', '*.txt'), recursive=True): + image_dir = osp.join(self._path, IcdarPath.IMAGES_DIR) + if osp.isdir(image_dir): + images = { osp.splitext(osp.relpath(p, image_dir))[0]: p + for p in find_images(image_dir, recursive=True) } + else: + images = {} + + for path in iglob(osp.join(self._path, '**', '*.txt'), recursive=True): item_id = osp.splitext(osp.relpath(path, self._path))[0] item_id = item_id.replace('\\', '/') if item_id.endswith('_GT'): item_id = item_id[:-3] - image_path = osp.join(self._path, IcdarPath.IMAGES_DIR, - item_id + IcdarPath.IMAGE_EXT) + if item_id not in items: items[item_id] = DatasetItem(item_id, subset=self._subset, - image=image_path) + image=images.get(item_id)) annotations = items[item_id].annotations colors = [(255, 255, 255)] @@ -180,7 +208,8 @@ def _load_segmentation_items(self): i = int(label_id) annotations.append(Mask(group=groups[i], image=self._lazy_extract_mask(mask, label_id), - attributes={ 'index': i - 1, 'color': ' '.join(str(p) for p in colors[i]), + attributes={ 'index': i - 1, + 'color': ' '.join(str(p) for p in colors[i]), 'text': chars[i], 'center': centers[i] } )) return items diff --git a/tests/test_camvid_format.py b/tests/test_camvid_format.py index dba9d0116d..9bf3b1b234 100644 --- a/tests/test_camvid_format.py +++ b/tests/test_camvid_format.py @@ -9,6 +9,7 @@ Extractor, LabelCategories, Mask) from datumaro.components.dataset import Dataset from datumaro.plugins.camvid_format import CamvidConverter, CamvidImporter +from datumaro.util.image import Image from datumaro.util.test_utils import (TestDir, compare_datasets, test_save_and_load) @@ -77,10 +78,10 @@ def test_can_detect_camvid(self): class CamvidConverterTest(TestCase): def _test_save_and_load(self, source_dataset, converter, test_dir, - target_dataset=None, importer_args=None): + target_dataset=None, importer_args=None, **kwargs): return test_save_and_load(self, source_dataset, converter, test_dir, importer='camvid', - target_dataset=target_dataset, importer_args=importer_args) + target_dataset=target_dataset, importer_args=importer_args, **kwargs) def test_can_save_camvid_segm(self): class TestExtractor(TestExtractorBase): @@ -243,3 +244,57 @@ def categories(self): self._test_save_and_load(SrcExtractor(), partial(CamvidConverter.convert, label_map='source'), test_dir, target_dataset=DstExtractor()) + + def test_can_save_and_load_image_with_arbitrary_extension(self): + class SrcExtractor(TestExtractorBase): + def __iter__(self): + return iter([ + DatasetItem(id='q/1', image=Image(path='q/1.JPEG', + data=np.zeros((4, 3, 3)))), + DatasetItem(id='a/b/c/2', image=Image( + path='a/b/c/2.bmp', data=np.ones((1, 5, 3)) + ), + annotations=[ + Mask(np.array([[0, 0, 0, 1, 0]]), + label=self._label('a')), + Mask(np.array([[0, 1, 1, 0, 0]]), + label=self._label('b')), + ]) + ]) + + def categories(self): + label_map = OrderedDict() + label_map['a'] = None + label_map['b'] = None + return Camvid.make_camvid_categories(label_map) + + class DstExtractor(TestExtractorBase): + def __iter__(self): + return iter([ + DatasetItem(id='q/1', image=Image(path='q/1.JPEG', + data=np.zeros((4, 3, 3)))), + DatasetItem(id='a/b/c/2', image=Image( + path='a/b/c/2.bmp', data=np.ones((1, 5, 3)) + ), + annotations=[ + Mask(np.array([[1, 0, 0, 0, 1]]), + label=self._label('background')), + Mask(np.array([[0, 0, 0, 1, 0]]), + label=self._label('a')), + Mask(np.array([[0, 1, 1, 0, 0]]), + label=self._label('b')), + ]) + ]) + + def categories(self): + label_map = OrderedDict() + label_map['background'] = None + label_map['a'] = None + label_map['b'] = None + return Camvid.make_camvid_categories(label_map) + + with TestDir() as test_dir: + self._test_save_and_load(SrcExtractor(), + partial(CamvidConverter.convert, save_images=True), + test_dir, require_images=True, + target_dataset=DstExtractor()) \ No newline at end of file diff --git a/tests/test_coco_format.py b/tests/test_coco_format.py index d57133d458..b884009b22 100644 --- a/tests/test_coco_format.py +++ b/tests/test_coco_format.py @@ -148,10 +148,10 @@ def test_can_detect(self): class CocoConverterTest(TestCase): def _test_save_and_load(self, source_dataset, converter, test_dir, - target_dataset=None, importer_args=None): + target_dataset=None, importer_args=None, **kwargs): return test_save_and_load(self, source_dataset, converter, test_dir, importer='coco', - target_dataset=target_dataset, importer_args=importer_args) + target_dataset=target_dataset, importer_args=importer_args, **kwargs) def test_can_save_and_load_captions(self): expected_dataset = Dataset.from_iterable([ @@ -554,7 +554,21 @@ def test_relative_paths(self): with TestDir() as test_dir: self._test_save_and_load(expected_dataset, - partial(CocoImageInfoConverter.convert, save_images=True), test_dir) + partial(CocoImageInfoConverter.convert, save_images=True), + test_dir, require_images=True) + + def test_can_save_and_load_image_with_arbitrary_extension(self): + expected = Dataset.from_iterable([ + DatasetItem(id='q/1', image=Image(path='q/1.JPEG', + data=np.zeros((4, 3, 3))), attributes={'id': 1}), + DatasetItem(id='a/b/c/2', image=Image(path='a/b/c/2.bmp', + data=np.zeros((3, 4, 3))), attributes={'id': 2}), + ]) + + with TestDir() as test_dir: + self._test_save_and_load(expected, + partial(CocoConverter.convert, save_images=True), + test_dir, require_images=True) def test_preserve_coco_ids(self): expected_dataset = Dataset.from_iterable([ @@ -564,7 +578,8 @@ def test_preserve_coco_ids(self): with TestDir() as test_dir: self._test_save_and_load(expected_dataset, - partial(CocoImageInfoConverter.convert, save_images=True), test_dir) + partial(CocoImageInfoConverter.convert, save_images=True), + test_dir, require_images=True) def test_annotation_attributes(self): expected_dataset = Dataset.from_iterable([ diff --git a/tests/test_cvat_format.py b/tests/test_cvat_format.py index c23c32a4b1..ced0930f04 100644 --- a/tests/test_cvat_format.py +++ b/tests/test_cvat_format.py @@ -142,10 +142,10 @@ def test_can_load_video(self): class CvatConverterTest(TestCase): def _test_save_and_load(self, source_dataset, converter, test_dir, - target_dataset=None, importer_args=None): + target_dataset=None, importer_args=None, **kwargs): return test_save_and_load(self, source_dataset, converter, test_dir, importer='cvat', - target_dataset=target_dataset, importer_args=importer_args) + target_dataset=target_dataset, importer_args=importer_args, **kwargs) def test_can_save_and_load(self): label_categories = LabelCategories() @@ -237,6 +237,27 @@ def test_can_save_and_load(self): partial(CvatConverter.convert, save_images=True), test_dir, target_dataset=target_dataset) + def test_relative_paths(self): + source_dataset = Dataset.from_iterable([ + DatasetItem(id='1', image=np.ones((4, 2, 3))), + DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3))), + DatasetItem(id='subdir2/1', image=np.ones((5, 4, 3))), + ]) + + target_dataset = Dataset.from_iterable([ + DatasetItem(id='1', image=np.ones((4, 2, 3)), + attributes={'frame': 0}), + DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3)), + attributes={'frame': 1}), + DatasetItem(id='subdir2/1', image=np.ones((5, 4, 3)), + attributes={'frame': 2}), + ], categories=[]) + + with TestDir() as test_dir: + self._test_save_and_load(source_dataset, + partial(CvatConverter.convert, save_images=True), test_dir, + target_dataset=target_dataset, require_images=True) + def test_can_save_dataset_with_cyrillic_and_spaces_in_filename(self): label_categories = LabelCategories() for i in range(10): @@ -269,38 +290,26 @@ def test_can_save_dataset_with_cyrillic_and_spaces_in_filename(self): with TestDir() as test_dir: self._test_save_and_load(source_dataset, partial(CvatConverter.convert, save_images=True), test_dir, - target_dataset=target_dataset) - - def test_relative_paths(self): - source_dataset = Dataset.from_iterable([ - DatasetItem(id='1', image=np.ones((4, 2, 3))), - DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3))), - DatasetItem(id='subdir2/1', image=np.ones((5, 4, 3))), - ], categories={ AnnotationType.label: LabelCategories() }) + target_dataset=target_dataset, require_images=True) - target_dataset = Dataset.from_iterable([ - DatasetItem(id='1', image=np.ones((4, 2, 3)), - attributes={'frame': 0}), - DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3)), - attributes={'frame': 1}), - DatasetItem(id='subdir2/1', image=np.ones((5, 4, 3)), - attributes={'frame': 2}), - ], categories={ - AnnotationType.label: LabelCategories() - }) + def test_can_save_and_load_image_with_arbitrary_extension(self): + expected = Dataset.from_iterable([ + DatasetItem('q/1', image=Image(path='q/1.JPEG', + data=np.zeros((4, 3, 3))), attributes={'frame': 1}), + DatasetItem('a/b/c/2', image=Image(path='a/b/c/2.bmp', + data=np.zeros((3, 4, 3))), attributes={'frame': 2}), + ], categories=[]) with TestDir() as test_dir: - self._test_save_and_load(source_dataset, - partial(CvatConverter.convert, save_images=True), test_dir, - target_dataset=target_dataset) + self._test_save_and_load(expected, + partial(CvatConverter.convert, save_images=True), + test_dir, require_images=True) def test_preserve_frame_ids(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='some/name1', image=np.ones((4, 2, 3)), attributes={'frame': 40}), - ], categories={ - AnnotationType.label: LabelCategories() - }) + ], categories=[]) with TestDir() as test_dir: self._test_save_and_load(expected_dataset, @@ -310,12 +319,12 @@ def test_reindex(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='some/name1', image=np.ones((4, 2, 3)), attributes={'frame': 40}), - ], categories={ AnnotationType.label: LabelCategories() }) + ]) expected_dataset = Dataset.from_iterable([ DatasetItem(id='some/name1', image=np.ones((4, 2, 3)), attributes={'frame': 0}), - ], categories={ AnnotationType.label: LabelCategories() }) + ], categories=[]) with TestDir() as test_dir: self._test_save_and_load(source_dataset, @@ -329,7 +338,7 @@ def test_inplace_save_writes_only_updated_data(self): DatasetItem(1, subset='a'), DatasetItem(2, subset='b'), DatasetItem(3, subset='c', image=np.ones((2, 2, 3))), - ], categories=[]) + ]) dataset.export(path, 'cvat', save_images=True) os.unlink(osp.join(path, 'a.xml')) os.unlink(osp.join(path, 'b.xml')) diff --git a/tests/test_datumaro_format.py b/tests/test_datumaro_format.py index 582bad7fdd..047950e6d3 100644 --- a/tests/test_datumaro_format.py +++ b/tests/test_datumaro_format.py @@ -109,7 +109,21 @@ def test_can_save_dataset_with_cyrillic_and_spaces_in_filename(self): with TestDir() as test_dir: self._test_save_and_load(test_dataset, - partial(DatumaroConverter.convert, save_images=True), test_dir) + partial(DatumaroConverter.convert, save_images=True), + test_dir) + + def test_can_save_and_load_image_with_arbitrary_extension(self): + expected = Dataset.from_iterable([ + DatasetItem(id='q/1', image=Image(path='q/1.JPEG', + data=np.zeros((4, 3, 3))), attributes={'frame': 1}), + DatasetItem(id='a/b/c/2', image=Image(path='a/b/c/2.bmp', + data=np.zeros((3, 4, 3))), attributes={'frame': 2}), + ]) + + with TestDir() as test_dir: + self._test_save_and_load(expected, + partial(DatumaroConverter.convert, save_images=True), + test_dir) def test_inplace_save_writes_only_updated_data(self): with TestDir() as path: diff --git a/tests/test_icdar_format.py b/tests/test_icdar_format.py index fa7150c2f5..7f0e99b0a7 100644 --- a/tests/test_icdar_format.py +++ b/tests/test_icdar_format.py @@ -3,13 +3,17 @@ from unittest import TestCase import numpy as np + from datumaro.components.extractor import (Bbox, Caption, DatasetItem, Mask, Polygon) from datumaro.components.project import Dataset from datumaro.plugins.icdar_format.converter import ( IcdarTextLocalizationConverter, IcdarTextSegmentationConverter, IcdarWordRecognitionConverter) -from datumaro.plugins.icdar_format.extractor import IcdarWordRecognitionImporter +from datumaro.plugins.icdar_format.extractor import ( + IcdarWordRecognitionImporter, IcdarTextLocalizationImporter, + IcdarTextSegmentationImporter) +from datumaro.util.image import Image from datumaro.util.test_utils import (TestDir, compare_datasets, test_save_and_load) @@ -17,10 +21,18 @@ DUMMY_DATASET_DIR = osp.join(osp.dirname(__file__), 'assets', 'icdar_dataset') class IcdarImporterTest(TestCase): - def test_can_detect(self): + def test_can_detect_word_recognition(self): self.assertTrue(IcdarWordRecognitionImporter.detect( osp.join(DUMMY_DATASET_DIR, 'word_recognition'))) + def test_can_detect_text_localization(self): + self.assertTrue(IcdarTextLocalizationImporter.detect( + osp.join(DUMMY_DATASET_DIR, 'text_localization'))) + + def test_can_detect_text_segmentation(self): + self.assertTrue(IcdarTextSegmentationImporter.detect( + osp.join(DUMMY_DATASET_DIR, 'text_segmentation'))) + def test_can_import_captions(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='word_1', subset='train', @@ -98,11 +110,11 @@ def test_can_import_masks(self): compare_datasets(self, expected_dataset, dataset) class IcdarConverterTest(TestCase): - def _test_save_and_load(self, source_dataset, converter, test_dir, - importer, target_dataset=None, importer_args=None, **kwargs): + def _test_save_and_load(self, source_dataset, converter, test_dir, importer, + target_dataset=None, importer_args=None, **kwargs): return test_save_and_load(self, source_dataset, converter, test_dir, - importer, target_dataset=target_dataset, importer_args=importer_args, - **kwargs) + importer, + target_dataset=target_dataset, importer_args=importer_args, **kwargs) def test_can_save_and_load_captions(self): expected_dataset = Dataset.from_iterable([ @@ -208,3 +220,21 @@ def test_can_save_dataset_with_cyrillic_and_spaces_in_filename(self): self._test_save_and_load(expected_dataset, partial(converter.convert, save_images=True), test_dir, importer, require_images=True) + + def test_can_save_and_load_image_with_arbitrary_extension(self): + expected = Dataset.from_iterable([ + DatasetItem(id='q/1', image=Image(path='q/1.JPEG', + data=np.zeros((4, 3, 3)))), + DatasetItem(id='a/b/c/2', image=Image(path='a/b/c/2.bmp', + data=np.zeros((3, 4, 3)))), + ]) + + for importer, converter in [ + ('icdar_word_recognition', IcdarWordRecognitionConverter), + ('icdar_text_localization', IcdarTextLocalizationConverter), + ('icdar_text_segmentation', IcdarTextSegmentationConverter), + ]: + with self.subTest(subformat=converter), TestDir() as test_dir: + self._test_save_and_load(expected, + partial(converter.convert, save_images=True), + test_dir, importer, require_images=True) \ No newline at end of file