diff --git a/workflows/roboflow_workflows/core_steps/visualizations/common/utils.py b/workflows/roboflow_workflows/core_steps/visualizations/common/utils.py index 0c012ac0f2..30f4c3daa8 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/common/utils.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/common/utils.py @@ -1,4 +1,9 @@ +import copy + +import numpy as np +import pycocotools.mask as mask_utils import supervision as sv +from roboflow_workflows.execution_engine.constants import RLE_MASK_KEY_IN_SV_DETECTIONS UNKNOWN_CLASS_COLOR = sv.Color.GREY UNKNOWN_CLASS_COLOR_RGB = UNKNOWN_CLASS_COLOR.as_rgb() @@ -37,6 +42,27 @@ def by_idx(self, idx: int) -> sv.Color: return self.colors[idx % len(self.colors)] +def ensure_dense_masks(detections: sv.Detections) -> sv.Detections: + """Decode RLE side-channel masks so Supervision annotators can read ``mask``. + + RLE-first blocks return ``mask=None`` with COCO RLE in ``data["rle_mask"]``. + Returns a shallow copy with ``mask`` set; the input is unchanged. + """ + if ( + detections.mask is not None + or RLE_MASK_KEY_IN_SV_DETECTIONS not in detections.data + ): + return detections + detections = copy.copy(detections) + detections.mask = np.array( + [ + mask_utils.decode(rle).astype(bool) + for rle in detections.data[RLE_MASK_KEY_IN_SV_DETECTIONS] + ] + ) + return detections + + def wrap_color_palette(palette: sv.ColorPalette) -> NegativeSafeColorPalette: """Wrap a palette so negative ids resolve to gray instead of raising. diff --git a/workflows/roboflow_workflows/core_steps/visualizations/halo/v1.py b/workflows/roboflow_workflows/core_steps/visualizations/halo/v1.py index bbc1ed85f5..8f55b60b49 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/halo/v1.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/halo/v1.py @@ -10,6 +10,7 @@ ColorableVisualizationBlock, ColorableVisualizationManifest, ) +from roboflow_workflows.core_steps.visualizations.common.utils import ensure_dense_masks from roboflow_workflows.execution_engine.entities.base import WorkflowImageData from roboflow_workflows.execution_engine.entities.types import ( FLOAT_ZERO_TO_ONE_KIND, @@ -177,6 +178,7 @@ def run( opacity, kernel_size, ) + predictions = ensure_dense_masks(predictions) scene = image.numpy_image if copy_image: scene = scene.copy() diff --git a/workflows/roboflow_workflows/core_steps/visualizations/halo/v2.py b/workflows/roboflow_workflows/core_steps/visualizations/halo/v2.py index 39aa1bb9bf..aa9999c89d 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/halo/v2.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/halo/v2.py @@ -7,6 +7,7 @@ ColorableVisualizationBlock, ColorableVisualizationManifest, ) +from roboflow_workflows.core_steps.visualizations.common.utils import ensure_dense_masks from roboflow_workflows.execution_engine.entities.base import WorkflowImageData from roboflow_workflows.execution_engine.entities.types import ( FLOAT_ZERO_TO_ONE_KIND, @@ -177,6 +178,7 @@ def run( opacity, kernel_size, ) + predictions = ensure_dense_masks(predictions) scene = image.numpy_image if copy_image: scene = scene.copy() diff --git a/workflows/roboflow_workflows/core_steps/visualizations/label/v1.py b/workflows/roboflow_workflows/core_steps/visualizations/label/v1.py index fb151c70fe..a0d01a58c8 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/label/v1.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/label/v1.py @@ -10,7 +10,10 @@ from roboflow_workflows.core_steps.visualizations.common.label_text import ( build_detection_labels, ) -from roboflow_workflows.core_steps.visualizations.common.utils import str_to_color +from roboflow_workflows.core_steps.visualizations.common.utils import ( + ensure_dense_masks, + str_to_color, +) from roboflow_workflows.execution_engine.entities.base import WorkflowImageData from roboflow_workflows.execution_engine.entities.types import ( FLOAT_KIND, @@ -262,6 +265,8 @@ def run( text_padding, border_radius, ) + if text_position == "CENTER_OF_MASS": + predictions = ensure_dense_masks(predictions) labels = build_detection_labels(predictions, text) scene = image.numpy_image if copy_image: diff --git a/workflows/roboflow_workflows/core_steps/visualizations/label/v2.py b/workflows/roboflow_workflows/core_steps/visualizations/label/v2.py index 0007cef5d1..b05707fa4c 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/label/v2.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/label/v2.py @@ -13,7 +13,10 @@ build_detection_labels, compute_adaptive_label_text_scale, ) -from roboflow_workflows.core_steps.visualizations.common.utils import str_to_color +from roboflow_workflows.core_steps.visualizations.common.utils import ( + ensure_dense_masks, + str_to_color, +) from roboflow_workflows.core_steps.visualizations.label.v1 import ( LONG_DESCRIPTION, SHORT_DESCRIPTION, @@ -255,6 +258,8 @@ def run( text_padding, border_radius, ) + if text_position == "CENTER_OF_MASS": + predictions = ensure_dense_masks(predictions) labels = build_detection_labels(predictions, text) annotated_image = annotator.annotate( scene=image.numpy_image.copy() if copy_image else image.numpy_image, diff --git a/workflows/roboflow_workflows/core_steps/visualizations/mask/v1.py b/workflows/roboflow_workflows/core_steps/visualizations/mask/v1.py index d19d16be1b..d4a43fd3aa 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/mask/v1.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/mask/v1.py @@ -1,8 +1,5 @@ -import copy from typing import List, Literal, Optional, Type, Union -import numpy as np -import pycocotools.mask as mask_utils import supervision as sv from pydantic import ConfigDict, Field from roboflow_workflows.core_steps.visualizations.common.base import OUTPUT_IMAGE_KEY @@ -10,6 +7,7 @@ ColorableVisualizationBlock, ColorableVisualizationManifest, ) +from roboflow_workflows.core_steps.visualizations.common.utils import ensure_dense_masks from roboflow_workflows.execution_engine.constants import RLE_MASK_KEY_IN_SV_DETECTIONS from roboflow_workflows.execution_engine.entities.base import WorkflowImageData from roboflow_workflows.execution_engine.entities.types import ( @@ -169,17 +167,7 @@ def run( opacity, ) - if ( - predictions.mask is None - and RLE_MASK_KEY_IN_SV_DETECTIONS in predictions.data - ): - predictions = copy.copy(predictions) - predictions.mask = np.array( - [ - mask_utils.decode(rle).astype(bool) - for rle in predictions.data[RLE_MASK_KEY_IN_SV_DETECTIONS] - ] - ) + predictions = ensure_dense_masks(predictions) scene = image.numpy_image if copy_image: diff --git a/workflows/roboflow_workflows/core_steps/visualizations/polygon/v1.py b/workflows/roboflow_workflows/core_steps/visualizations/polygon/v1.py index 12e8dd5b46..368d159570 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/polygon/v1.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/polygon/v1.py @@ -10,6 +10,7 @@ ColorableVisualizationBlock, ColorableVisualizationManifest, ) +from roboflow_workflows.core_steps.visualizations.common.utils import ensure_dense_masks from roboflow_workflows.execution_engine.entities.base import WorkflowImageData from roboflow_workflows.execution_engine.entities.types import ( INSTANCE_SEGMENTATION_PREDICTION_KIND, @@ -165,6 +166,7 @@ def run( thickness, ) + predictions = ensure_dense_masks(predictions) scene = image.numpy_image if copy_image: scene = scene.copy() diff --git a/workflows/roboflow_workflows/core_steps/visualizations/polygon/v2.py b/workflows/roboflow_workflows/core_steps/visualizations/polygon/v2.py index e3ec164ec1..afee33bddc 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/polygon/v2.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/polygon/v2.py @@ -7,6 +7,7 @@ ColorableVisualizationBlock, ColorableVisualizationManifest, ) +from roboflow_workflows.core_steps.visualizations.common.utils import ensure_dense_masks from roboflow_workflows.execution_engine.entities.base import WorkflowImageData from roboflow_workflows.execution_engine.entities.types import ( INSTANCE_SEGMENTATION_PREDICTION_KIND, @@ -164,6 +165,7 @@ def run( thickness, ) + predictions = ensure_dense_masks(predictions) scene = image.numpy_image if copy_image: scene = scene.copy() diff --git a/workflows/roboflow_workflows/core_steps/visualizations/rich_label/v1.py b/workflows/roboflow_workflows/core_steps/visualizations/rich_label/v1.py index 44ba80957b..2584e0f9b0 100644 --- a/workflows/roboflow_workflows/core_steps/visualizations/rich_label/v1.py +++ b/workflows/roboflow_workflows/core_steps/visualizations/rich_label/v1.py @@ -20,7 +20,10 @@ build_detection_labels, compute_adaptive_rich_font_size, ) -from roboflow_workflows.core_steps.visualizations.common.utils import str_to_color +from roboflow_workflows.core_steps.visualizations.common.utils import ( + ensure_dense_masks, + str_to_color, +) from roboflow_workflows.execution_engine.entities.base import WorkflowImageData from roboflow_workflows.execution_engine.entities.types import ( INTEGER_KIND, @@ -350,6 +353,8 @@ def run( border_radius, max_line_length, ) + if text_position == "CENTER_OF_MASS": + predictions = ensure_dense_masks(predictions) labels = build_detection_labels(predictions, text) annotated_image = annotator.annotate( diff --git a/workflows/tests/unit_tests/core_steps/visualizations/test_halo.py b/workflows/tests/unit_tests/core_steps/visualizations/test_halo.py index 735777ded0..41785ce3d0 100644 --- a/workflows/tests/unit_tests/core_steps/visualizations/test_halo.py +++ b/workflows/tests/unit_tests/core_steps/visualizations/test_halo.py @@ -1,4 +1,5 @@ import numpy as np +import pycocotools.mask as mask_utils import pytest import supervision as sv from pydantic import ValidationError @@ -146,3 +147,43 @@ def test_halo_visualization_block_v2() -> None: assert not np.array_equal( output.get("image").numpy_image, np.zeros((1000, 1000, 3), dtype=np.uint8) ) + + +def _rle_square_detections() -> sv.Detections: + mask = np.zeros((100, 100), dtype=np.uint8) + mask[20:60, 20:60] = 1 + rle = mask_utils.encode(np.asfortranarray(mask)) + rle["counts"] = rle["counts"].decode("utf-8") + return sv.Detections( + xyxy=np.array([[20, 20, 60, 60]], dtype=np.float64), + mask=None, + class_id=np.array([1]), + data={ + "class_name": np.array(["cat"]), + "rle_mask": np.array([rle], dtype=object), + }, + ) + + +@pytest.mark.parametrize( + "block_cls", [HaloVisualizationBlockV1, HaloVisualizationBlockV2] +) +def test_halo_visualization_renders_rle_masks(block_cls) -> None: + output = block_cls().run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + numpy_image=np.zeros((100, 100, 3), dtype=np.uint8), + ), + predictions=_rle_square_detections(), + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=[], + color_axis="CLASS", + opacity=0.8, + kernel_size=5, + ) + + result = output["image"].numpy_image + assert result.shape == (100, 100, 3) + assert not np.array_equal(result, np.zeros((100, 100, 3), dtype=np.uint8)) diff --git a/workflows/tests/unit_tests/core_steps/visualizations/test_label.py b/workflows/tests/unit_tests/core_steps/visualizations/test_label.py index a0e03143b7..6428a5b6cb 100644 --- a/workflows/tests/unit_tests/core_steps/visualizations/test_label.py +++ b/workflows/tests/unit_tests/core_steps/visualizations/test_label.py @@ -1,4 +1,5 @@ import numpy as np +import pycocotools.mask as mask_utils import pytest import supervision as sv from pydantic import ValidationError @@ -270,3 +271,43 @@ def test_label_visualization_block_with_area_converted_missing_data() -> None: assert "image" in output assert hasattr(output.get("image"), "numpy_image") assert output.get("image").numpy_image.shape == (1000, 1000, 3) + + +def test_label_visualization_center_of_mass_with_rle_masks() -> None: + mask = np.zeros((100, 100), dtype=np.uint8) + mask[20:60, 20:60] = 1 + rle = mask_utils.encode(np.asfortranarray(mask)) + rle["counts"] = rle["counts"].decode("utf-8") + detections = sv.Detections( + xyxy=np.array([[20, 20, 60, 60]], dtype=np.float64), + mask=None, + class_id=np.array([1]), + data={ + "class_name": np.array(["cat"]), + "rle_mask": np.array([rle], dtype=object), + }, + ) + + output = LabelVisualizationBlockV1().run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + numpy_image=np.zeros((100, 100, 3), dtype=np.uint8), + ), + predictions=detections, + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + text="Class", + text_position="CENTER_OF_MASS", + text_color="WHITE", + text_scale=1.0, + text_thickness=1, + text_padding=10, + border_radius=0, + ) + + result = output["image"].numpy_image + assert result.shape == (100, 100, 3) + assert not np.array_equal(result, np.zeros((100, 100, 3), dtype=np.uint8)) diff --git a/workflows/tests/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py b/workflows/tests/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py index 9fbc5dfc64..8ec204da4b 100644 --- a/workflows/tests/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py +++ b/workflows/tests/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py @@ -1,4 +1,5 @@ import numpy as np +import pycocotools.mask as mask_utils import pytest import supervision as sv import torch @@ -277,3 +278,44 @@ def test_label_v2_tensor_native_empty_predictions_passthrough_is_device_resident assert ( shared_output["image"].tensor_image.data_ptr() == image.tensor_image.data_ptr() ) + + +def test_label_v2_center_of_mass_with_rle_masks() -> None: + mask = np.zeros((100, 100), dtype=np.uint8) + mask[20:60, 20:60] = 1 + rle = mask_utils.encode(np.asfortranarray(mask)) + rle["counts"] = rle["counts"].decode("utf-8") + detections = sv.Detections( + xyxy=np.array([[20, 20, 60, 60]], dtype=np.float64), + mask=None, + class_id=np.array([1]), + data={ + "class_name": np.array(["cat"]), + "rle_mask": np.array([rle], dtype=object), + }, + ) + + output = LabelVisualizationBlockV2().run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + numpy_image=np.zeros((100, 100, 3), dtype=np.uint8), + ), + predictions=detections, + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + text="Class", + text_position="CENTER_OF_MASS", + text_color="WHITE", + text_size_mode="Manual", + text_scale=1.0, + text_thickness=1, + text_padding=10, + border_radius=0, + ) + + result = output["image"].numpy_image + assert result.shape == (100, 100, 3) + assert not np.array_equal(result, np.zeros((100, 100, 3), dtype=np.uint8)) diff --git a/workflows/tests/unit_tests/core_steps/visualizations/test_polygon.py b/workflows/tests/unit_tests/core_steps/visualizations/test_polygon.py index 15d3a6f847..1fcede00a1 100644 --- a/workflows/tests/unit_tests/core_steps/visualizations/test_polygon.py +++ b/workflows/tests/unit_tests/core_steps/visualizations/test_polygon.py @@ -1,4 +1,5 @@ import numpy as np +import pycocotools.mask as mask_utils import pytest import supervision as sv from pydantic import ValidationError @@ -187,3 +188,64 @@ def test_polygon_annotator_draws_at_correct_coordinates() -> None: assert ( not origin_region.any() ), "Polygon was drawn near origin — coordinate offset is broken" + + +def _rle_circle_detections() -> sv.Detections: + height, width, cx, cy, radius = 100, 100, 50, 50, 15 + yy, xx = np.ogrid[:height, :width] + mask = (xx - cx) ** 2 + (yy - cy) ** 2 <= radius**2 + rle = mask_utils.encode(np.asfortranarray(mask.astype(np.uint8))) + rle["counts"] = rle["counts"].decode("utf-8") + return sv.Detections( + xyxy=np.array([[35, 35, 65, 65]], dtype=np.float64), + mask=None, + class_id=np.array([1]), + data={ + "class_name": np.array(["cat"]), + "rle_mask": np.array([rle], dtype=object), + }, + ) + + +def _interior_contour_pixel(mask: np.ndarray, xyxy) -> tuple[int, int]: + x1, y1, x2, y2 = (int(v) for v in xyxy) + height, width = mask.shape + for y in range(1, height - 1): + for x in range(1, width - 1): + if not mask[y, x]: + continue + if x in (x1, x2) or y in (y1, y2): + continue + if not ( + mask[y - 1, x] and mask[y + 1, x] and mask[y, x - 1] and mask[y, x + 1] + ): + return x, y + raise AssertionError("expected a contour pixel inside the box") + + +@pytest.mark.parametrize( + "block_cls", [PolygonVisualizationBlockV1, PolygonVisualizationBlockV2] +) +def test_polygon_visualization_draws_rle_contour_not_box(block_cls) -> None: + detections = _rle_circle_detections() + yy, xx = np.ogrid[:100, :100] + circle = (xx - 50) ** 2 + (yy - 50) ** 2 <= 15**2 + contour_x, contour_y = _interior_contour_pixel(circle, detections.xyxy[0]) + + output = block_cls().run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + numpy_image=np.zeros((100, 100, 3), dtype=np.uint8), + ), + predictions=detections, + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=[], + color_axis="CLASS", + thickness=1, + ) + + result = output["image"].numpy_image + assert result[contour_y, contour_x].any(), "expected the mask contour to be drawn" + assert not result[35, 35].any(), "box corner should stay empty for an RLE circle" diff --git a/workflows/tests/unit_tests/core_steps/visualizations/test_rich_label.py b/workflows/tests/unit_tests/core_steps/visualizations/test_rich_label.py index 0ed5995ed3..64677242ca 100644 --- a/workflows/tests/unit_tests/core_steps/visualizations/test_rich_label.py +++ b/workflows/tests/unit_tests/core_steps/visualizations/test_rich_label.py @@ -1,4 +1,5 @@ import numpy as np +import pycocotools.mask as mask_utils import pytest import supervision as sv import torch @@ -226,6 +227,32 @@ def test_rich_label_visualization_block(bundled_fonts) -> None: ), "Image should be modified by label rendering" +def test_rich_label_center_of_mass_with_rle_masks(bundled_fonts) -> None: + mask = np.zeros((400, 400), dtype=np.uint8) + mask[40:160, 40:160] = 1 + rle = mask_utils.encode(np.asfortranarray(mask)) + rle["counts"] = rle["counts"].decode("utf-8") + detections = sv.Detections( + xyxy=np.array([[40, 40, 160, 160]], dtype=np.float64), + mask=None, + class_id=np.array([0]), + data={ + "class_name": np.array(["cat"]), + "rle_mask": np.array([rle], dtype=object), + }, + ) + + output = _run_block( + RichLabelVisualizationBlockV1(), + predictions=detections, + text_position="CENTER_OF_MASS", + ) + + result = output["image"].numpy_image + assert result.shape == (400, 400, 3) + assert not np.array_equal(result, np.zeros((400, 400, 3), dtype=np.uint8)) + + def test_rich_label_visualization_block_raises_on_unknown_font_delivered_at_runtime() -> ( None ): diff --git a/workflows/tests/unit_tests/core_steps/visualizations/test_str_to_color.py b/workflows/tests/unit_tests/core_steps/visualizations/test_str_to_color.py index 2aade1c25b..977ef33f04 100644 --- a/workflows/tests/unit_tests/core_steps/visualizations/test_str_to_color.py +++ b/workflows/tests/unit_tests/core_steps/visualizations/test_str_to_color.py @@ -1,6 +1,11 @@ +import numpy as np +import pycocotools.mask as mask_utils import pytest import supervision as sv -from roboflow_workflows.core_steps.visualizations.common.utils import str_to_color +from roboflow_workflows.core_steps.visualizations.common.utils import ( + ensure_dense_masks, + str_to_color, +) def test_str_to_color_with_hex_color() -> None: @@ -56,3 +61,57 @@ def test_str_to_color_with_invalid_color() -> None: # when with pytest.raises(ValueError): _ = str_to_color(color) + + +def _rle_from_mask(mask: np.ndarray) -> dict: + rle = mask_utils.encode(np.asfortranarray(mask.astype(np.uint8))) + rle["counts"] = rle["counts"].decode("utf-8") + return rle + + +def test_ensure_dense_masks_returns_input_when_mask_already_set() -> None: + mask = np.zeros((1, 8, 8), dtype=bool) + mask[0, 2:6, 2:6] = True + detections = sv.Detections( + xyxy=np.array([[2, 2, 6, 6]], dtype=np.float64), + mask=mask, + class_id=np.array([1]), + ) + + result = ensure_dense_masks(detections) + + assert result is detections + assert result.mask is mask + + +def test_ensure_dense_masks_returns_input_when_no_rle_key() -> None: + detections = sv.Detections( + xyxy=np.array([[2, 2, 6, 6]], dtype=np.float64), + mask=None, + class_id=np.array([1]), + ) + + result = ensure_dense_masks(detections) + + assert result is detections + assert result.mask is None + + +def test_ensure_dense_masks_decodes_rle_without_mutating_input() -> None: + mask = np.zeros((16, 16), dtype=bool) + mask[4:12, 4:12] = True + detections = sv.Detections( + xyxy=np.array([[4, 4, 12, 12]], dtype=np.float64), + mask=None, + class_id=np.array([1]), + data={"rle_mask": np.array([_rle_from_mask(mask)], dtype=object)}, + ) + + result = ensure_dense_masks(detections) + + assert detections.mask is None + assert result is not detections + assert result.mask is not None + assert result.mask.shape == (1, 16, 16) + assert result.mask.dtype == bool + assert np.array_equal(result.mask[0], mask)