1114 rle support for coco

docs/datasets.md → docs/datasets/core.md

@@ -1,5 +1,6 @@
 ---
 comments: true
+status: new
 ---
 
 # Datasets

docs/datasets/utils.md

@@ -0,0 +1,18 @@
+---
+comments: true
+status: new
+---
+
+# Datasets Utils
+
+<div class="md-typeset">
+ <h2><a href="#supervision.dataset.utils.rle_to_mask">rle_to_mask</a></h2>
+</div>
+
+:::supervision.dataset.utils.rle_to_mask
+
+<div class="md-typeset">
+ <h2><a href="#supervision.dataset.utils.mask_to_rle">mask_to_rle</a></h2>
+</div>
+
+:::supervision.dataset.utils.mask_to_rle

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

supervision/__init__.py

@@ -34,6 +34,7 @@
  ClassificationDataset,
  DetectionDataset,
 )
+from supervision.dataset.utils import mask_to_rle, rle_to_mask
 from supervision.detection.annotate import BoxAnnotator
 from supervision.detection.core import Detections
 from supervision.detection.line_zone import LineZone, LineZoneAnnotator

supervision/dataset/core.py

@@ -430,6 +430,10 @@ def as_coco(
  """
  Exports the dataset to COCO format. This method saves the
  images and their corresponding annotations in COCO format.
+ The format of the mask is determined automatically:
+ when a mask consists of multiple disconnected elements
+ or has holes the RLE format is used,
+ otherwise, the mask is encoded as a polygon.
 
  Args:
  images_directory_path (Optional[str]): The path to the directory

supervision/dataset/formats/coco.py

@@ -5,10 +5,13 @@
 
 import cv2
 import numpy as np
+import numpy.typing as npt
 
 from supervision.dataset.utils import (
  approximate_mask_with_polygons,
  map_detections_class_id,
+ mask_to_rle,
+ rle_to_mask,
 )
 from supervision.detection.core import Detections
 from supervision.detection.utils import polygon_to_mask
@@ -57,13 +60,24 @@ def group_coco_annotations_by_image_id(
  return annotations
 
 
-def _polygons_to_masks(
- polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
-) -> np.ndarray:
+def coco_annotations_to_masks(
+ image_annotations: List[dict], resolution_wh: Tuple[int, int]
+) -> npt.NDArray[np.bool_]:
  return np.array(
  [
- polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
- for polygon in polygons
+ rle_to_mask(
+ rle=np.array(image_annotation["segmentation"]["counts"]),
+ resolution_wh=resolution_wh,
+ )
+ if image_annotation["iscrowd"]
+ else polygon_to_mask(
+ polygon=np.reshape(
+ np.asarray(image_annotation["segmentation"], dtype=np.int32),
+ (-1, 2),
+ ),
+ resolution_wh=resolution_wh,
+ )
+ for image_annotation in image_annotations
  ],
  dtype=bool,
  )
@@ -83,20 +97,33 @@ def coco_annotations_to_detections(
  xyxy[:, 2:4] += xyxy[:, 0:2]
 
  if with_masks:
- polygons = [
- np.reshape(
- np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2)
- )
- for image_annotation in image_annotations
- ]
- mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
+ mask = coco_annotations_to_masks(
+ image_annotations=image_annotations, resolution_wh=resolution_wh
+ )
  return Detections(
  class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
  )
 
  return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
 
 
+def _mask_has_holes(mask: np.ndarray) -> bool:
+ mask_uint8 = mask.astype(np.uint8)
+ _, hierarchy = cv2.findContours(mask_uint8, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
+ parent_countour_index = 3
+ for h in hierarchy[0]:
+ if h[parent_countour_index] != -1:
+ return True
+ return False
+
+
+def _mask_has_multiple_segments(mask: np.ndarray) -> bool:
+ mask_uint8 = mask.astype(np.uint8)
+ labels = np.zeros_like(mask_uint8, dtype=np.int32)
+ number_of_labels, _ = cv2.connectedComponents(mask_uint8, labels, connectivity=4)
+ return number_of_labels > 2
+
+
 def detections_to_coco_annotations(
  detections: Detections,
  image_id: int,
@@ -108,24 +135,37 @@ def detections_to_coco_annotations(
  coco_annotations = []
  for xyxy, mask, _, class_id, _, _ in detections:
  box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
- polygon = []
+ segmentation = []
+ iscrowd = 0
  if mask is not None:
- polygon = list(
- 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,
- )[0].flatten()
+ iscrowd = _mask_has_holes(mask=mask) or _mask_has_multiple_segments(
+ mask=mask
  )
+
+ if iscrowd:
+ segmentation = {
+ "counts": mask_to_rle(mask=mask),
+ "size": list(mask.shape[:2]),
+ }
+ else:
+ segmentation = [
+ list(
+ 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,
+ )[0].flatten()
+ )
+ ] # multicomponent masks supported only for rle format
  coco_annotation = {
  "id": annotation_id,
  "image_id": image_id,
  "category_id": int(class_id),
  "bbox": [xyxy[0], xyxy[1], box_width, box_height],
  "area": box_width * box_height,
- "segmentation": [polygon] if polygon else [],
- "iscrowd": 0,
+ "segmentation": segmentation,
+ "iscrowd": iscrowd,
  }
  coco_annotations.append(coco_annotation)
  annotation_id += 1

supervision/dataset/utils.py

@@ -2,10 +2,11 @@
 import os
 import random
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, TypeVar
+from typing import Dict, List, Optional, Tuple, TypeVar, Union
 
 import cv2
 import numpy as np
+import numpy.typing as npt
 
 from supervision.detection.core import Detections
 from supervision.detection.utils import (
@@ -129,3 +130,115 @@ def train_test_split(
 
  split_index = int(len(data) * train_ratio)
  return data[:split_index], data[split_index:]
+
+
+def rle_to_mask(
+ 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 (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.
+
+ Returns:
+ 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).
+
+ Examples:
+ ```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), (
+ "the sum of the number of pixels in the RLE must be the same "
+ "as the number of pixels in the expected mask"
+ )
+
+ zero_one_values = np.zeros(shape=(rle.size, 1), dtype=np.uint8)
+ zero_one_values[1::2] = 1
+
+ decoded_rle = np.repeat(zero_one_values, rle, axis=0)
+ decoded_rle = np.append(
+ decoded_rle, np.zeros(width * height - len(decoded_rle), dtype=np.uint8)
+ )
+ return decoded_rle.reshape((height, width), order="F")
+
+
+def mask_to_rle(mask: npt.NDArray[np.bool_]) -> List[int]:
+ """
+ Converts a binary mask into a run-length encoding (RLE).
+
+ Args:
+ mask (npt.NDArray[np.bool_]): 2D binary mask where `True` indicates foreground
+ object and `False` indicates background.
+
+ Returns:
+ 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 input mask is not 2D or is empty.
+
+ Examples:
+ ```python
+ import numpy as np
+ import supervision as sv
+
+ 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"
+
+ on_value_change_indices = np.where(
+ mask.ravel(order="F") != np.roll(mask.ravel(order="F"), 1)
+ )[0]
+
+ on_value_change_indices = np.append(on_value_change_indices, mask.size)
+ # need to add 0 at the beginning when the same value is in the first and
+ # last element of the flattened mask
+ if on_value_change_indices[0] != 0:
+ on_value_change_indices = np.insert(on_value_change_indices, 0, 0)
+
+ rle = np.diff(on_value_change_indices)
+
+ if mask[0][0] == 1:
+ rle = np.insert(rle, 0, 0)
+
+ return list(rle)