Merge pull request #1465 from Rusteam/yolo-conf

Adding support for importing/exporting confidence in YOLO formats

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
 
     <target> <x-center> <y-center> <width> <height>
+    <target> <x-center> <y-center> <width> <height> <confidence>
 
 where `<target>` 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
 
     <target> <x-center> <y-center> <width> <height>
+    <target> <x-center> <y-center> <width> <height> <confidence>
 
 where `<target>` is the zero-based integer index of the object class label from
 `names` and the bounding box coordinates are expressed as

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
 
     <target> <x-center> <y-center> <width> <height>
+    <target> <x-center> <y-center> <width> <height> <confidence> # if include_confidence=True
 
 where `<target>` 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
 
     <target> <x-center> <y-center> <width> <height>
+    <target> <x-center> <y-center> <width> <height> <confidence> # if include_confidence=True
 
 where `<target>` is the zero-based integer index of the object class label from
 `names` and the bounding box coordinates are expressed as

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::
 
         <target> <x-center> <y-center> <width> <height>
+        <target> <x-center> <y-center> <width> <height> <confidence>
 
     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):

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):