From 326a6640b41918f9dd16542c903b7325655854a3 Mon Sep 17 00:00:00 2001 From: "rustem.galiullin" Date: Thu, 9 Dec 2021 14:38:22 +0400 Subject: [PATCH] import and export yolov5 datasets with confidence (optional), update tests and docs --- .../user_guide/dataset_creation/datasets.rst | 6 +- docs/source/user_guide/export_datasets.rst | 8 ++- fiftyone/utils/yolo.py | 45 ++++++++++-- tests/unittests/import_export_tests.py | 70 +++++++++++-------- 4 files changed, 88 insertions(+), 41 deletions(-) diff --git a/docs/source/user_guide/dataset_creation/datasets.rst b/docs/source/user_guide/dataset_creation/datasets.rst index 57fd9e8f51..305dd0f47b 100644 --- a/docs/source/user_guide/dataset_creation/datasets.rst +++ b/docs/source/user_guide/dataset_creation/datasets.rst @@ -1973,11 +1973,12 @@ omitted, in which case the `data/` directory is listed to determine the available images. The TXT files in `data/` are space-delimited files where each row corresponds -to an object in the image of the same name, in the following format: +to an object in the image of the same name, in one the following formats: .. code-block:: text + where `` is the zero-based integer index of the object class label from `obj.names` and the bounding box coordinates are expressed as @@ -2214,11 +2215,12 @@ specific split being imported or exported is specified by the `split` argument to :class:`fiftyone.utils.yolo.YOLOv5DatasetImporter`. The TXT files in `labels/` are space-delimited files where each row corresponds -to an object in the image of the same name, in the following format: +to an object in the image of the same name, in one the following formats: .. code-block:: text + where `` is the zero-based integer index of the object class label from `names` and the bounding box coordinates are expressed as diff --git a/docs/source/user_guide/export_datasets.rst b/docs/source/user_guide/export_datasets.rst index c6989e2fb6..ce753caff0 100644 --- a/docs/source/user_guide/export_datasets.rst +++ b/docs/source/user_guide/export_datasets.rst @@ -2058,12 +2058,13 @@ and `images.txt` contains the list of images in `data/`: ... and the TXT files in `data/` are space-delimited files where each row -corresponds to an object in the image of the same name, in the following -format: +corresponds to an object in the image of the same name, in one of the following +formats: .. code-block:: text + # if include_confidence=True where `` is the zero-based integer index of the object class label from `obj.names` and the bounding box coordinates are expressed as @@ -2224,11 +2225,12 @@ specific split being imported or exported is specified by the `split` argument to :class:`fiftyone.utils.yolo.YOLOv5DatasetExporter`. The TXT files in `labels/` are space-delimited files where each row corresponds -to an object in the image of the same name, in the following format: +to an object in the image of the same name, in one of the following formats: .. code-block:: text + # if include_confidence=True where `` is the zero-based integer index of the object class label from `names` and the bounding box coordinates are expressed as diff --git a/fiftyone/utils/yolo.py b/fiftyone/utils/yolo.py index 86b382e064..ac6b49ad8b 100644 --- a/fiftyone/utils/yolo.py +++ b/fiftyone/utils/yolo.py @@ -572,6 +572,12 @@ class YOLOv4DatasetExporter( classes (None): the list of possible class labels. If not provided, this list will be extracted when :meth:`log_collection` is called, if possible + include_confidence (False): whether to include detection confidences in + the export. The supported values are: + + - ``False``: (default) do not include confidences + - ``True``: always include confidences + image_format (None): the image format to use when writing in-memory images to disk. By default, ``fiftyone.config.default_image_ext`` is used @@ -586,6 +592,7 @@ def __init__( images_path=None, export_media=None, classes=None, + include_confidence=False, image_format=None, ): data_path, export_media = self._parse_data_path( @@ -619,6 +626,7 @@ def __init__( self.images_path = images_path self.export_media = export_media self.classes = classes + self.include_confidence = include_confidence self.image_format = image_format self._classes = None @@ -693,6 +701,7 @@ def export_sample(self, image_or_path, detections, metadata=None): out_labels_path, self._labels_map_rev, dynamic_classes=self._dynamic_classes, + include_confidence=self.include_confidence, ) def close(self, *args): @@ -778,6 +787,12 @@ class YOLOv5DatasetExporter( classes (None): the list of possible class labels. If not provided, this list will be extracted when :meth:`log_collection` is called, if possible + include_confidence (False): whether to include detection confidences in + the export. The supported values are: + + - ``False``: (default) do not include confidences + - ``True``: always include confidences + image_format (None): the image format to use when writing in-memory images to disk. By default, ``fiftyone.config.default_image_ext`` is used @@ -792,6 +807,7 @@ def __init__( yaml_path=None, export_media=None, classes=None, + include_confidence=False, image_format=None, ): data_path, export_media = self._parse_data_path( @@ -821,6 +837,7 @@ def __init__( self.yaml_path = yaml_path self.export_media = export_media self.classes = classes + self.include_confidence = include_confidence self.image_format = image_format self._classes = None @@ -884,6 +901,7 @@ def export_sample(self, image_or_path, detections, metadata=None): out_labels_path, self._labels_map_rev, dynamic_classes=self._dynamic_classes, + include_confidence=self.include_confidence, ) def close(self, *args): @@ -917,7 +935,7 @@ class YOLOAnnotationWriter(object): """Class for writing annotations in YOLO-style TXT format.""" def write( - self, detections, txt_path, labels_map_rev, dynamic_classes=False + self, detections, txt_path, labels_map_rev, dynamic_classes=False, include_confidence=False ): """Writes the detections to disk. @@ -928,6 +946,7 @@ def write( integers dynamic_classes (False): whether to dynamically add new labels to ``labels_map_rev`` + include_confidence (False): whether to include confidence in exported file """ rows = [] for detection in detections.detections: @@ -946,7 +965,8 @@ def write( else: target = labels_map_rev[label] - row = _make_yolo_row(detection.bounding_box, target) + row = _make_yolo_row(detection.bounding_box, target, + confidence=detection.confidence if include_confidence else None) rows.append(row) _write_file_lines(rows, txt_path) @@ -956,9 +976,10 @@ def load_yolo_annotations(txt_path, classes): """Loads the YOLO-style annotations from the given TXT file. The txt file should be a space-delimited file where each row corresponds - to an object in the following format:: + to an object in one the following formats:: + where ``target`` is the zero-based integer index of the object class label from ``classes`` and the bounding box coordinates are expressed as relative @@ -1014,7 +1035,14 @@ def _get_yolo_v5_labels_path(image_path): def _parse_yolo_row(row, classes): - target, xc, yc, w, h = row.split() + row_vals = row.split() + if len(row_vals) == 5: + (target, xc, yc, w, h), conf = row_vals, None + elif len(row_vals) == 6: + target, xc, yc, w, h, conf = row_vals + conf = float(conf) + else: + raise NotImplementedError(f"rows with length {len(row_vals)} are not supported") try: label = classes[int(target)] @@ -1028,14 +1056,17 @@ def _parse_yolo_row(row, classes): float(h), ] - return fol.Detection(label=label, bounding_box=bounding_box) + return fol.Detection(label=label, bounding_box=bounding_box, confidence=conf) -def _make_yolo_row(bounding_box, target): +def _make_yolo_row(bounding_box, target, confidence=None): xtl, ytl, w, h = bounding_box xc = xtl + 0.5 * w yc = ytl + 0.5 * h - return "%d %f %f %f %f" % (target, xc, yc, w, h) + if confidence is None: + return "%d %f %f %f %f" % (target, xc, yc, w, h) + else: + return "%d %f %f %f %f %f" % (target, xc, yc, w, h, confidence) def _read_yaml_file(path): diff --git a/tests/unittests/import_export_tests.py b/tests/unittests/import_export_tests.py index 80615c9dbc..e463858a00 100644 --- a/tests/unittests/import_export_tests.py +++ b/tests/unittests/import_export_tests.py @@ -1013,24 +1013,30 @@ def test_yolov4_dataset(self): export_dir = self._new_dir() - dataset.export( - export_dir=export_dir, - dataset_type=fo.types.YOLOv4Dataset, - label_field="predictions", - ) + for with_confidence in [False, True]: + dataset.export( + export_dir=export_dir, + dataset_type=fo.types.YOLOv4Dataset, + label_field="predictions", + include_confidence=with_confidence, + ) - dataset2 = fo.Dataset.from_dir( - dataset_dir=export_dir, - dataset_type=fo.types.YOLOv4Dataset, - label_field="predictions", - include_all_data=True, - ) + dataset2 = fo.Dataset.from_dir( + dataset_dir=export_dir, + dataset_type=fo.types.YOLOv4Dataset, + label_field="predictions", + include_all_data=True, + ) - self.assertEqual(len(dataset), len(dataset2)) - self.assertEqual( - dataset.count("predictions.detections"), - dataset2.count("predictions.detections"), - ) + self.assertEqual(len(dataset), len(dataset2)) + self.assertEqual( + dataset.count("predictions.detections"), + dataset2.count("predictions.detections"), + ) + self.assertEqual( + dataset.bounds("predictions.detections.confidence") if with_confidence else (None, None), + dataset2.bounds("predictions.detections.confidence"), + ) # Labels-only @@ -1065,21 +1071,27 @@ def test_yolov5_dataset(self): export_dir = self._new_dir() - dataset.export( - export_dir=export_dir, dataset_type=fo.types.YOLOv5Dataset, - ) + for with_confidence in [False, True]: + dataset.export( + export_dir=export_dir, dataset_type=fo.types.YOLOv5Dataset, + include_confidence=with_confidence, + ) - dataset2 = fo.Dataset.from_dir( - dataset_dir=export_dir, - dataset_type=fo.types.YOLOv5Dataset, - label_field="predictions", - ) + dataset2 = fo.Dataset.from_dir( + dataset_dir=export_dir, + dataset_type=fo.types.YOLOv5Dataset, + label_field="predictions", + ) - self.assertEqual(len(dataset), len(dataset2)) - self.assertEqual( - dataset.count("predictions.detections"), - dataset2.count("predictions.detections"), - ) + self.assertEqual(len(dataset), len(dataset2)) + self.assertEqual( + dataset.count("predictions.detections"), + dataset2.count("predictions.detections"), + ) + self.assertEqual( + dataset.bounds("predictions.detections.confidence") if with_confidence else (None, None), + dataset2.bounds("predictions.detections.confidence"), + ) class ImageSegmentationDatasetTests(ImageDatasetTests):