From 48035c14e2de8fc6f5f3d7a1bc633f50217eda09 Mon Sep 17 00:00:00 2001 From: Ishaan <103370340+0xD4rky@users.noreply.github.com> Date: Tue, 31 Mar 2026 15:21:30 +0530 Subject: [PATCH] fixed precision converting annotations with `"force_mask=True"` (#1746) * changing polygon format conversion * reformating polygon type conversion * handling mask dimensions * fix: resolve unresolved PR #1746 review comments * refactor(tests): restructure YOLO polygon mask precision tests with fixtures --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: jirka <6035284+Borda@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Codex --- src/supervision/dataset/formats/yolo.py | 7 +- tests/dataset/formats/test_yolo.py | 125 ++++++++++++++++++++++++ 2 files changed, 130 insertions(+), 2 deletions(-) diff --git a/src/supervision/dataset/formats/yolo.py b/src/supervision/dataset/formats/yolo.py index 0d78773f..9abaff45 100644 --- a/src/supervision/dataset/formats/yolo.py +++ b/src/supervision/dataset/formats/yolo.py @@ -53,7 +53,10 @@ def _polygons_to_masks( ) -> npt.NDArray[np.bool_]: return np.array( [ - polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh) + polygon_to_mask( + polygon=np.round(polygon).astype(np.int32), + resolution_wh=resolution_wh, + ) for polygon in polygons ], dtype=bool, @@ -129,7 +132,7 @@ def yolo_annotations_to_detections( return Detections(class_id=class_id, xyxy=xyxy, data=data) polygons = [ - np.round(polygon * np.array(resolution_wh, dtype=np.float32)).astype(int) + polygon * np.array(resolution_wh, dtype=np.float32) for polygon in relative_polygon ] mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh) diff --git a/tests/dataset/formats/test_yolo.py b/tests/dataset/formats/test_yolo.py index 5ca52087..86047dc9 100644 --- a/tests/dataset/formats/test_yolo.py +++ b/tests/dataset/formats/test_yolo.py @@ -3,6 +3,7 @@ from __future__ import annotations import os import tempfile from contextlib import ExitStack as DoesNotRaise +from pathlib import Path import numpy as np import pytest @@ -413,3 +414,127 @@ def test_load_yolo_annotations_segmentation_produces_masks() -> None: assert detection.mask is not None, ( "Segmentation annotations with is_obb=False must produce mask arrays" ) + + +def test_polygons_to_masks_multiple_polygons_shape() -> None: + """Regression test for #1746: _polygons_to_masks must return shape (N, H, W). + + The original PR rewrite processed only a single polygon and always returned + shape (1, H, W), breaking multi-polygon detections. + """ + from supervision.dataset.formats.yolo import _polygons_to_masks + + resolution_wh = (100, 100) + # Fractional pixel coords ensure the rounding path inside the function is exercised + polygon_a = np.array( + [[10.5, 20.5], [10.5, 50.5], [40.5, 50.5], [40.5, 20.5]], dtype=np.float32 + ) + polygon_b = np.array( + [[60.3, 30.7], [60.3, 70.3], [90.3, 70.3], [90.3, 30.7]], dtype=np.float32 + ) + + masks = _polygons_to_masks( + polygons=[polygon_a, polygon_b], resolution_wh=resolution_wh + ) + + assert masks.shape == (2, 100, 100), f"Expected (2, 100, 100), got {masks.shape}" + assert masks.dtype == np.bool_ + assert masks[0].any(), "Polygon A produced an empty mask" + assert masks[1].any(), "Polygon B produced an empty mask" + assert not np.any(masks[0] & masks[1]), ( + "Non-overlapping polygons produced overlapping masks" + ) + + +@pytest.fixture +def yolo_mask_round_trip_sample( + tmp_path: Path, +) -> tuple[str, str, str, tuple[int, int], str]: + """Create a minimal YOLO segmentation sample for round-trip mask tests.""" + images_dir = tmp_path / "images" + labels_dir = tmp_path / "labels" + images_dir.mkdir() + labels_dir.mkdir() + + # Odd resolution ensures coord * dim is non-integer (e.g. 0.25 * 101 = 25.25) + resolution_wh = (101, 97) + Image.new("RGB", resolution_wh).save(images_dir / "test.jpg") + + original_line = "0 0.25000 0.40000 0.25000 0.60000 0.45000 0.60000 0.45000 0.40000" + (labels_dir / "test.txt").write_text(original_line + "\n") + + data_yaml_path = tmp_path / "data.yaml" + data_yaml_path.write_text("names: ['class0']\n") + + return ( + str(images_dir), + str(labels_dir), + str(data_yaml_path), + resolution_wh, + original_line, + ) + + +def test_yolo_polygon_mask_precision_no_coord_drift_loads_mask( + yolo_mask_round_trip_sample: tuple[str, str, str, tuple[int, int], str], +) -> None: + """YOLO load with force_masks=True should produce a non-empty mask.""" + images_dir, labels_dir, data_yaml_path, _, _ = yolo_mask_round_trip_sample + + _, _, annotations = load_yolo_annotations( + images_directory_path=images_dir, + annotations_directory_path=labels_dir, + data_yaml_path=data_yaml_path, + force_masks=True, + ) + + assert len(annotations) == 1 + detection = next(iter(annotations.values())) + assert detection.mask is not None + assert detection.mask.shape[0] == 1 + assert detection.mask[0].any() + + +def test_yolo_polygon_mask_precision_no_coord_drift_round_trip_iou( + yolo_mask_round_trip_sample: tuple[str, str, str, tuple[int, int], str], +) -> None: + """YOLO load/save round-trip should keep segmentation mask geometry stable.""" + images_dir, labels_dir, data_yaml_path, resolution_wh, original_line = ( + yolo_mask_round_trip_sample + ) + + _, _, annotations = load_yolo_annotations( + images_directory_path=images_dir, + annotations_directory_path=labels_dir, + data_yaml_path=data_yaml_path, + force_masks=True, + ) + detection = next(iter(annotations.values())) + + image_arr = np.zeros((resolution_wh[1], resolution_wh[0], 3), dtype=np.uint8) + saved_lines = detections_to_yolo_annotations( + detections=detection, image_shape=image_arr.shape + ) + + assert len(saved_lines) == 1 + original_detection = yolo_annotations_to_detections( + lines=[original_line], resolution_wh=resolution_wh, with_masks=True + ) + saved_detection = yolo_annotations_to_detections( + lines=saved_lines, resolution_wh=resolution_wh, with_masks=True + ) + + assert original_detection.mask is not None + assert saved_detection.mask is not None + + original_mask = original_detection.mask[0] + saved_mask = saved_detection.mask[0] + intersection = np.logical_and(original_mask, saved_mask).sum() + union = np.logical_or(original_mask, saved_mask).sum() + assert union > 0 + # Keep polygon round-trip drift bounded while avoiding vertex-order assumptions. + iou = intersection / union + assert iou > 0.95, ( + f"Mask IoU {iou:.6f} too low after YOLO load/save round-trip — " + "precision regression in polygon mask conversion" + )