Fix saving attribute in WiderFace extractor (cvat-ai#251)

* add fixes * update changelog
TOsmanov · May 26, 2021 · bac10c7 · bac10c7
1 parent 58892de
commit bac10c7
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git 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 (<https://github.com/openvinotoolkit/datumaro/pull/216>)
 - Empty lines in YOLO annotations are ignored (<https://github.com/openvinotoolkit/datumaro/pull/221>)
 - Export in VOC format when no image info is available (<https://github.com/openvinotoolkit/datumaro/pull/239>)
+- Fixed saving attribute in WiderFace extractor (<https://github.com/openvinotoolkit/datumaro/pull/251>)
 
 ### Security
 -

diff --git 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
@@ -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',