Merge pull request #100 from roboflow/feature/support_for_yolo_dataset_export

feature/support for yolo dataset export

Author: SkalskiP

docs/changelog.md

@@ -1,3 +1,9 @@
+### 0.8.0 <small>May 17, 2023</small>
+
+- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: support for Dataset inheritance. Current `Dataset` got renamed to `DetectionDataset` and make it inherit from `BaseDataset`.
+- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: ability to save datasets in YOLO format using `DetectionDataset.as_yolo`.
+- Changed [[#100](https://github.com/roboflow/supervision/pull/100)]: default value of `approximation_percentage` parameter from `0.75` to `0.0` in `DetectionDataset.as_yolo` and `DetectionDataset.as_pascal_voc`.
+
 ### 0.7.0 <small>May 11, 2023</small>
 
 - Added [[#91](https://github.com/roboflow/supervision/pull/91)]: `Detections.from_yolo_nas` to enable seamless integration with [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) model.

docs/dataset/core.md

@@ -1,8 +1,8 @@
 !!! warning
 
-    `Dataset` API is still fluid and may change. If you use Dataset in your project until further notice, freeze the 
-    `supervision` version in your `requirements.txt`.
+    Dataset API is still fluid and may change. If you use Dataset API in your project until further notice, freeze the 
+    `supervision` version in your `requirements.txt` or `setup.py`.
 
-## Dataset
+## DetectionDataset
 
-:::supervision.dataset.core.Dataset
+:::supervision.dataset.core.DetectionDataset

setup.py

@@ -26,7 +26,8 @@ def get_version():
     install_requires=[
         'numpy>=1.20.0',
         'opencv-python',
-        'matplotlib'
+        'matplotlib',
+        'pyyaml'
     ],
     packages=find_packages(exclude=("tests",)),
     extras_require={

supervision/__init__.py

@@ -1,6 +1,6 @@
-__version__ = "0.7.0"
+__version__ = "0.8.0"
 
-from supervision.dataset.core import Dataset
+from supervision.dataset.core import BaseDataset, DetectionDataset
 from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
 from supervision.detection.core import Detections
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator

supervision/dataset/core.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Iterator
+from typing import Dict, Iterator, List, Optional, Tuple
 
 import cv2
 import numpy as np
@@ -11,15 +11,25 @@
     detections_to_pascal_voc,
     load_pascal_voc_annotations,
 )
-from supervision.dataset.formats.yolo import load_yolo_annotations
+from supervision.dataset.formats.yolo import (
+    load_yolo_annotations,
+    save_data_yaml,
+    save_yolo_annotations,
+)
+from supervision.dataset.ultils import save_dataset_images
 from supervision.detection.core import Detections
 from supervision.file import list_files_with_extensions
 
 
 @dataclass
-class Dataset:
+class BaseDataset:
+    pass
+
+
+@dataclass
+class DetectionDataset(BaseDataset):
     """
-    Dataclass containing information about the dataset.
+    Dataclass containing information about object detection dataset.
 
     Attributes:
         classes (List[str]): List containing dataset class names.
@@ -57,7 +67,7 @@ def as_pascal_voc(
         annotations_directory_path: Optional[str] = None,
         min_image_area_percentage: float = 0.0,
         max_image_area_percentage: float = 1.0,
-        approximation_percentage: float = 0.75,
+        approximation_percentage: float = 0.0,
     ) -> None:
         """
         Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
@@ -107,7 +117,7 @@ def as_pascal_voc(
     @classmethod
     def from_pascal_voc(
         cls, images_directory_path: str, annotations_directory_path: str
-    ) -> Dataset:
+    ) -> DetectionDataset:
         """
         Creates a Dataset instance from PASCAL VOC formatted data.
 
@@ -116,7 +126,7 @@ def from_pascal_voc(
             annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.
 
         Returns:
-            Dataset: A Dataset instance containing the loaded images and annotations.
+            DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
 
         Example:
             ```python
@@ -168,7 +178,7 @@ def from_pascal_voc(
         annotations = {
             image_name: detections for image_name, detections, _ in raw_annotations
         }
-        return Dataset(classes=classes, images=images, annotations=annotations)
+        return DetectionDataset(classes=classes, images=images, annotations=annotations)
 
     @classmethod
     def from_yolo(
@@ -177,7 +187,7 @@ def from_yolo(
         annotations_directory_path: str,
         data_yaml_path: str,
         force_masks: bool = False,
-    ) -> Dataset:
+    ) -> DetectionDataset:
         """
         Creates a Dataset instance from YOLO formatted data.
 
@@ -188,7 +198,7 @@ def from_yolo(
             force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
 
         Returns:
-            Dataset: A Dataset instance containing the loaded images and annotations.
+            DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
 
         Example:
             ```python
@@ -219,4 +229,50 @@ def from_yolo(
             data_yaml_path=data_yaml_path,
             force_masks=force_masks,
         )
-        return Dataset(classes=classes, images=images, annotations=annotations)
+        return DetectionDataset(classes=classes, images=images, annotations=annotations)
+
+    def as_yolo(
+        self,
+        images_directory_path: Optional[str] = None,
+        annotations_directory_path: Optional[str] = None,
+        data_yaml_path: Optional[str] = None,
+        min_image_area_percentage: float = 0.0,
+        max_image_area_percentage: float = 1.0,
+        approximation_percentage: float = 0.0,
+    ) -> None:
+        """
+        Exports the dataset to YOLO (You Only Look Once) format. This method saves the images and their corresponding
+        annotations in YOLO format, which is a simple text file that describes an object in the image. It also allows
+        for the optional saving of a data.yaml file, used in YOLOv5, that contains metadata about the dataset.
+
+        The method allows filtering the detections based on their area percentage and offers an option for polygon approximation.
+
+        Args:
+            images_directory_path (Optional[str]): The path to the directory where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_directory_path (Optional[str]): The path to the directory where the annotations in
+                YOLO format should be saved. If not provided, annotations will not be saved.
+            data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
+                If not provided, the file will not be saved.
+            min_image_area_percentage (float): The minimum percentage of detection area relative to
+                the image area for a detection to be included.
+            max_image_area_percentage (float): The maximum percentage of detection area relative to
+                the image area for a detection to be included.
+            approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
+                in the range [0, 1). This is useful for simplifying the annotations.
+        """
+        if images_directory_path is not None:
+            save_dataset_images(
+                images_directory_path=images_directory_path, images=self.images
+            )
+        if annotations_directory_path is not None:
+            save_yolo_annotations(
+                annotations_directory_path=annotations_directory_path,
+                images=self.images,
+                annotations=self.annotations,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+        if data_yaml_path is not None:
+            save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)

supervision/dataset/formats/pascal_voc.py

@@ -63,9 +63,6 @@ def detections_to_pascal_voc(
         str: An XML string in Pascal VOC format representing the detections.
     """
     height, width, depth = image_shape
-    image_area = height * width
-    minimum_detection_area = min_image_area_percentage * image_area
-    maximum_detection_area = max_image_area_percentage * image_area
 
     # Create root element
     annotation = Element("annotation")

supervision/dataset/formats/yolo.py

@@ -1,12 +1,15 @@
+import os
 from pathlib import Path
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
+import yaml
 
+from supervision.dataset.ultils import approximate_mask_with_polygons
 from supervision.detection.core import Detections
 from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
-from supervision.file import list_files_with_extensions, read_txt_file
+from supervision.file import list_files_with_extensions, read_txt_file, save_text_file
 
 
 def _parse_box(values: List[str]) -> np.ndarray:
@@ -78,6 +81,11 @@ def _extract_class_names(file_path: str) -> List[str]:
     return names
 
 
+def _image_name_to_annotation_name(image_name: str) -> str:
+    base_name, _ = os.path.splitext(image_name)
+    return base_name + ".txt"
+
+
 def yolo_annotations_to_detections(
     lines: List[str], resolution_wh: Tuple[int, int], with_masks: bool
 ) -> Detections:
@@ -130,7 +138,7 @@ def load_yolo_annotations(
         force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
 
     Returns:
-        Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image paths as keys and images as values, and a dictionary with image paths as keys and corresponding Detections instances as values.
+        Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values.
     """
     image_paths = list_files_with_extensions(
         directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
@@ -156,6 +164,93 @@ def load_yolo_annotations(
             lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
         )
 
-        images[str(image_path)] = image
-        annotations[str(image_path)] = annotation
+        images[image_path.name] = image
+        annotations[image_path.name] = annotation
     return classes, images, annotations
+
+
+def object_to_yolo(
+    xyxy: np.ndarray,
+    class_id: int,
+    image_shape: Tuple[int, int, int],
+    polygon: Optional[np.ndarray] = None,
+) -> str:
+    h, w, _ = image_shape
+    if polygon is None:
+        xyxy_relative = xyxy / np.array([w, h, w, h], dtype=np.float32)
+        x_min, y_min, x_max, y_max = xyxy_relative
+        x_center = (x_min + x_max) / 2
+        y_center = (y_min + y_max) / 2
+        width = x_max - x_min
+        height = y_max - y_min
+        return f"{int(class_id)} {x_center:.5f} {y_center:.5f} {width:.5f} {height:.5f}"
+    else:
+        polygon_relative = polygon / np.array([w, h], dtype=np.float32)
+        polygon_relative = polygon_relative.reshape(-1)
+        polygon_parsed = " ".join([f"{value:.5f}" for value in polygon_relative])
+        return f"{int(class_id)} {polygon_parsed}"
+
+
+def detections_to_yolo_annotations(
+    detections: Detections,
+    image_shape: Tuple[int, int, int],
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> List[str]:
+    annotation = []
+    for xyxy, mask, _, class_id, _ in detections:
+        if mask is not None:
+            polygons = approximate_mask_with_polygons(
+                mask=mask,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+            for polygon in polygons:
+                xyxy = polygon_to_xyxy(polygon=polygon)
+                next_object = object_to_yolo(
+                    xyxy=xyxy,
+                    class_id=class_id,
+                    image_shape=image_shape,
+                    polygon=polygon,
+                )
+                annotation.append(next_object)
+        else:
+            next_object = object_to_yolo(
+                xyxy=xyxy, class_id=class_id, image_shape=image_shape
+            )
+            annotation.append(next_object)
+    return annotation
+
+
+def save_yolo_annotations(
+    annotations_directory_path: str,
+    images: Dict[str, np.ndarray],
+    annotations: Dict[str, Detections],
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> None:
+    Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
+    for image_name, image in images.items():
+        detections = annotations[image_name]
+        yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
+        yolo_annotations_path = os.path.join(
+            annotations_directory_path, yolo_annotations_name
+        )
+        lines = detections_to_yolo_annotations(
+            detections=detections,
+            image_shape=image.shape,
+            min_image_area_percentage=min_image_area_percentage,
+            max_image_area_percentage=max_image_area_percentage,
+            approximation_percentage=approximation_percentage,
+        )
+        save_text_file(lines=lines, file_path=yolo_annotations_path)
+
+
+def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
+    data = {"nc": len(classes), "names": classes}
+    Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
+    with open(data_yaml_path, "w") as outfile:
+        yaml.dump(data, outfile, default_flow_style=False)

supervision/dataset/ultils.py

@@ -1,5 +1,8 @@
-from typing import List
+import os
+from pathlib import Path
+from typing import Dict, List
 
+import cv2
 import numpy as np
 
 from supervision.detection.utils import (
@@ -15,7 +18,7 @@ def approximate_mask_with_polygons(
     max_image_area_percentage: float = 1.0,
     approximation_percentage: float = 0.75,
 ) -> List[np.ndarray]:
-    height, width = mask
+    height, width = mask.shape
     image_area = height * width
     minimum_detection_area = min_image_area_percentage * image_area
     maximum_detection_area = max_image_area_percentage * image_area
@@ -35,3 +38,13 @@ def approximate_mask_with_polygons(
         approximate_polygon(polygon=polygon, percentage=approximation_percentage)
         for polygon in polygons
     ]
+
+
+def save_dataset_images(
+    images_directory_path: str, images: Dict[str, np.ndarray]
+) -> None:
+    Path(images_directory_path).mkdir(parents=True, exist_ok=True)
+
+    for image_name, image in images.items():
+        target_image_path = os.path.join(images_directory_path, image_name)
+        cv2.imwrite(target_image_path, image)