diff --git a/CHANGELOG.md b/CHANGELOG.md index e85315d9d12f..2be78b941526 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Allowed adding "difficult", "truncated", "occluded" attributes when converting to Pascal VOC if these attributes are not present () - Empty lines in YOLO annotations are ignored () - Export in VOC format when no image info is available () +- Fixed saving attribute in WiderFace extractor () ### Security - diff --git a/datumaro/plugins/widerface_format.py b/datumaro/plugins/widerface_format.py index 96796b09f350..a8439dc83ce0 100644 --- a/datumaro/plugins/widerface_format.py +++ b/datumaro/plugins/widerface_format.py @@ -10,6 +10,7 @@ from datumaro.components.converter import Converter from datumaro.components.extractor import (AnnotationType, Bbox, DatasetItem, Importer, Label, LabelCategories, SourceExtractor) +from datumaro.util import str_to_bool class WiderFacePath: @@ -119,7 +120,10 @@ def _load_items(self, path): i = 4 for attr in WiderFacePath.BBOX_ATTRIBUTES: if bbox_list[i] != '-': - attributes[attr] = bbox_list[i] + if bbox_list[i] in ['True', 'False']: + attributes[attr] = str_to_bool(bbox_list[i]) + else: + attributes[attr] = bbox_list[i] i += 1 annotations.append(Bbox( diff --git a/tests/test_widerface_format.py b/tests/test_widerface_format.py index 8e2586999ef2..4d733ee82a4e 100644 --- a/tests/test_widerface_format.py +++ b/tests/test_widerface_format.py @@ -43,7 +43,9 @@ def test_can_save_and_load(self): 'blur': '2', 'expression': '1', 'illumination': '0', 'occluded': '0', 'pose': '1', 'invalid': '0'}), Bbox(0, 2, 3, 2, label=0, attributes={ - 'occluded': 'False'}), + 'occluded': False}), + Bbox(0, 3, 4, 2, label=0, attributes={ + 'occluded': True}), Bbox(0, 2, 4, 2, label=0), Bbox(0, 7, 3, 2, label=0, attributes={ 'blur': '2', 'expression': '1', 'illumination': '0',