diff --git a/datumaro/datumaro/plugins/yolo_format/converter.py b/datumaro/datumaro/plugins/yolo_format/converter.py index 59a46a6260c5..de30d7d785e9 100644 --- a/datumaro/datumaro/plugins/yolo_format/converter.py +++ b/datumaro/datumaro/plugins/yolo_format/converter.py @@ -75,21 +75,24 @@ def __call__(self, extractor, save_dir): image_paths = OrderedDict() for item in subset: - image_name = '%s.jpg' % item.id - image_paths[item.id] = osp.join('data', - osp.basename(subset_dir), image_name) + if not item.has_image: + raise Exception("Failed to export item '%s': " + "item has no image info" % item.id) + height, width = item.image.size + image_name = item.image.filename + item_name = osp.splitext(item.image.filename)[0] if self._save_images: if item.has_image and item.image.has_data: + if not item_name: + item_name = item.id + image_name = item_name + '.jpg' save_image(osp.join(subset_dir, image_name), item.image.data) else: log.warning("Item '%s' has no image" % item.id) - - if not item.has_image: - raise Exception("Failed to export item '%s': " - "item has no image info" % item.id) - height, width = item.image.size + image_paths[item.id] = osp.join('data', + osp.basename(subset_dir), image_name) yolo_annotation = '' for bbox in item.annotations: @@ -102,7 +105,7 @@ def __call__(self, extractor, save_dir): yolo_bb = ' '.join('%.6f' % p for p in yolo_bb) yolo_annotation += '%s %s\n' % (bbox.label, yolo_bb) - annotation_path = osp.join(subset_dir, '%s.txt' % item.id) + annotation_path = osp.join(subset_dir, '%s.txt' % item_name) with open(annotation_path, 'w') as f: f.write(yolo_annotation) diff --git a/datumaro/datumaro/plugins/yolo_format/extractor.py b/datumaro/datumaro/plugins/yolo_format/extractor.py index cb52327717c1..7840b26c5ca1 100644 --- a/datumaro/datumaro/plugins/yolo_format/extractor.py +++ b/datumaro/datumaro/plugins/yolo_format/extractor.py @@ -10,7 +10,7 @@ from datumaro.components.extractor import (SourceExtractor, Extractor, DatasetItem, AnnotationType, Bbox, LabelCategories ) -from datumaro.util.image import lazy_image +from datumaro.util.image import Image from .format import YoloPath @@ -33,16 +33,31 @@ def __len__(self): def categories(self): return self._parent.categories() - def __init__(self, config_path): + def __init__(self, config_path, image_info=None): super().__init__() if not osp.isfile(config_path): - raise Exception("Can't read dataset descriptor file '%s'" % \ + raise Exception("Can't read dataset descriptor file '%s'" % config_path) rootpath = osp.dirname(config_path) self._path = rootpath + assert image_info is None or isinstance(image_info, (str, dict)) + if image_info is None: + image_info = osp.join(rootpath, YoloPath.IMAGE_META_FILE) + if not osp.isfile(image_info): + image_info = {} + if isinstance(image_info, str): + if not osp.isfile(image_info): + raise Exception("Can't read image meta file '%s'" % image_info) + with open(image_info) as f: + image_info = {} + for line in f: + image_name, h, w = line.strip().split() + image_info[image_name] = (int(h), int(w)) + self._image_info = image_info + with open(config_path, 'r') as f: config_lines = f.readlines() @@ -77,10 +92,10 @@ def __init__(self, config_path): subset.items = OrderedDict( (osp.splitext(osp.basename(p))[0], p.strip()) for p in f) - for image_path in subset.items.values(): + for item_id, image_path in subset.items.items(): image_path = self._make_local_path(image_path) - if not osp.isfile(image_path): - raise Exception("Can't find image '%s'" % image_path) + if not osp.isfile(image_path) and item_id not in image_info: + raise Exception("Can't find image '%s'" % item_id) subsets[subset_name] = subset @@ -103,8 +118,10 @@ def _get(self, item_id, subset_name): if isinstance(item, str): image_path = self._make_local_path(item) - image = lazy_image(image_path) - h, w, _ = image().shape + image_size = self._image_info.get(item_id) + image = Image(path=image_path, size=image_size) + h, w = image.size + anno_path = osp.splitext(image_path)[0] + '.txt' annotations = self._parse_annotations(anno_path, w, h) diff --git a/datumaro/datumaro/plugins/yolo_format/format.py b/datumaro/datumaro/plugins/yolo_format/format.py index 8d44a9ba8fbb..c88c99d442d2 100644 --- a/datumaro/datumaro/plugins/yolo_format/format.py +++ b/datumaro/datumaro/plugins/yolo_format/format.py @@ -6,4 +6,6 @@ class YoloPath: DEFAULT_SUBSET_NAME = 'train' - SUBSET_NAMES = ['train', 'valid'] \ No newline at end of file + SUBSET_NAMES = ['train', 'valid'] + + IMAGE_META_FILE = 'images.meta' \ No newline at end of file diff --git a/datumaro/datumaro/plugins/yolo_format/importer.py b/datumaro/datumaro/plugins/yolo_format/importer.py index 26532d091286..fcee669dc6b9 100644 --- a/datumaro/datumaro/plugins/yolo_format/importer.py +++ b/datumaro/datumaro/plugins/yolo_format/importer.py @@ -15,18 +15,20 @@ def __call__(self, path, **extra_params): from datumaro.components.project import Project # cyclic import project = Project() - if not osp.exists(path): - raise Exception("Failed to find 'yolo' dataset at '%s'" % path) - 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: + raise Exception("Failed to find 'yolo' dataset at '%s'" % path) + for config_path in config_paths: log.info("Found a dataset at '%s'" % config_path) - source_name = osp.splitext(osp.basename(config_path))[0] + source_name = '%s_%s' % ( + osp.basename(osp.dirname(config_path)), + osp.splitext(osp.basename(config_path))[0]) project.add_source(source_name, { 'url': config_path, 'format': 'yolo', diff --git a/datumaro/tests/test_yolo_format.py b/datumaro/tests/test_yolo_format.py index 3af08781c368..e9a95108a9ee 100644 --- a/datumaro/tests/test_yolo_format.py +++ b/datumaro/tests/test_yolo_format.py @@ -59,7 +59,7 @@ class TestExtractor(Extractor): def __iter__(self): return iter([ DatasetItem(id=1, subset='train', - image=Image(size=(10, 15)), + image=Image(path='1.jpg', size=(10, 15)), annotations=[ Bbox(0, 2, 4, 2, label=2), Bbox(3, 3, 2, 3, label=4), @@ -84,3 +84,33 @@ def categories(self): parsed_dataset = YoloImporter()(test_dir).make_dataset() compare_datasets(self, source_dataset, parsed_dataset) + + def test_can_load_dataset_with_exact_image_info(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: + source_dataset = TestExtractor() + + YoloConverter()(source_dataset, test_dir) + + parsed_dataset = YoloImporter()(test_dir, + image_info={'1': (10, 15)}).make_dataset() + + compare_datasets(self, source_dataset, parsed_dataset)