1114 rle support for coco

supervision/dataset/formats/coco.py

@@ -9,6 +9,7 @@
 from supervision.dataset.utils import (
  approximate_mask_with_polygons,
  map_detections_class_id,
+ rle_to_mask,
 )
 from supervision.detection.core import Detections
 from supervision.detection.utils import polygon_to_mask
@@ -69,6 +70,26 @@ def _polygons_to_masks(
  )
 
 
+def _rles_to_masks(
+ rles: List[np.ndarray], resolution_wh: Tuple[int, int]
+) -> np.ndarray:
+ return np.array(
+ [rle_to_mask(rle=rle, resolution_wh=resolution_wh) for rle in rles],
+ dtype=bool,
+ )
+
+
+def _concatenate_annotation_masks(mask_polygon, mask_rle):
+ if mask_polygon.ndim == 3 and mask_rle.ndim == 3:
+ return np.concatenate((mask_polygon, mask_rle))
+ elif mask_polygon.ndim == 3:
+ return mask_polygon
+ elif mask_rle.ndim == 3:
+ return mask_rle
+ else:
+ None
+
+
 def coco_annotations_to_detections(
  image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool
 ) -> Detections:
@@ -88,10 +109,25 @@ def coco_annotations_to_detections(
  np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2)
  )
  for image_annotation in image_annotations
+ if not image_annotation["iscrowd"]
+ ]
+ mask_polygon = _polygons_to_masks(
+ polygons=polygons, resolution_wh=resolution_wh
+ )
+
+ rles = [
+ np.array(image_annotation["segmentation"]["counts"])
+ for image_annotation in image_annotations
+ if image_annotation["iscrowd"]
  ]
- mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
+ mask_rle = _rles_to_masks(rles=rles, resolution_wh=resolution_wh)
+
  return Detections(
- class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
+ class_id=np.asarray(class_ids, dtype=int),
+ xyxy=xyxy,
+ mask=_concatenate_annotation_masks(
+ mask_polygon=mask_polygon, mask_rle=mask_rle
+ ),
  )
 
  return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
@@ -108,24 +144,26 @@ 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 = []
  if mask is not None:
- polygon = list(
+ 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()
  )
+ # todo: flag for when to use RLE?
+ # segmentation = {"counts": mask_to_rle(binary_mask=mask), "size": list(mask.shape[:2])}
  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] if segmentation else [],
+ "iscrowd": 0, ## todo: iscrowd depends on flag 1 if RLE 0 if polygon
  }
  coco_annotations.append(coco_annotation)
  annotation_id += 1

supervision/dataset/utils.py

@@ -1,6 +1,7 @@
 import copy
 import os
 import random
+from itertools import groupby
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, TypeVar
 
@@ -129,3 +130,23 @@ def train_test_split(
 
  split_index = int(len(data) * train_ratio)
  return data[:split_index], data[split_index:]
+
+
+def rle_to_mask(rle: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
+ width, height = resolution_wh
+
+ zero_one_values = np.zeros_like(rle)
+ zero_one_values[1::2] = 1
+
+ decoded_rle = np.repeat(zero_one_values, rle)
+ return decoded_rle.reshape((height, width), order="F")
+
+
+def mask_to_rle(binary_mask: np.ndarray) -> list:
+ rle = []
+ for _, group in groupby(binary_mask.ravel(order="F")):
+ rle.append(len(list(group)))
+
+ if binary_mask[0][0] == 1:
+ rle = [0] + rle
+ return rle

test/dataset/formats/test_coco.py

@@ -20,14 +20,17 @@ def mock_cock_coco_annotation(
  category_id: int = 0,
  bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
  area: float = 0.0,
+ segmentation: List[list] = None,
+ iscrowd: bool = False,
 ) -> dict:
  return {
  "id": annotation_id,
  "image_id": image_id,
  "category_id": category_id,
  "bbox": list(bbox),
  "area": area,
- "iscrowd": 0,
+ "segmentation": segmentation,
+ "iscrowd": int(iscrowd),
  }
 
 
@@ -226,6 +229,137 @@ def test_group_coco_annotations_by_image_id(
  ),
  DoesNotRaise(),
  ), # two image annotations
+ (
+ [
+ mock_cock_coco_annotation(
+ category_id=0,
+ bbox=(0, 0, 10, 10),
+ area=10 * 10,
+ segmentation=[[0, 0, 4, 0, 4, 5, 9, 5, 9, 9, 0, 9]],
+ )
+ ],
+ (20, 20),
+ True,
+ Detections(
+ xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
+ class_id=np.array([0], dtype=int),
+ mask=np.array(
+ [
+ 0 if i >= 10 or j >= 10 or (i < 5 and j >= 5) else 1
+ for i in range(0, 20)
+ for j in range(0, 20)
+ ]
+ ).reshape((1, 20, 20)),
+ ),
+ DoesNotRaise(),
+ ), # single image annotations with mask, segmentation mask in L-like shape, like below:
+ # 1 0 0 0
+ # 1 1 0 0
+ # 0 0 0 0
+ # 0 0 0 0
+ (
+ [
+ mock_cock_coco_annotation(
+ category_id=0,
+ bbox=(0, 0, 10, 10),
+ area=10 * 10,
+ segmentation={
+ "size": [20, 20],
+ "counts": [
+ 0,
+ 10,
+ 10,
+ 10,
+ 10,
+ 10,
+ 10,
+ 10,
+ 10,
+ 10,
+ 15,
+ 5,
+ 15,
+ 5,
+ 15,
+ 5,
+ 15,
+ 5,
+ 15,
+ 5,
+ 210,
+ ],
+ },
+ iscrowd=True,
+ )
+ ],
+ (20, 20),
+ True,
+ Detections(
+ xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
+ class_id=np.array([0], dtype=int),
+ mask=np.array(
+ [
+ 0 if i >= 10 or j >= 10 or (i < 5 and j >= 5) else 1
+ for i in range(0, 20)
+ for j in range(0, 20)
+ ]
+ ).reshape((1, 20, 20)),
+ ),
+ DoesNotRaise(),
+ ), # single image annotations with mask, RLE segmentation mask in L-like shape, like below:
+ # 1 0 0 0
+ # 1 1 0 0
+ # 0 0 0 0
+ # 0 0 0 0
+ (
+ [
+ mock_cock_coco_annotation(
+ category_id=0,
+ bbox=(0, 0, 10, 10),
+ area=10 * 10,
+ segmentation=[[0, 0, 4, 0, 4, 5, 9, 5, 9, 9, 0, 9]],
+ ),
+ mock_cock_coco_annotation(
+ category_id=0,
+ bbox=(5, 0, 5, 5),
+ area=5 * 5,
+ segmentation={
+ "size": [20, 20],
+ "counts": [100, 5, 15, 5, 15, 5, 15, 5, 15, 5, 215],
+ },
+ iscrowd=True,
+ ),
+ ],
+ (20, 20),
+ True,
+ Detections(
+ xyxy=np.array([[0, 0, 10, 10], [5, 0, 10, 5]], dtype=np.float32),
+ class_id=np.array([0, 0], dtype=int),
+ mask=np.array(
+ [
+ np.array(
+ [
+ 0 if i >= 10 or j >= 10 or (i < 5 and j >= 5) else 1
+ for i in range(0, 20)
+ for j in range(0, 20)
+ ]
+ ).reshape((20, 20)),
+ np.array(
+ [
+ 1 if j > 4 and j < 10 and i < 5 else 0
+ for i in range(0, 20)
+ for j in range(0, 20)
+ ]
+ ).reshape((20, 20)),
+ ]
+ ),
+ ),
+ DoesNotRaise(),
+ ), # two image annotations with mask, one mask as polygon in in L-like shape, second as RLE in shape of square, like below (P = polygon, R = RLE):
+ # P R 0 0
+ # P P 0 0
+ # 0 0 0 0
+ # 0 0 0 0
  ],
 )
 def test_coco_annotations_to_detections(