From 2b0b0cf3fb179b78aa82e98a580523622a933450 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Fri, 21 Aug 2026 17:06:45 -0400 Subject: [PATCH 1/2] Import RLE segmentation masks as outlines DIVE stores geometry rather than rasters, so a decoded mask becomes its contour. Both counts spellings are read, the list form and the LEB128 string pycocotools writes; an undecodable mask still warns. --- docs/DataFormats.md | 8 +- server/dive_utils/serializers/kwcoco.py | 119 +++++++++++++++++-- server/tests/test_deserialize_kwcoco_json.py | 86 +++++++++++++- 3 files changed, 199 insertions(+), 14 deletions(-) diff --git a/docs/DataFormats.md b/docs/DataFormats.md index 2113e367a..49b15cbcc 100644 --- a/docs/DataFormats.md +++ b/docs/DataFormats.md @@ -521,8 +521,12 @@ For COCO files not produced by DIVE: * Partially supported: * COCO has no direct equivalent for DIVE groups, so groups are not represented in COCO export. * Partially supported: - * Run-length encoded segmentations (RLE): bounding boxes and other fields import, - but masks are skipped and a warning is shown. + * Run-length encoded segmentations (RLE): the mask is decoded and imported as its + outline, since DIVE stores geometry rather than rasters. Both COCO counts + spellings are read: a list of run lengths, and the LEB128 string pycocotools + writes. Holes are not representable and are dropped, and a mask that cannot be + decoded is skipped with a warning, as before. Web import only; desktop import + still skips RLE. ### Example COCO Annotation with DIVE Extensions diff --git a/server/dive_utils/serializers/kwcoco.py b/server/dive_utils/serializers/kwcoco.py index dda69e0c7..ff1438ad9 100644 --- a/server/dive_utils/serializers/kwcoco.py +++ b/server/dive_utils/serializers/kwcoco.py @@ -15,10 +15,14 @@ from . import viame RLE_SEGMENTATION_WARNING = ( - 'The COCO file included run-length encoded segmentation masks that are not supported. ' - 'Bounding boxes and other annotation data were imported, but masks were skipped.' + 'The COCO file included run-length encoded segmentation masks that could not be decoded. ' + 'Bounding boxes and other annotation data were imported, but those masks were skipped.' ) +# A mask larger than this is refused rather than allocated; 8K x 8K is already +# far beyond anything DIVE displays. +_RLE_MAX_PIXELS = 64 * 1024 * 1024 + PROB_TOP_K = 10 PROB_EPSILON = 0.001 @@ -165,6 +169,88 @@ def _is_rle_segmentation(annotation: dict, segmentation=None) -> bool: return bool(annotation.get('iscrowd', False)) or isinstance(segmentation, dict) +def _decode_rle_counts(counts) -> Optional[List[int]]: + """Run lengths from either COCO counts spelling. + + Uncompressed COCO writes a list of integers; pycocotools writes the same + runs LEB128-encoded into a string. + """ + if isinstance(counts, (list, tuple)): + if all(isinstance(count, int) and not isinstance(count, bool) and count >= 0 + for count in counts): + return list(counts) + return None + if not isinstance(counts, (str, bytes)): + return None + + text = counts.decode('ascii') if isinstance(counts, bytes) else counts + runs: List[int] = [] + position = 0 + while position < len(text): + value = 0 + shift = 0 + more = True + while more: + if position >= len(text): + return None + char = ord(text[position]) - 48 + value |= (char & 0x1F) << shift + more = bool(char & 0x20) + position += 1 + shift += 5 + # The final chunk carries the sign in bit 0x10 (rleFrString). + if not more and char & 0x10: + value |= -1 << shift + # Runs past the first two are deltas against the run two places back. + if len(runs) > 2: + value += runs[-2] + runs.append(value) + return runs if all(run >= 0 for run in runs) else None + + +def _rle_polygon_coords(segmentation) -> List[List[Tuple[float, float]]]: + """Trace a COCO RLE mask into image-space polygon contours. + + DIVE stores geometry, not rasters, so an imported mask becomes its outline. + Holes are not representable and are dropped. + """ + if not isinstance(segmentation, dict): + return [] + size = segmentation.get('size') + if not (isinstance(size, (list, tuple)) and len(size) == 2): + return [] + height, width = size + if not (isinstance(height, int) and isinstance(width, int)): + return [] + if height <= 0 or width <= 0 or height * width > _RLE_MAX_PIXELS: + return [] + + runs = _decode_rle_counts(segmentation.get('counts')) + if runs is None or sum(runs) != height * width: + return [] + + import cv2 + import numpy as np + + flat = np.zeros(height * width, dtype=np.uint8) + position = 0 + for index, run in enumerate(runs): + if index % 2: # odd runs are foreground + flat[position:position + run] = 1 + position += run + # COCO run-length order is column-major. + mask = flat.reshape((height, width), order='F') + + found = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + coord_lists = [] + for contour in found[-2]: + points = contour.reshape(-1, 2) + if len(points) >= 3: + coord_lists.append([(float(x), float(y)) for x, y in points]) + coord_lists.sort(key=len, reverse=True) + return coord_lists + + def _extract_polygon_coords_lists(segmentation) -> List[List[Tuple[float, float]]]: """Parse COCO / KWCOCO polygon segmentations into coordinate lists.""" if not segmentation or isinstance(segmentation, dict): @@ -201,9 +287,10 @@ def _bbox_from_points(points: List[Tuple[float, float]]) -> List[float]: def _annotation_has_importable_bounds(annotation: dict) -> bool: if _has_valid_bbox(annotation): return True - if _is_rle_segmentation(annotation): - return False - return bool(_extract_polygon_coords_lists(annotation.get('segmentation', []))) + segmentation = annotation.get('segmentation', []) + if _is_rle_segmentation(annotation, segmentation): + return bool(_rle_polygon_coords(segmentation)) + return bool(_extract_polygon_coords_lists(segmentation)) def _missing_bounds_error(annotation_ids: List) -> str: @@ -214,7 +301,7 @@ def _missing_bounds_error(annotation_ids: List) -> str: f'they have no bbox and ' f'no usable polygon segmentation (ids: {shown}{extra}). ' 'Provide bbox [x, y, width, height] or polygon segmentation as [[x1, y1, ...]]. ' - 'Annotations with only RLE segmentation masks still require a bbox.' + 'An RLE mask supplies bounds only when it can be decoded.' ) @@ -222,7 +309,11 @@ def _resolve_coco_bbox(annotation: dict) -> List[float]: if _has_valid_bbox(annotation): return list(annotation['bbox']) - coord_lists = _extract_polygon_coords_lists(annotation.get('segmentation', [])) + segmentation = annotation.get('segmentation', []) + if _is_rle_segmentation(annotation, segmentation): + coord_lists = _rle_polygon_coords(segmentation) + else: + coord_lists = _extract_polygon_coords_lists(segmentation) all_points = [point for coords in coord_lists for point in coords] if all_points: return _bbox_from_points(all_points) @@ -333,10 +424,16 @@ def _parse_annotation( # parse polygons segmentation = annotation.get('segmentation', []) - rle_skipped = _is_rle_segmentation(annotation, segmentation) - - if segmentation and not rle_skipped: - coord_lists = _extract_polygon_coords_lists(segmentation) + rle_skipped = False + + if segmentation: + if _is_rle_segmentation(annotation, segmentation): + coord_lists = _rle_polygon_coords(segmentation) + # Only undecodable masks are reported; a traced one is not a loss + # worth warning about. + rle_skipped = not coord_lists + else: + coord_lists = _extract_polygon_coords_lists(segmentation) if coord_lists: viame.create_geoJSONFeature(features, 'Polygon', coord_lists[0]) diff --git a/server/tests/test_deserialize_kwcoco_json.py b/server/tests/test_deserialize_kwcoco_json.py index 43321e0ca..c7079fe19 100644 --- a/server/tests/test_deserialize_kwcoco_json.py +++ b/server/tests/test_deserialize_kwcoco_json.py @@ -938,7 +938,7 @@ def test_import_missing_bbox_raises_descriptive_error(): kwcoco.load_coco_as_tracks_and_attributes(coco) message = str(exc.value) assert "no bbox and no usable polygon" in message - assert "RLE segmentation masks still require a bbox" in message + assert "An RLE mask supplies bounds only when it can be decoded" in message def test_import_polygon_without_bbox_derives_bounds(): @@ -1239,3 +1239,87 @@ def test_frame_rate_absent_or_unusable(): assert kwcoco.frame_rate_from_coco( _fps_document([{'id': 1, 'annotation_fps': fps}]) ) is None + + +def _rle_to_string(cnts): + """pycocotools rleToString, so the decoder is tested against real output.""" + out = [] + for i, count in enumerate(cnts): + x = int(count) + if i > 2: + x -= int(cnts[i - 2]) + more = True + while more: + chunk = x & 0x1F + x >>= 5 + more = (x != -1) if (chunk & 0x10) else (x != 0) + if more: + chunk |= 0x20 + out.append(chr(chunk + 48)) + return ''.join(out) + + +def _square_mask_runs(): + """Column-major run lengths for a 6x6 square at (3, 2) in a 10x10 mask.""" + runs, current, length = [], 0, 0 + for column in range(10): + for row in range(10): + value = 1 if (2 <= row < 8 and 3 <= column < 9) else 0 + if value == current: + length += 1 + else: + runs.append(length) + current = value + length = 1 + runs.append(length) + return runs + + +def _rle_document(counts): + return { + 'images': [{'id': 1, 'file_name': 'frame_000000.png', 'frame_index': 0}], + 'annotations': [{ + 'id': 1, 'image_id': 1, 'category_id': 1, 'track_id': 1, + 'segmentation': {'counts': counts, 'size': [10, 10]}, 'iscrowd': 1, + }], + 'categories': [{'id': 1, 'name': 'fish'}], + } + + +@pytest.mark.parametrize('as_string', [False, True]) +def test_rle_masks_import_as_outlines(as_string): + """DIVE stores geometry, so a decoded mask arrives as its outline.""" + runs = _square_mask_runs() + counts = _rle_to_string(runs) if as_string else runs + tracks, _, warnings, _ = kwcoco.load_coco_as_tracks_and_attributes(_rle_document(counts)) + + feature = tracks['tracks']['1']['features'][0] + polygon = [ + geometry for geometry in feature['geometry']['features'] + if geometry['geometry']['type'] == 'Polygon' + ] + assert polygon, 'expected a polygon traced from the mask' + coords = polygon[0]['geometry']['coordinates'][0] + xs = [point[0] for point in coords] + ys = [point[1] for point in coords] + assert (min(xs), max(xs), min(ys), max(ys)) == (3, 8, 2, 7) + # The mask carried no bbox, so it supplied the bounds itself. + assert feature['bounds'] == [3, 2, 8, 7] + assert kwcoco.RLE_SEGMENTATION_WARNING not in warnings + + +def test_undecodable_rle_still_warns(): + """Run lengths that do not fill the mask are reported, not guessed at.""" + document = _rle_document([5]) + document['annotations'][0]['bbox'] = [0, 0, 4, 4] + tracks, _, warnings, _ = kwcoco.load_coco_as_tracks_and_attributes(document) + + assert kwcoco.RLE_SEGMENTATION_WARNING in warnings + assert tracks['tracks']['1']['features'][0]['bounds'] == [0, 0, 4, 4] + + +def test_decode_rle_counts_rejects_junk(): + assert kwcoco._decode_rle_counts([1, -2]) is None + assert kwcoco._decode_rle_counts([1, 'x']) is None + assert kwcoco._decode_rle_counts(None) is None + assert kwcoco._decode_rle_counts(_rle_to_string([4, 2, 4])) == [4, 2, 4] From 78d3cdffb16be0fa0744942538f848202b4c7f51 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Fri, 21 Aug 2026 17:13:28 -0400 Subject: [PATCH 2/2] Trace mask outlines with numpy, not opencv opencv is only a dev dependency, so the deployed server has numpy alone. Moore-neighbour tracing walks the boundary instead, which costs the perimeter rather than the area. --- docs/DataFormats.md | 3 +- server/dive_utils/serializers/kwcoco.py | 89 ++++++++++++++++++++++--- 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/docs/DataFormats.md b/docs/DataFormats.md index 49b15cbcc..974c154e8 100644 --- a/docs/DataFormats.md +++ b/docs/DataFormats.md @@ -526,7 +526,8 @@ For COCO files not produced by DIVE: spellings are read: a list of run lengths, and the LEB128 string pycocotools writes. Holes are not representable and are dropped, and a mask that cannot be decoded is skipped with a warning, as before. Web import only; desktop import - still skips RLE. + still skips RLE. Decoding needs no extra dependency: the outline is traced with + numpy alone. ### Example COCO Annotation with DIVE Extensions diff --git a/server/dive_utils/serializers/kwcoco.py b/server/dive_utils/serializers/kwcoco.py index ff1438ad9..e1ff73f37 100644 --- a/server/dive_utils/serializers/kwcoco.py +++ b/server/dive_utils/serializers/kwcoco.py @@ -208,6 +208,64 @@ def _decode_rle_counts(counts) -> Optional[List[int]]: return runs if all(run >= 0 for run in runs) else None +# Clockwise Moore neighbourhood, as (dx, dy) starting from due east. +_MOORE_OFFSETS = ( + (1, 0), (1, 1), (0, 1), (-1, 1), (-1, 0), (-1, -1), (0, -1), (1, -1), +) + + +def _trace_contour(mask, start, visited) -> List[Tuple[float, float]]: + """Moore-neighbour trace of one component's outer boundary. + + Pure numpy: the server has numpy but not opencv, and walking the boundary + costs the perimeter rather than the area. + """ + height, width = mask.shape + contour = [start] + visited[start[1], start[0]] = True + # Entering the start pixel from the west, so begin the search north of it. + previous = (start[0] - 1, start[1]) + current = start + + while True: + back = (previous[0] - current[0], previous[1] - current[1]) + try: + index = _MOORE_OFFSETS.index(back) + except ValueError: + index = 0 + found = None + for step in range(1, 9): + offset = _MOORE_OFFSETS[(index + step) % 8] + candidate = (current[0] + offset[0], current[1] + offset[1]) + if not (0 <= candidate[0] < width and 0 <= candidate[1] < height): + continue + if mask[candidate[1], candidate[0]]: + found = candidate + break + previous = candidate + if found is None: # isolated pixel + break + if found == start and len(contour) > 1: + break + contour.append(found) + visited[found[1], found[0]] = True + previous = current + current = found + if len(contour) > 4 * height * width: # cannot happen; refuses to spin + break + + return [(float(x), float(y)) for x, y in contour] + + +def _polygon_area(points: List[Tuple[float, float]]) -> float: + """Shoelace area of a closed contour.""" + total = 0.0 + for index, (x, y) in enumerate(points): + next_x, next_y = points[(index + 1) % len(points)] + total += x * next_y - next_x * y + return abs(total) / 2.0 + + def _rle_polygon_coords(segmentation) -> List[List[Tuple[float, float]]]: """Trace a COCO RLE mask into image-space polygon contours. @@ -229,25 +287,38 @@ def _rle_polygon_coords(segmentation) -> List[List[Tuple[float, float]]]: if runs is None or sum(runs) != height * width: return [] - import cv2 import numpy as np - flat = np.zeros(height * width, dtype=np.uint8) + flat = np.zeros(height * width, dtype=bool) position = 0 for index, run in enumerate(runs): if index % 2: # odd runs are foreground - flat[position:position + run] = 1 + flat[position:position + run] = True position += run # COCO run-length order is column-major. mask = flat.reshape((height, width), order='F') + if not mask.any(): + return [] + + # A boundary pixel is foreground with at least one background 4-neighbour. + padded = np.zeros((height + 2, width + 2), dtype=bool) + padded[1:-1, 1:-1] = mask + interior = ( + padded[:-2, 1:-1] & padded[2:, 1:-1] & padded[1:-1, :-2] & padded[1:-1, 2:] + ) + boundary = mask & ~interior - found = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + visited = np.zeros_like(mask) coord_lists = [] - for contour in found[-2]: - points = contour.reshape(-1, 2) - if len(points) >= 3: - coord_lists.append([(float(x), float(y)) for x, y in points]) - coord_lists.sort(key=len, reverse=True) + for y, x in zip(*np.nonzero(boundary)): + if visited[y, x]: + continue + contour = _trace_contour(mask, (int(x), int(y)), visited) + if len(contour) >= 3: + coord_lists.append(contour) + # Largest by enclosed area, not by point count: a long thin outline can + # carry more points than a bigger blob, and callers take the first. + coord_lists.sort(key=_polygon_area, reverse=True) return coord_lists