From d21c98a4d2fda92300d4931f952fcc0e91669959 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 13 May 2024 12:06:38 +0200 Subject: [PATCH] docs updated; `rle_to_mask` and `mask_to_rle` added to `__init__.py` --- docs/{datasets.md => datasets/core.md} | 1 + docs/datasets/utils.md | 18 ++++++++ mkdocs.yml | 4 +- supervision/__init__.py | 4 ++ supervision/dataset/utils.py | 60 ++++++++++++++++++-------- test/dataset/test_utils.py | 12 ++++-- 6 files changed, 76 insertions(+), 23 deletions(-) rename docs/{datasets.md => datasets/core.md} (97%) create mode 100644 docs/datasets/utils.md 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 + +
+

rle_to_mask

+
+ +:::supervision.dataset.utils.rle_to_mask + +
+

mask_to_rle

+
+ +:::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_],