diff --git a/docs/datasets.md b/docs/datasets/core.md
similarity index 97%
rename from docs/datasets.md
rename to docs/datasets/core.md
index 73931515..03d0c196 100644
--- a/docs/datasets.md
+++ b/docs/datasets/core.md
@@ -1,5 +1,6 @@
---
comments: true
+status: new
---
# Datasets
diff --git a/docs/datasets/utils.md b/docs/datasets/utils.md
new file mode 100644
index 00000000..6be56303
--- /dev/null
+++ b/docs/datasets/utils.md
@@ -0,0 +1,18 @@
+---
+comments: true
+status: new
+---
+
+# Datasets Utils
+
+
+
+:::supervision.dataset.utils.rle_to_mask
+
+
+
+:::supervision.dataset.utils.mask_to_rle
diff --git a/mkdocs.yml b/mkdocs.yml
index cf206a82..281c40c9 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -61,7 +61,9 @@ nav:
- Detection Smoother: detection/tools/smoother.md
- Save Detections: detection/tools/save_detections.md
- Trackers: trackers.md
- - Datasets: datasets.md
+ - Datasets:
+ - Core: datasets/core.md
+ - Utils: datasets/utils.md
- Utils:
- Video: utils/video.md
- Image: utils/image.md
diff --git a/supervision/__init__.py b/supervision/__init__.py
index bb526514..b7f7a5e8 100644
--- a/supervision/__init__.py
+++ b/supervision/__init__.py
@@ -34,6 +34,10 @@ from supervision.dataset.core import (
ClassificationDataset,
DetectionDataset,
)
+from supervision.dataset.utils import (
+ rle_to_mask,
+ mask_to_rle,
+)
from supervision.detection.annotate import BoxAnnotator
from supervision.detection.core import Detections
from supervision.detection.line_zone import LineZone, LineZoneAnnotator
diff --git a/supervision/dataset/utils.py b/supervision/dataset/utils.py
index b0b5577e..ed7062b9 100644
--- a/supervision/dataset/utils.py
+++ b/supervision/dataset/utils.py
@@ -2,7 +2,7 @@ import copy
import os
import random
from pathlib import Path
-from typing import Dict, List, Optional, Tuple, TypeVar
+from typing import Dict, List, Optional, Tuple, Union, TypeVar
import cv2
import numpy as np
@@ -133,33 +133,42 @@ def train_test_split(
def rle_to_mask(
- rle: npt.NDArray[np.int_], resolution_wh: Tuple[int, int]
+ rle: Union[npt.NDArray[np.int_], List[int]], resolution_wh: Tuple[int, int]
) -> npt.NDArray[np.bool_]:
"""
Converts run-length encoding (RLE) to a binary mask.
Args:
- rle (npt.NDArray[np.int_]): The 1D RLE array, the format used in the COCO
- dataset (column-wise encoding, values of an array with even indices
- represent the number of pixels assigned as background,
+ rle (Union[npt.NDArray[np.int_], List[int]]): The 1D RLE array, the format
+ used in the COCO dataset (column-wise encoding, values of an array with
+ even indices represent the number of pixels assigned as background,
values of an array with odd indices represent the number of pixels
assigned as foreground object).
resolution_wh (Tuple[int, int]): The width (w) and height (h)
- of the desired binary mask resolution.
+ of the desired binary mask.
Returns:
- npt.NDArray[np.bool_]: The generated 2D Boolean mask of shape (h,w),
- where the foreground object is marked with `True`'s and the rest
- is filled with `False`'s.
+ The generated 2D Boolean mask of shape `(h, w)`, where the foreground object is
+ marked with `True`'s and the rest is filled with `False`'s.
Raises:
AssertionError: If the sum of pixels encoded in RLE differs from the
- number of pixels in the expected mask (computed based on resolution_wh).
+ number of pixels in the expected mask (computed based on resolution_wh).
Examples:
- rle = [2, 2, 2], resolution_wh = [3, 2] -> mask = [[False, True, False],
- [False, True, False]]
+ ```python
+ import supervision as sv
+
+ sv.rle_to_mask([2, 2, 2], (3, 2))
+ # array([
+ # [False, True, False],
+ # [False, True, False]
+ # ])
+ ```
"""
+ if isinstance(rle, list):
+ rle = np.array(rle, dtype=int)
+
width, height = resolution_wh
assert width * height == np.sum(rle), (
@@ -186,20 +195,33 @@ def mask_to_rle(mask: npt.NDArray[np.bool_]) -> List[int]:
object and `False` indicates background.
Returns:
- List[int]: the run-length encoded mask. Values of a list with even indices
+ The run-length encoded mask. Values of a list with even indices
represent the number of pixels assigned as background (`False`), values
of a list with odd indices represent the number of pixels assigned
as foreground object (`True`).
Raises:
- AssertionError: If imput mask is not 2D or is empty.
+ AssertionError: If input mask is not 2D or is empty.
- Examples:
- mask = [[False, True, True], -> rle = [2, 4]
- [False, True, True]]
+ Examples:
+ ```python
+ import numpy as np
+ import supervision as sv
- mask = [[True, True, True], -> rle = [0, 6]
- [True, True, True]]
+ mask = np.array([
+ [False, True, True],
+ [False, True, True]
+ ])
+ sv.mask_to_rle(mask)
+ # [2, 4]
+
+ mask = np.array([
+ [True, True, True],
+ [True, True, True]
+ ])
+ sv.mask_to_rle(mask)
+ # [0, 6]
+ ```
"""
assert mask.ndim == 2, "Input mask must be 2D"
assert mask.size != 0, "Input mask cannot be empty"
diff --git a/test/dataset/test_utils.py b/test/dataset/test_utils.py
index e30b4e5f..41e1da5b 100644
--- a/test/dataset/test_utils.py
+++ b/test/dataset/test_utils.py
@@ -286,7 +286,7 @@ def test_map_detections_class_id(
), # raises AssertionError because mask is empty
],
)
-def test_mask_to_rle_conversion(
+def test_mask_to_rle(
mask: npt.NDArray[np.bool_], expected_rle: List[int], exception: Exception
) -> None:
with exception:
@@ -302,7 +302,13 @@ def test_mask_to_rle_conversion(
[3, 3],
np.zeros((3, 3)).astype(bool),
DoesNotRaise(),
- ), # mask with background only (mask with only False values)
+ ), # mask with background only (mask with only False values); rle as array
+ (
+ [9],
+ [3, 3],
+ np.zeros((3, 3)).astype(bool),
+ DoesNotRaise(),
+ ), # mask with background only (mask with only False values); rle as list
(
np.array([0, 9]),
[3, 3],
@@ -346,7 +352,7 @@ def test_mask_to_rle_conversion(
# number of pixels in expected mask (width x height).
],
)
-def test_rle_to_mask_convertion(
+def test_rle_to_mask(
rle: npt.NDArray[np.int_],
resolution_wh: Tuple[int, int],
expected_mask: npt.NDArray[np.bool_],