From f489c176a6cbec4b0f749012b1c07c003922dfb6 Mon Sep 17 00:00:00 2001 From: Anastasia Yasakova Date: Tue, 19 Jan 2021 12:25:26 +0300 Subject: [PATCH] Add label support in WiderFace dataset format (#90) * Add label support in WiderFace dataset format * add labels.txt --- CHANGELOG.md | 2 +- datumaro/plugins/widerface_format.py | 71 ++++++++++++++++++++--- tests/assets/widerface_dataset/labels.txt | 2 + tests/test_widerface_format.py | 33 +++++++---- 4 files changed, 86 insertions(+), 22 deletions(-) create mode 100644 tests/assets/widerface_dataset/labels.txt diff --git a/CHANGELOG.md b/CHANGELOG.md index 91a17fe4e6b2..7f8f08fe9b3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added -- `WiderFace` dataset format () +- `WiderFace` dataset format (, ) - Function to transform annotations to labels () - Task-specific Splitter (, ) - `VGGFace2` dataset format (, ) diff --git a/datumaro/plugins/widerface_format.py b/datumaro/plugins/widerface_format.py index 7a9efea52258..58c25540b1f7 100644 --- a/datumaro/plugins/widerface_format.py +++ b/datumaro/plugins/widerface_format.py @@ -9,7 +9,7 @@ from datumaro.components.converter import Converter from datumaro.components.extractor import (AnnotationType, Bbox, DatasetItem, - Importer, SourceExtractor) + Importer, Label, LabelCategories, SourceExtractor) class WiderFacePath: @@ -17,6 +17,8 @@ class WiderFacePath: ANNOTATIONS_DIR = 'wider_face_split' IMAGES_DIR = 'images' SUBSET_DIR = 'WIDER_' + LABELS_FILE = 'labels.txt' + IMAGES_DIR_NO_LABEL = 'no_label' BBOX_ATTRIBUTES = ['blur', 'expression', 'illumination', 'occluded', 'pose', 'invalid'] @@ -33,8 +35,31 @@ def __init__(self, path): subset = subset.split('_')[2] super().__init__(subset=subset) + self._categories = self._load_categories() self._items = list(self._load_items(path).values()) + def _load_categories(self): + self._categories[AnnotationType.label] = LabelCategories() + label_cat = LabelCategories() + path = osp.join(self._dataset_dir, WiderFacePath.LABELS_FILE) + if osp.isfile(path): + with open(path, encoding='utf-8') as labels_file: + labels = [s.strip() for s in labels_file] + for label in labels: + label_cat.add(label) + else: + subset_path = osp.join(self._dataset_dir, + WiderFacePath.SUBSET_DIR + self._subset, + WiderFacePath.IMAGES_DIR) + if osp.isdir(subset_path): + for images_dir in sorted(os.listdir(subset_path)): + if osp.isdir(osp.join(subset_path, images_dir)) and \ + images_dir != WiderFacePath.IMAGES_DIR_NO_LABEL: + if '--' in images_dir: + images_dir = images_dir.split('--')[1] + label_cat.add(images_dir) + return { AnnotationType.label: label_cat } + def _load_items(self, path): items = {} with open(path, 'r') as f: @@ -48,10 +73,19 @@ def _load_items(self, path): image_path = osp.join(self._dataset_dir, WiderFacePath.SUBSET_DIR + self._subset, WiderFacePath.IMAGES_DIR, image[:-1]) item_id = image[:-(len(WiderFacePath.IMAGE_EXT) + 1)] + annotations = [] + if '/' in item_id: + label_name = item_id.split('/')[0] + if '--' in label_name: + label_name = label_name.split('--')[1] + if label_name != WiderFacePath.IMAGES_DIR_NO_LABEL: + label = \ + self._categories[AnnotationType.label].find(label_name)[0] + annotations.append(Label(label=label)) + item_id = item_id[len(item_id.split('/')[0]) + 1:] bbox_count = lines[image_id + 1] bbox_lines = lines[image_id + 2 : image_id + int(bbox_count) + 2] - annotations = [] for bbox in bbox_lines: bbox_list = bbox.split() if len(bbox_list) >= 4: @@ -63,8 +97,8 @@ def _load_items(self, path): attributes[attr] = int(bbox_list[i]) i += 1 annotations.append(Bbox( - int(bbox_list[0]), int(bbox_list[1]), - int(bbox_list[2]), int(bbox_list[3]), + float(bbox_list[0]), float(bbox_list[1]), + float(bbox_list[2]), float(bbox_list[3]), attributes = attributes )) @@ -83,18 +117,37 @@ class WiderFaceConverter(Converter): def apply(self): save_dir = self._save_dir - os.makedirs(save_dir, exist_ok=True) + label_categories = self._extractor.categories()[AnnotationType.label] + + labels_path = osp.join(save_dir, WiderFacePath.LABELS_FILE) + with open(labels_path, 'w', encoding='utf-8') as f: + f.write('\n'.join(label.name for label in label_categories)) + for subset_name, subset in self._extractor.subsets().items(): subset_dir = osp.join(save_dir, WiderFacePath.SUBSET_DIR + subset_name) wider_annotation = '' for item in subset: - wider_annotation += '%s\n' % (item.id + WiderFacePath.IMAGE_EXT) - if item.has_image and self._save_images: - self._save_image(item, osp.join(save_dir, subset_dir, - WiderFacePath.IMAGES_DIR, item.id + WiderFacePath.IMAGE_EXT)) + labels = [a.label for a in item.annotations + if a.type == AnnotationType.label] + if labels: + wider_annotation += '%s\n' % (str(labels[0]) + '--' \ + + label_categories[labels[0]].name + '/' \ + + item.id + WiderFacePath.IMAGE_EXT) + if item.has_image and self._save_images: + self._save_image(item, osp.join(save_dir, subset_dir, + WiderFacePath.IMAGES_DIR, str(labels[0]) + '--' \ + + label_categories[labels[0]].name + '/' + item.id \ + + WiderFacePath.IMAGE_EXT)) + else: + wider_annotation += '%s\n' % (WiderFacePath.IMAGES_DIR_NO_LABEL \ + + '/' + item.id + WiderFacePath.IMAGE_EXT) + if item.has_image and self._save_images: + self._save_image(item, osp.join(save_dir, subset_dir, + WiderFacePath.IMAGES_DIR, WiderFacePath.IMAGES_DIR_NO_LABEL \ + + '/' + item.id + WiderFacePath.IMAGE_EXT)) bboxes = [a for a in item.annotations if a.type == AnnotationType.bbox] diff --git a/tests/assets/widerface_dataset/labels.txt b/tests/assets/widerface_dataset/labels.txt new file mode 100644 index 000000000000..d5ed8e24452a --- /dev/null +++ b/tests/assets/widerface_dataset/labels.txt @@ -0,0 +1,2 @@ +Parade +Handshaking \ No newline at end of file diff --git a/tests/test_widerface_format.py b/tests/test_widerface_format.py index ed51379ba6bd..d93f7bf240b4 100644 --- a/tests/test_widerface_format.py +++ b/tests/test_widerface_format.py @@ -2,7 +2,8 @@ from unittest import TestCase import numpy as np -from datumaro.components.extractor import Bbox, DatasetItem +from datumaro.components.extractor import (AnnotationType, Bbox, DatasetItem, + Label, LabelCategories) from datumaro.components.dataset import Dataset from datumaro.plugins.widerface_format import WiderFaceConverter, WiderFaceImporter from datumaro.util.test_utils import TestDir, compare_datasets @@ -17,6 +18,7 @@ def test_can_save_and_load(self): Bbox(0, 1, 2, 3, attributes = { 'blur': 2, 'expression': 0, 'illumination': 0, 'occluded': 0, 'pose': 2, 'invalid': 0}), + Label(0), ] ), DatasetItem(id='2', subset='train', image=np.ones((10, 10, 3)), @@ -30,6 +32,7 @@ def test_can_save_and_load(self): Bbox(2, 1, 2, 3, attributes = { 'blur': 2, 'expression': 0, 'illumination': 0, 'occluded': 0, 'pose': 0, 'invalid': 1}), + Label(1), ] ), @@ -47,11 +50,14 @@ def test_can_save_and_load(self): ), DatasetItem(id='4', subset='val', image=np.ones((8, 8, 3))), - ]) + ], categories={ + AnnotationType.label: LabelCategories.from_iterable( + 'label_' + str(i) for i in range(3)), + }) with TestDir() as test_dir: WiderFaceConverter.convert(source_dataset, test_dir, save_images=True) - parsed_dataset = WiderFaceImporter()(test_dir).make_dataset() + parsed_dataset = Dataset.import_from(test_dir, 'wider_face') compare_datasets(self, source_dataset, parsed_dataset) @@ -65,11 +71,11 @@ def test_can_save_dataset_with_no_subsets(self): 'occluded': 0, 'pose': 2, 'invalid': 0}), ] ), - ]) + ], categories=[]) with TestDir() as test_dir: WiderFaceConverter.convert(source_dataset, test_dir, save_images=True) - parsed_dataset = WiderFaceImporter()(test_dir).make_dataset() + parsed_dataset = Dataset.import_from(test_dir, 'wider_face') compare_datasets(self, source_dataset, parsed_dataset) @@ -85,7 +91,7 @@ def test_can_save_dataset_with_non_widerface_attributes(self): 'non-widerface attribute': 0}), ] ), - ]) + ], categories=[]) target_dataset = Dataset.from_iterable([ DatasetItem(id='a/b/1', image=np.ones((8, 8, 3)), @@ -96,11 +102,11 @@ def test_can_save_dataset_with_non_widerface_attributes(self): Bbox(1, 1, 2, 2), ] ), - ]) + ], categories=[]) with TestDir() as test_dir: WiderFaceConverter.convert(source_dataset, test_dir, save_images=True) - parsed_dataset = WiderFaceImporter()(test_dir).make_dataset() + parsed_dataset = Dataset.import_from(test_dir, 'wider_face') compare_datasets(self, target_dataset, parsed_dataset) @@ -112,15 +118,16 @@ def test_can_detect(self): def test_can_import(self): expected_dataset = Dataset.from_iterable([ - DatasetItem(id='0--Parade/0_Parade_image_01', subset='train', + DatasetItem(id='0_Parade_image_01', subset='train', image=np.ones((10, 15, 3)), annotations=[ Bbox(1, 2, 2, 2, attributes = { 'blur': 0, 'expression': 0, 'illumination': 0, 'occluded': 0, 'pose': 0, 'invalid': 0}), + Label(0), ] ), - DatasetItem(id='1--Handshaking/1_Handshaking_image_02', subset='train', + DatasetItem(id='1_Handshaking_image_02', subset='train', image=np.ones((10, 15, 3)), annotations=[ Bbox(1, 1, 2, 2, attributes = { @@ -129,9 +136,10 @@ def test_can_import(self): Bbox(5, 1, 2, 2, attributes = { 'blur': 0, 'expression': 0, 'illumination': 1, 'occluded': 0, 'pose': 0, 'invalid': 0}), + Label(1), ] ), - DatasetItem(id='0--Parade/0_Parade_image_03', subset='val', + DatasetItem(id='0_Parade_image_03', subset='val', image=np.ones((10, 15, 3)), annotations=[ Bbox(0, 0, 1, 1, attributes = { @@ -143,9 +151,10 @@ def test_can_import(self): Bbox(5, 6, 1, 1, attributes = { 'blur': 2, 'expression': 0, 'illumination': 0, 'occluded': 0, 'pose': 2, 'invalid': 0}), + Label(0), ] ), - ]) + ], categories= ['Parade', 'Handshaking']) dataset = Dataset.import_from(DUMMY_DATASET_DIR, 'wider_face')