From 7c59f80961a88bc0fac1db0269e585bf5ac4e29b Mon Sep 17 00:00:00 2001 From: Albert van Houten Date: Thu, 28 May 2026 16:04:13 +0200 Subject: [PATCH 1/5] refactor(mask_converters): optimize polygon rasterization and dtype handling Signed-off-by: Albert van Houten --- .../converters/mask_converters.py | 87 ++++++++---------- .../converters/test_mask_converters.py | 88 +++++++++++++++++++ 2 files changed, 124 insertions(+), 51 deletions(-) diff --git a/src/datumaro/experimental/converters/mask_converters.py b/src/datumaro/experimental/converters/mask_converters.py index 40e83dd835..59f66c1658 100644 --- a/src/datumaro/experimental/converters/mask_converters.py +++ b/src/datumaro/experimental/converters/mask_converters.py @@ -92,7 +92,11 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: output_column_name = self.output_mask.name output_shape_column_name = self.output_mask.name + "_shape" - def polygons_to_mask(polygons_data: list, labels_data: list, img_info: dict) -> tuple[list[int], list[int]]: + # Pre-compute dtype and normalize flag once + _numpy_dtype = polars_to_numpy_dtype(self.output_mask.field.dtype) + _normalize = self.input_polygon.field.normalize + + def polygons_to_mask(polygons_data: list, labels_data: list, img_info: dict) -> tuple[np.ndarray, list[int]]: """Rasterize polygons into indexed mask using OpenCV contour filling. The mask uses: @@ -104,11 +108,10 @@ def polygons_to_mask(polygons_data: list, labels_data: list, img_info: dict) -> image_height = img_info["height"] # Initialize mask with background index - numpy_dtype = polars_to_numpy_dtype(self.output_mask.field.dtype) mask = np.full( shape=(image_height, image_width), fill_value=self.background_index, - dtype=numpy_dtype, + dtype=_numpy_dtype, ) # Rasterize each polygon @@ -117,7 +120,7 @@ def polygons_to_mask(polygons_data: list, labels_data: list, img_info: dict) -> class_index = labels_data[i] # Denormalize coordinates if needed - if self.input_polygon.field.normalize: + if _normalize: coords = coords.copy() coords[:, 0] *= image_width coords[:, 1] *= image_height @@ -126,14 +129,7 @@ def polygons_to_mask(polygons_data: list, labels_data: list, img_info: dict) -> contour = coords.astype(np.int32) # Fill polygon with class index + 1 (to reserve 0 for background) - # This means label 0 becomes mask index 1, label 1 becomes mask index 2, etc. - cv2.drawContours( - mask, - [contour], - 0, - int(class_index) + 1, # +1 to shift labels and reserve 0 for background - thickness=cv2.FILLED, - ) + cv2.fillPoly(mask, [contour], int(class_index) + 1) return mask.reshape(-1), [image_height, image_width] @@ -153,7 +149,7 @@ def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: return pl.struct( pl.Series(results_batch_polygons).alias("mask"), - pl.Series(results_batch_shape, dtype=pl.List(pl.Int32())).alias("shape"), + pl.Series("shape", results_batch_shape, dtype=pl.List(pl.Int32())), eager=True, ) @@ -217,53 +213,39 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: output_column_name = self.output_instance_mask.name output_shape_column_name = self.output_instance_mask.name + "_shape" - def polygons_to_instance_masks(polygons_data: list, img_info: dict) -> tuple[list[bool], list[int]]: - """Rasterize polygons into instance masks using OpenCV contour filling.""" - # Extract image dimensions + # Pre-compute dtype and normalize flag once + _numpy_dtype = polars_to_numpy_dtype(self.output_instance_mask.field.dtype) + _normalize = self.input_polygon.field.normalize + # Always use uint8 for internal Polars storage (List(Boolean) construction is extremely slow) + _use_uint8_storage = self.output_instance_mask.field.dtype == pl.Boolean() + _storage_pl_dtype = pl.UInt8() if _use_uint8_storage else self.output_instance_mask.field.dtype + + def polygons_to_instance_masks(polygons_data: list, img_info: dict) -> tuple[np.ndarray, list[int]]: + """Rasterize polygons into instance masks using cv2.fillPoly (batch-optimized).""" image_width = img_info["width"] image_height = img_info["height"] + n = len(polygons_data) - # Convert dtype - use uint8 for OpenCV, then convert to bool - numpy_dtype = polars_to_numpy_dtype(self.output_instance_mask.field.dtype) - - if len(polygons_data) == 0: - # No polygons, return empty mask with shape (0, H, W) - empty_mask = np.array([], dtype=numpy_dtype) - return empty_mask.tolist(), [0, image_height, image_width] + if n == 0: + empty_mask = np.array([], dtype=np.uint8) + return empty_mask, [0, image_height, image_width] - # Create instance masks for each polygon - instance_masks = [] + # Pre-allocate the full (N, H, W) output in one shot as uint8 + stacked_masks = np.zeros((n, image_height, image_width), dtype=np.uint8) - for polygon_data in polygons_data: + for i, polygon_data in enumerate(polygons_data): coords = polygon_data.to_numpy() - # Initialize mask for this instance (use uint8 for OpenCV compatibility) - mask = np.zeros((image_height, image_width), dtype=np.uint8) - - # Denormalize coordinates if needed - if self.input_polygon.field.normalize: + if _normalize: coords = coords.copy() coords[:, 0] *= image_width coords[:, 1] *= image_height - # Convert to OpenCV contour format contour = coords.astype(np.int32) + # fillPoly is faster than drawContours for single-polygon fills + cv2.fillPoly(stacked_masks[i], [contour], 1) - # Fill polygon with 1 for instance mask - cv2.drawContours( - mask, - [contour], - 0, - 1, # Fill with 1 for binary instance mask - thickness=cv2.FILLED, - ) - - # Convert to the target dtype (e.g., bool) - mask = mask.astype(numpy_dtype) - instance_masks.append(mask) - - # Stack into (N, H, W) tensor - stacked_masks = np.stack(instance_masks, axis=0) + # Keep as uint8 for Polars storage; dtype cast happens after collection return stacked_masks.reshape(-1), list(stacked_masks.shape) # Apply conversion using map_batches @@ -293,14 +275,17 @@ def apply_conversion_batch(batch_df: pl.DataFrame, **kwargs) -> pl.DataFrame: # ] ).map_batches( apply_conversion_batch, - return_dtype=pl.Struct( - {"mask": pl.List(self.output_instance_mask.field.dtype), "shape": pl.List(pl.Int32())} - ), + return_dtype=pl.Struct({"mask": pl.List(_storage_pl_dtype), "shape": pl.List(pl.Int32())}), ) + mask_col = mask_data.struct.field("mask") + # Cast from uint8 storage back to target dtype if needed + if _use_uint8_storage: + mask_col = mask_col.cast(pl.List(self.output_instance_mask.field.dtype)) + return df.with_columns( [ - mask_data.struct.field("mask").alias(output_column_name), + mask_col.alias(output_column_name), mask_data.struct.field("shape").alias(output_shape_column_name), ] ) diff --git a/tests/unit/experimental/converters/test_mask_converters.py b/tests/unit/experimental/converters/test_mask_converters.py index 23f856a17a..9f7c849b2e 100644 --- a/tests/unit/experimental/converters/test_mask_converters.py +++ b/tests/unit/experimental/converters/test_mask_converters.py @@ -501,3 +501,91 @@ def get_invalid_mask(): # Check that it raises error for invalid shape with pytest.raises(ValueError, match="Mask array must be 2D \(H,W\), got shape \(1, 2, 2\)"): converter_instance.convert(df) + + +def test_polygon_to_instance_mask_converter_uint8_dtype(): + """Test instance mask converter with UInt8 output dtype (no boolean cast path).""" + polygon_coords1 = [[10.0, 10.0], [20.0, 10.0], [15.0, 20.0]] + polygon_coords2 = [[30.0, 30.0], [40.0, 30.0], [40.0, 40.0], [30.0, 40.0]] + + polygon_series = pl.Series([polygon_coords1, polygon_coords2], dtype=pl.List(pl.Array(pl.Float32, 2))) + + df = pl.DataFrame( + { + "polygons": [polygon_series], + "image_info": [{"width": 80, "height": 80}], + } + ) + + converter_instance = PolygonToInstanceMaskConverter() + setattr( + converter_instance, + "input_polygon", + AttributeSpec(name="polygons", field=PolygonField(dtype=pl.Float32(), format="xy", normalize=False)), + ) + setattr(converter_instance, "input_image_info", AttributeSpec(name="image_info", field=ImageInfoField())) + setattr( + converter_instance, + "output_instance_mask", + AttributeSpec(name="instance_mask", field=InstanceMaskField(dtype=pl.UInt8())), + ) + converter_instance.filter_output_spec() + + result_df = converter_instance.convert(df) + + mask_data = np.array(result_df["instance_mask"][0]) + mask_shape = result_df["instance_mask_shape"][0] + masks = mask_data.reshape(mask_shape) + + assert masks.shape == (2, 80, 80) + assert masks.dtype == np.uint8 + assert masks[0, 15, 15] == 1 # Inside triangle + assert masks[0, 5, 5] == 0 # Outside triangle + assert masks[1, 35, 35] == 1 # Inside rectangle + assert not np.any(masks[0] & masks[1]) + + +def test_polygon_to_instance_mask_converter_multi_sample_batch(): + """Test instance mask converter with multiple samples in a single DataFrame.""" + poly1 = pl.Series( + [[[10.0, 10.0], [20.0, 10.0], [15.0, 20.0]]], + dtype=pl.List(pl.Array(pl.Float32, 2)), + ) + poly2 = pl.Series( + [[[30.0, 30.0], [50.0, 30.0], [50.0, 50.0], [30.0, 50.0]]], + dtype=pl.List(pl.Array(pl.Float32, 2)), + ) + + df = pl.DataFrame( + { + "polygons": [poly1, poly2], + "image_info": [{"width": 60, "height": 60}, {"width": 60, "height": 60}], + } + ) + + converter_instance = PolygonToInstanceMaskConverter() + setattr( + converter_instance, + "input_polygon", + AttributeSpec(name="polygons", field=PolygonField(dtype=pl.Float32(), format="xy", normalize=False)), + ) + setattr(converter_instance, "input_image_info", AttributeSpec(name="image_info", field=ImageInfoField())) + setattr( + converter_instance, + "output_instance_mask", + AttributeSpec(name="instance_mask", field=InstanceMaskField(dtype=pl.Boolean())), + ) + converter_instance.filter_output_spec() + + result_df = converter_instance.convert(df) + + # Check first sample (triangle) + mask1 = np.array(result_df["instance_mask"][0]).reshape(result_df["instance_mask_shape"][0]) + assert mask1.shape == (1, 60, 60) + assert mask1[0, 15, 15] # Inside triangle + + # Check second sample (rectangle) + mask2 = np.array(result_df["instance_mask"][1]).reshape(result_df["instance_mask_shape"][1]) + assert mask2.shape == (1, 60, 60) + assert mask2[0, 40, 40] # Inside rectangle + assert not mask2[0, 5, 5] # Outside rectangle From 783ededb6db1bf2ef7462c522f50ed9e31f146e6 Mon Sep 17 00:00:00 2001 From: Albert van Houten Date: Thu, 28 May 2026 16:31:33 +0200 Subject: [PATCH 2/5] refactor(mask_converters): enhance dtype handling for polygon rasterization and add uint16 test Signed-off-by: Albert van Houten --- .../converters/mask_converters.py | 10 +++-- .../converters/test_mask_converters.py | 39 +++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/src/datumaro/experimental/converters/mask_converters.py b/src/datumaro/experimental/converters/mask_converters.py index 59f66c1658..043d6d6200 100644 --- a/src/datumaro/experimental/converters/mask_converters.py +++ b/src/datumaro/experimental/converters/mask_converters.py @@ -227,10 +227,10 @@ def polygons_to_instance_masks(polygons_data: list, img_info: dict) -> tuple[np. n = len(polygons_data) if n == 0: - empty_mask = np.array([], dtype=np.uint8) + empty_mask = np.array([], dtype=np.uint8 if _use_uint8_storage else _numpy_dtype) return empty_mask, [0, image_height, image_width] - # Pre-allocate the full (N, H, W) output in one shot as uint8 + # Pre-allocate the full (N, H, W) output in one shot as uint8 (required by OpenCV) stacked_masks = np.zeros((n, image_height, image_width), dtype=np.uint8) for i, polygon_data in enumerate(polygons_data): @@ -245,7 +245,11 @@ def polygons_to_instance_masks(polygons_data: list, img_info: dict) -> tuple[np. # fillPoly is faster than drawContours for single-polygon fills cv2.fillPoly(stacked_masks[i], [contour], 1) - # Keep as uint8 for Polars storage; dtype cast happens after collection + # For Boolean output, keep as uint8 for fast Polars List(UInt8) storage; + # the cast to Boolean happens after collection. For other dtypes, cast now. + if not _use_uint8_storage and _numpy_dtype != np.uint8: + stacked_masks = stacked_masks.astype(_numpy_dtype) + return stacked_masks.reshape(-1), list(stacked_masks.shape) # Apply conversion using map_batches diff --git a/tests/unit/experimental/converters/test_mask_converters.py b/tests/unit/experimental/converters/test_mask_converters.py index 9f7c849b2e..5a1d1f31fd 100644 --- a/tests/unit/experimental/converters/test_mask_converters.py +++ b/tests/unit/experimental/converters/test_mask_converters.py @@ -545,6 +545,45 @@ def test_polygon_to_instance_mask_converter_uint8_dtype(): assert not np.any(masks[0] & masks[1]) +def test_polygon_to_instance_mask_converter_uint16_dtype(): + """Test instance mask converter with UInt16 output dtype to verify proper casting.""" + polygon_coords1 = [[10.0, 10.0], [20.0, 10.0], [15.0, 20.0]] + + polygon_series = pl.Series([polygon_coords1], dtype=pl.List(pl.Array(pl.Float32, 2))) + + df = pl.DataFrame( + { + "polygons": [polygon_series], + "image_info": [{"width": 50, "height": 50}], + } + ) + + converter_instance = PolygonToInstanceMaskConverter() + setattr( + converter_instance, + "input_polygon", + AttributeSpec(name="polygons", field=PolygonField(dtype=pl.Float32(), format="xy", normalize=False)), + ) + setattr(converter_instance, "input_image_info", AttributeSpec(name="image_info", field=ImageInfoField())) + setattr( + converter_instance, + "output_instance_mask", + AttributeSpec(name="instance_mask", field=InstanceMaskField(dtype=pl.UInt16())), + ) + converter_instance.filter_output_spec() + + result_df = converter_instance.convert(df) + + mask_data = np.array(result_df["instance_mask"][0]) + mask_shape = result_df["instance_mask_shape"][0] + masks = mask_data.reshape(mask_shape) + + assert masks.shape == (1, 50, 50) + assert masks.dtype == np.uint16 + assert masks[0, 15, 15] == 1 # Inside triangle + assert masks[0, 5, 5] == 0 # Outside + + def test_polygon_to_instance_mask_converter_multi_sample_batch(): """Test instance mask converter with multiple samples in a single DataFrame.""" poly1 = pl.Series( From bbd8f4cf2a8f571ba44a8af069a00da4d973affb Mon Sep 17 00:00:00 2001 From: Albert van Houten Date: Fri, 29 May 2026 11:44:32 +0200 Subject: [PATCH 3/5] Make polygon to mask functions public Signed-off-by: Albert van Houten --- .../experimental/converters/__init__.py | 4 + .../converters/mask_converters.py | 215 ++++++++++++------ 2 files changed, 148 insertions(+), 71 deletions(-) diff --git a/src/datumaro/experimental/converters/__init__.py b/src/datumaro/experimental/converters/__init__.py index 061a5818d7..5598ef0000 100644 --- a/src/datumaro/experimental/converters/__init__.py +++ b/src/datumaro/experimental/converters/__init__.py @@ -49,6 +49,8 @@ MaskChannelsFirstConverter, PolygonToInstanceMaskConverter, PolygonToMaskConverter, + polygons_to_instance_masks, + polygons_to_mask, ) from datumaro.experimental.converters.media_converters import ( ImageInfoToMediaInfoConverter, @@ -162,4 +164,6 @@ "copy_columns_with_shape", "find_conversion_path", "list_eval_ref", + "polygons_to_instance_masks", + "polygons_to_mask", ] diff --git a/src/datumaro/experimental/converters/mask_converters.py b/src/datumaro/experimental/converters/mask_converters.py index 043d6d6200..63b396b78c 100644 --- a/src/datumaro/experimental/converters/mask_converters.py +++ b/src/datumaro/experimental/converters/mask_converters.py @@ -1,6 +1,7 @@ # Copyright (C) 2025 Intel Corporation # # SPDX-License-Identifier: MIT +from collections.abc import Mapping, Sequence from typing import Any import cv2 @@ -22,6 +23,131 @@ from datumaro.experimental.type_registry import polars_to_numpy_dtype from datumaro.util.mask_tools import generate_colormap +# A single polygon's coordinates: anything convertible by ``np.asarray`` into an +# ``(N, 2)`` array of ``(x, y)`` points (e.g. ``numpy.ndarray``, ``polars.Series``, +# or a nested ``list``/``tuple`` of points). +PolygonCoords = Any + + +def polygons_to_mask( + polygons_data: Sequence[PolygonCoords], + labels_data: Sequence[int], + img_info: Mapping[str, int], + *, + dtype: np.dtype = np.uint8, + normalize: bool = False, + background_index: int = 0, +) -> tuple[np.ndarray, list[int]]: + """Rasterize polygons into a single indexed (semantic) mask using OpenCV contour filling. + + The mask uses: + + - ``background_index`` (default ``0``): Background (empty areas) + - Index 1+: Polygon class labels (shifted by 1 to reserve 0 for background) + + Args: + polygons_data: Sequence of polygon coordinates. Each element is + converted with ``np.asarray`` into an ``(N, 2)`` array of + ``(x, y)`` points, so ``numpy.ndarray``, ``polars.Series``, or + nested ``list``/``tuple`` are all accepted. + labels_data: Class label index for each polygon. + img_info: Mapping with ``"width"`` and ``"height"`` keys. + dtype: Output numpy dtype for the mask. + normalize: If True, polygon coordinates are treated as normalized + (in ``[0, 1]``) and are scaled by the image dimensions. + background_index: Fill value used for background pixels. + + Returns: + A tuple ``(flattened_mask, shape)`` where ``flattened_mask`` is the + flattened ``(H * W,)`` mask array and ``shape`` is ``[height, width]``. + """ + # Extract image dimensions + image_width = img_info["width"] + image_height = img_info["height"] + + # Initialize mask with background index + mask = np.full( + shape=(image_height, image_width), + fill_value=background_index, + dtype=dtype, + ) + + # Rasterize each polygon + for i, polygon_data in enumerate(polygons_data): + coords = polygon_data.to_numpy() if hasattr(polygon_data, "to_numpy") else np.asarray(polygon_data) + class_index = labels_data[i] + + # Denormalize coordinates if needed + if normalize: + coords = coords.astype(np.float64, copy=True) + coords[:, 0] *= image_width + coords[:, 1] *= image_height + + # Convert to OpenCV contour format + contour = coords.astype(np.int32) + + # Fill polygon with class index + 1 (to reserve 0 for background) + cv2.fillPoly(mask, [contour], int(class_index) + 1) + + return mask.reshape(-1), [image_height, image_width] + + +def polygons_to_instance_masks( + polygons_data: Sequence[PolygonCoords], + img_info: Mapping[str, int], + *, + dtype: np.dtype = np.uint8, + normalize: bool = False, +) -> tuple[np.ndarray, list[int]]: + """Rasterize polygons into binary instance masks using OpenCV contour filling. + + Produces a stack of binary masks of shape ``(N, H, W)`` where ``N`` is the + number of polygons. Each mask represents a single instance without category + information. + + Args: + polygons_data: Sequence of polygon coordinates. Each element is + converted with ``np.asarray`` into an ``(N, 2)`` array of + ``(x, y)`` points, so ``numpy.ndarray``, ``polars.Series``, or + nested ``list``/``tuple`` are all accepted. + img_info: Mapping with ``"width"`` and ``"height"`` keys. + dtype: Output numpy dtype for the masks. + normalize: If True, polygon coordinates are treated as normalized + (in ``[0, 1]``) and are scaled by the image dimensions. + + Returns: + A tuple ``(flattened_masks, shape)`` where ``flattened_masks`` is the + flattened ``(N * H * W,)`` mask array and ``shape`` is + ``[num_instances, height, width]``. + """ + image_width = img_info["width"] + image_height = img_info["height"] + n = len(polygons_data) + + if n == 0: + empty_mask = np.array([], dtype=dtype) + return empty_mask, [0, image_height, image_width] + + # Pre-allocate the full (N, H, W) output in one shot as uint8 (required by OpenCV) + stacked_masks = np.zeros((n, image_height, image_width), dtype=np.uint8) + + for i, polygon_data in enumerate(polygons_data): + coords = polygon_data.to_numpy() if hasattr(polygon_data, "to_numpy") else np.asarray(polygon_data) + + if normalize: + coords = coords.astype(np.float64, copy=True) + coords[:, 0] *= image_width + coords[:, 1] *= image_height + + contour = coords.astype(np.int32) + # fillPoly is faster than drawContours for single-polygon fills + cv2.fillPoly(stacked_masks[i], [contour], 1) + + if dtype != np.uint8: + stacked_masks = stacked_masks.astype(dtype) + + return stacked_masks.reshape(-1), list(stacked_masks.shape) + @converter(lazy=True) class PolygonToMaskConverter(Converter): @@ -95,43 +221,7 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: # Pre-compute dtype and normalize flag once _numpy_dtype = polars_to_numpy_dtype(self.output_mask.field.dtype) _normalize = self.input_polygon.field.normalize - - def polygons_to_mask(polygons_data: list, labels_data: list, img_info: dict) -> tuple[np.ndarray, list[int]]: - """Rasterize polygons into indexed mask using OpenCV contour filling. - - The mask uses: - - Index 0: Background (empty areas) - - Index 1+: Polygon class labels (shifted by 1 to reserve 0 for background) - """ - # Extract image dimensions - image_width = img_info["width"] - image_height = img_info["height"] - - # Initialize mask with background index - mask = np.full( - shape=(image_height, image_width), - fill_value=self.background_index, - dtype=_numpy_dtype, - ) - - # Rasterize each polygon - for i, polygon_data in enumerate(polygons_data): - coords = polygon_data.to_numpy() - class_index = labels_data[i] - - # Denormalize coordinates if needed - if _normalize: - coords = coords.copy() - coords[:, 0] *= image_width - coords[:, 1] *= image_height - - # Convert to OpenCV contour format - contour = coords.astype(np.int32) - - # Fill polygon with class index + 1 (to reserve 0 for background) - cv2.fillPoly(mask, [contour], int(class_index) + 1) - - return mask.reshape(-1), [image_height, image_width] + _background_index = self.background_index # Apply conversion using map_batches def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: @@ -143,7 +233,14 @@ def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: results_batch_polygons = [] results_batch_shape = [] for polygons, labels, img_infos in zip(batch_polygons, batch_labels, batch_img_infos): - mask_data, shape_data = polygons_to_mask(polygons, labels, img_infos) + mask_data, shape_data = polygons_to_mask( + polygons, + labels, + img_infos, + dtype=_numpy_dtype, + normalize=_normalize, + background_index=_background_index, + ) results_batch_polygons.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) @@ -219,38 +316,9 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: # Always use uint8 for internal Polars storage (List(Boolean) construction is extremely slow) _use_uint8_storage = self.output_instance_mask.field.dtype == pl.Boolean() _storage_pl_dtype = pl.UInt8() if _use_uint8_storage else self.output_instance_mask.field.dtype - - def polygons_to_instance_masks(polygons_data: list, img_info: dict) -> tuple[np.ndarray, list[int]]: - """Rasterize polygons into instance masks using cv2.fillPoly (batch-optimized).""" - image_width = img_info["width"] - image_height = img_info["height"] - n = len(polygons_data) - - if n == 0: - empty_mask = np.array([], dtype=np.uint8 if _use_uint8_storage else _numpy_dtype) - return empty_mask, [0, image_height, image_width] - - # Pre-allocate the full (N, H, W) output in one shot as uint8 (required by OpenCV) - stacked_masks = np.zeros((n, image_height, image_width), dtype=np.uint8) - - for i, polygon_data in enumerate(polygons_data): - coords = polygon_data.to_numpy() - - if _normalize: - coords = coords.copy() - coords[:, 0] *= image_width - coords[:, 1] *= image_height - - contour = coords.astype(np.int32) - # fillPoly is faster than drawContours for single-polygon fills - cv2.fillPoly(stacked_masks[i], [contour], 1) - - # For Boolean output, keep as uint8 for fast Polars List(UInt8) storage; - # the cast to Boolean happens after collection. For other dtypes, cast now. - if not _use_uint8_storage and _numpy_dtype != np.uint8: - stacked_masks = stacked_masks.astype(_numpy_dtype) - - return stacked_masks.reshape(-1), list(stacked_masks.shape) + # For Boolean output, rasterize as uint8 for fast Polars List(UInt8) storage; + # the cast to Boolean happens after collection. + _storage_numpy_dtype = np.uint8 if _use_uint8_storage else _numpy_dtype # Apply conversion using map_batches def apply_conversion_batch(batch_df: pl.DataFrame, **kwargs) -> pl.DataFrame: # noqa: ARG001 @@ -262,7 +330,12 @@ def apply_conversion_batch(batch_df: pl.DataFrame, **kwargs) -> pl.DataFrame: # results_batch_shape = [] for polygons, img_info in zip(batch_polygons, batch_img_infos): - mask_data, shape_data = polygons_to_instance_masks(polygons, img_info) + mask_data, shape_data = polygons_to_instance_masks( + polygons, + img_info, + dtype=_storage_numpy_dtype, + normalize=_normalize, + ) results_batch_mask.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) From 9b4db477bb09c5a89fd47a170f41b20897412f50 Mon Sep 17 00:00:00 2001 From: Albert van Houten Date: Mon, 1 Jun 2026 09:59:27 +0200 Subject: [PATCH 4/5] feat(mask_converters): add mask_info parameter for customizable output mask size Signed-off-by: Albert van Houten --- .../converters/mask_converters.py | 61 ++++++- .../converters/test_mask_converters.py | 166 ++++++++++++++++++ 2 files changed, 223 insertions(+), 4 deletions(-) diff --git a/src/datumaro/experimental/converters/mask_converters.py b/src/datumaro/experimental/converters/mask_converters.py index 63b396b78c..d0d22cadb8 100644 --- a/src/datumaro/experimental/converters/mask_converters.py +++ b/src/datumaro/experimental/converters/mask_converters.py @@ -37,6 +37,7 @@ def polygons_to_mask( dtype: np.dtype = np.uint8, normalize: bool = False, background_index: int = 0, + mask_info: Mapping[str, int] | None = None, ) -> tuple[np.ndarray, list[int]]: """Rasterize polygons into a single indexed (semantic) mask using OpenCV contour filling. @@ -56,6 +57,10 @@ def polygons_to_mask( normalize: If True, polygon coordinates are treated as normalized (in ``[0, 1]``) and are scaled by the image dimensions. background_index: Fill value used for background pixels. + mask_info: Optional mapping with ``"width"`` and ``"height"`` keys + describing the output mask size. Defaults to ``img_info``. When it + differs from the image size, polygon coordinates are scaled to fit + the mask. Returns: A tuple ``(flattened_mask, shape)`` where ``flattened_mask`` is the @@ -65,9 +70,19 @@ def polygons_to_mask( image_width = img_info["width"] image_height = img_info["height"] + # Resolve target mask dimensions (default to the image dimensions) + mask_info = mask_info if mask_info is not None else img_info + out_width = mask_info["width"] + out_height = mask_info["height"] + + # Scale factors to map image-space coordinates into mask-space + scale_x = out_width / image_width + scale_y = out_height / image_height + needs_scaling = scale_x != 1.0 or scale_y != 1.0 + # Initialize mask with background index mask = np.full( - shape=(image_height, image_width), + shape=(out_height, out_width), fill_value=background_index, dtype=dtype, ) @@ -83,13 +98,19 @@ def polygons_to_mask( coords[:, 0] *= image_width coords[:, 1] *= image_height + # Scale coordinates to mask dimensions if they differ from the image + if needs_scaling: + coords = coords.astype(np.float64, copy=True) + coords[:, 0] *= scale_x + coords[:, 1] *= scale_y + # Convert to OpenCV contour format contour = coords.astype(np.int32) # Fill polygon with class index + 1 (to reserve 0 for background) cv2.fillPoly(mask, [contour], int(class_index) + 1) - return mask.reshape(-1), [image_height, image_width] + return mask.reshape(-1), [out_height, out_width] def polygons_to_instance_masks( @@ -98,6 +119,7 @@ def polygons_to_instance_masks( *, dtype: np.dtype = np.uint8, normalize: bool = False, + mask_info: Mapping[str, int] | None = None, ) -> tuple[np.ndarray, list[int]]: """Rasterize polygons into binary instance masks using OpenCV contour filling. @@ -114,6 +136,10 @@ def polygons_to_instance_masks( dtype: Output numpy dtype for the masks. normalize: If True, polygon coordinates are treated as normalized (in ``[0, 1]``) and are scaled by the image dimensions. + mask_info: Optional mapping with ``"width"`` and ``"height"`` keys + describing the output mask size. Defaults to ``img_info``. When it + differs from the image size, polygon coordinates are scaled to fit + the mask. Returns: A tuple ``(flattened_masks, shape)`` where ``flattened_masks`` is the @@ -124,12 +150,22 @@ def polygons_to_instance_masks( image_height = img_info["height"] n = len(polygons_data) + # Resolve target mask dimensions (default to the image dimensions) + mask_info = mask_info if mask_info is not None else img_info + out_width = mask_info["width"] + out_height = mask_info["height"] + + # Scale factors to map image-space coordinates into mask-space + scale_x = out_width / image_width + scale_y = out_height / image_height + needs_scaling = scale_x != 1.0 or scale_y != 1.0 + if n == 0: empty_mask = np.array([], dtype=dtype) - return empty_mask, [0, image_height, image_width] + return empty_mask, [0, out_height, out_width] # Pre-allocate the full (N, H, W) output in one shot as uint8 (required by OpenCV) - stacked_masks = np.zeros((n, image_height, image_width), dtype=np.uint8) + stacked_masks = np.zeros((n, out_height, out_width), dtype=np.uint8) for i, polygon_data in enumerate(polygons_data): coords = polygon_data.to_numpy() if hasattr(polygon_data, "to_numpy") else np.asarray(polygon_data) @@ -139,6 +175,12 @@ def polygons_to_instance_masks( coords[:, 0] *= image_width coords[:, 1] *= image_height + # Scale coordinates to mask dimensions if they differ from the image + if needs_scaling: + coords = coords.astype(np.float64, copy=True) + coords[:, 0] *= scale_x + coords[:, 1] *= scale_y + contour = coords.astype(np.int32) # fillPoly is faster than drawContours for single-polygon fills cv2.fillPoly(stacked_masks[i], [contour], 1) @@ -165,6 +207,9 @@ class PolygonToMaskConverter(Converter): # Configuration options background_index: int = 0 # Background value + # Optional output mask size as a mapping with "width"/"height" keys. + # When None, the mask matches the image size (img_info). + mask_info: Mapping[str, int] | None = None def filter_output_spec(self) -> bool: """ @@ -222,6 +267,7 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: _numpy_dtype = polars_to_numpy_dtype(self.output_mask.field.dtype) _normalize = self.input_polygon.field.normalize _background_index = self.background_index + _mask_info = self.mask_info # Apply conversion using map_batches def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: @@ -240,6 +286,7 @@ def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: dtype=_numpy_dtype, normalize=_normalize, background_index=_background_index, + mask_info=_mask_info, ) results_batch_polygons.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) @@ -283,6 +330,10 @@ class PolygonToInstanceMaskConverter(Converter): input_image_info: AttributeSpec[ImageInfoField] output_instance_mask: AttributeSpec[InstanceMaskField] + # Optional output mask size as a mapping with "width"/"height" keys. + # When None, the mask matches the image size (img_info). + mask_info: Mapping[str, int] | None = None + def filter_output_spec(self) -> bool: """Configure output specification for instance mask format.""" # Configure output for instance mask format @@ -319,6 +370,7 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: # For Boolean output, rasterize as uint8 for fast Polars List(UInt8) storage; # the cast to Boolean happens after collection. _storage_numpy_dtype = np.uint8 if _use_uint8_storage else _numpy_dtype + _mask_info = self.mask_info # Apply conversion using map_batches def apply_conversion_batch(batch_df: pl.DataFrame, **kwargs) -> pl.DataFrame: # noqa: ARG001 @@ -335,6 +387,7 @@ def apply_conversion_batch(batch_df: pl.DataFrame, **kwargs) -> pl.DataFrame: # img_info, dtype=_storage_numpy_dtype, normalize=_normalize, + mask_info=_mask_info, ) results_batch_mask.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) diff --git a/tests/unit/experimental/converters/test_mask_converters.py b/tests/unit/experimental/converters/test_mask_converters.py index 5a1d1f31fd..4a859eb65c 100644 --- a/tests/unit/experimental/converters/test_mask_converters.py +++ b/tests/unit/experimental/converters/test_mask_converters.py @@ -12,6 +12,7 @@ PolygonToInstanceMaskConverter, PolygonToMaskConverter, ) +from datumaro.experimental.converters.mask_converters import polygons_to_instance_masks, polygons_to_mask from datumaro.experimental.fields import ( ImageInfoField, InstanceMaskCallableField, @@ -628,3 +629,168 @@ def test_polygon_to_instance_mask_converter_multi_sample_batch(): assert mask2.shape == (1, 60, 60) assert mask2[0, 40, 40] # Inside rectangle assert not mask2[0, 5, 5] # Outside rectangle + + +# --------------------------------------------------------------------------- +# Tests for generating masks at a different size than the image +# --------------------------------------------------------------------------- + + +def test_polygons_to_mask_default_size_matches_image(): + """Without mask_info the mask matches the image dimensions.""" + # Square polygon covering the left/top quadrant of a 20x20 image. + polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] + + mask_data, shape = polygons_to_mask(polygons, [0], {"width": 20, "height": 20}) + mask = mask_data.reshape(shape) + + assert shape == [20, 20] + assert mask[5, 5] == 1 # Inside the square (label 0 -> stored as 1) + assert mask[15, 15] == 0 # Background + + +def test_polygons_to_mask_smaller_than_image(): + """A smaller mask scales the polygon coordinates down accordingly.""" + polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] + + mask_data, shape = polygons_to_mask( + polygons, [0], {"width": 20, "height": 20}, mask_info={"width": 10, "height": 10} + ) + mask = mask_data.reshape(shape) + + assert shape == [10, 10] + # The 10x10 square in a 20x20 image scales to a 5x5 square in a 10x10 mask. + assert mask[2, 2] == 1 # Inside the scaled square + assert mask[8, 8] == 0 # Background (outside scaled square) + + +def test_polygons_to_mask_larger_than_image(): + """A larger mask scales the polygon coordinates up accordingly.""" + polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] + + mask_data, shape = polygons_to_mask( + polygons, [0], {"width": 20, "height": 20}, mask_info={"width": 40, "height": 40} + ) + mask = mask_data.reshape(shape) + + assert shape == [40, 40] + # The 10x10 square scales to a 20x20 square in a 40x40 mask. + assert mask[10, 10] == 1 # Inside the scaled-up square + assert mask[30, 30] == 0 # Background + + +def test_polygons_to_mask_non_uniform_size(): + """Width and height can be scaled independently.""" + polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] + + mask_data, shape = polygons_to_mask( + polygons, [0], {"width": 20, "height": 20}, mask_info={"width": 40, "height": 10} + ) + mask = mask_data.reshape(shape) + + assert shape == [10, 40] # [height, width] + # x scales by 2 (20->40), y scales by 0.5 (20->10): square becomes 20 wide, 5 tall. + assert mask[2, 10] == 1 # Inside the stretched square + assert mask[8, 30] == 0 # Background + + +def test_polygons_to_instance_masks_custom_size(): + """Instance masks can also be rasterized at a different size than the image.""" + polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] + + mask_data, shape = polygons_to_instance_masks( + polygons, {"width": 20, "height": 20}, mask_info={"width": 40, "height": 40} + ) + masks = mask_data.reshape(shape) + + assert shape == [1, 40, 40] + assert masks[0, 10, 10] == 1 # Inside the scaled-up instance + assert masks[0, 35, 35] == 0 # Background + + +def test_polygons_to_instance_masks_empty_custom_size(): + """An empty polygon list still respects the requested mask size.""" + mask_data, shape = polygons_to_instance_masks( + [], {"width": 20, "height": 20}, mask_info={"width": 30, "height": 15} + ) + + assert shape == [0, 15, 30] + assert mask_data.size == 0 + + +def _make_polygon_to_mask_converter(mask_info=None): + converter_instance = PolygonToMaskConverter(mask_info=mask_info) # type: ignore[call-arg] + setattr( + converter_instance, + "input_polygon", + AttributeSpec(name="polygons", field=PolygonField(dtype=pl.Float32(), format="xy", normalize=False)), + ) + setattr( + converter_instance, + "input_labels", + AttributeSpec(name="labels", field=LabelField(dtype=pl.UInt32(), multi_label=True)), + ) + setattr(converter_instance, "input_image_info", AttributeSpec(name="image_info", field=ImageInfoField())) + setattr(converter_instance, "output_mask", AttributeSpec(name="mask", field=MaskField(dtype=pl.UInt8()))) + converter_instance.filter_output_spec() + return converter_instance + + +def test_polygon_to_mask_converter_custom_size(): + """The PolygonToMaskConverter honors the mask_info option.""" + polygon_coords = [[0.0, 0.0], [10.0, 0.0], [10.0, 10.0], [0.0, 10.0]] + polygon_series = pl.Series([polygon_coords], dtype=pl.List(pl.Array(pl.Float32, 2))) + + df = pl.DataFrame( + { + "polygons": [polygon_series], + "labels": [[0]], + "image_info": [{"width": 20, "height": 20}], + } + ) + + converter_instance = _make_polygon_to_mask_converter(mask_info={"width": 10, "height": 10}) + result_df = converter_instance.convert(df) + + mask = np.array(result_df["mask"][0]).reshape(result_df["mask_shape"][0]) + + assert mask.shape == (10, 10) + assert mask[2, 2] == 1 # Inside the scaled square + assert mask[8, 8] == 0 # Background + + +def test_polygon_to_instance_mask_converter_custom_size(): + """The PolygonToInstanceMaskConverter honors the mask_info option.""" + polygon_coords = [[0.0, 0.0], [10.0, 0.0], [10.0, 10.0], [0.0, 10.0]] + polygon_series = pl.Series([polygon_coords], dtype=pl.List(pl.Array(pl.Float32, 2))) + + df = pl.DataFrame( + { + "polygons": [polygon_series], + "image_info": [{"width": 20, "height": 20}], + } + ) + + converter_instance = PolygonToInstanceMaskConverter( # type: ignore[call-arg] + mask_info={"width": 40, "height": 40} + ) + setattr( + converter_instance, + "input_polygon", + AttributeSpec(name="polygons", field=PolygonField(dtype=pl.Float32(), format="xy", normalize=False)), + ) + setattr(converter_instance, "input_image_info", AttributeSpec(name="image_info", field=ImageInfoField())) + setattr( + converter_instance, + "output_instance_mask", + AttributeSpec(name="instance_mask", field=InstanceMaskField(dtype=pl.Boolean())), + ) + converter_instance.filter_output_spec() + + result_df = converter_instance.convert(df) + + masks = np.array(result_df["instance_mask"][0]).reshape(result_df["instance_mask_shape"][0]) + + assert masks.shape == (1, 40, 40) + assert masks[0, 10, 10] # Inside the scaled-up instance + assert not masks[0, 35, 35] # Background From cc18efdbbe194b0478d5a24bc784f340335e4c26 Mon Sep 17 00:00:00 2001 From: Albert van Houten Date: Mon, 1 Jun 2026 11:30:27 +0200 Subject: [PATCH 5/5] refactor(mask_converters): replace img_info with img_size and mask_info with mask_size for clarity Signed-off-by: Albert van Houten --- .../converters/mask_converters.py | 57 ++++------ .../converters/test_mask_converters.py | 102 ++---------------- 2 files changed, 26 insertions(+), 133 deletions(-) diff --git a/src/datumaro/experimental/converters/mask_converters.py b/src/datumaro/experimental/converters/mask_converters.py index d0d22cadb8..129538dbc9 100644 --- a/src/datumaro/experimental/converters/mask_converters.py +++ b/src/datumaro/experimental/converters/mask_converters.py @@ -1,7 +1,7 @@ # Copyright (C) 2025 Intel Corporation # # SPDX-License-Identifier: MIT -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Any import cv2 @@ -32,12 +32,12 @@ def polygons_to_mask( polygons_data: Sequence[PolygonCoords], labels_data: Sequence[int], - img_info: Mapping[str, int], + img_size: tuple[int, int], *, dtype: np.dtype = np.uint8, normalize: bool = False, background_index: int = 0, - mask_info: Mapping[str, int] | None = None, + mask_size: tuple[int, int] | None = None, ) -> tuple[np.ndarray, list[int]]: """Rasterize polygons into a single indexed (semantic) mask using OpenCV contour filling. @@ -52,28 +52,24 @@ def polygons_to_mask( ``(x, y)`` points, so ``numpy.ndarray``, ``polars.Series``, or nested ``list``/``tuple`` are all accepted. labels_data: Class label index for each polygon. - img_info: Mapping with ``"width"`` and ``"height"`` keys. + img_size: Image size as a ``(width, height)`` tuple. dtype: Output numpy dtype for the mask. normalize: If True, polygon coordinates are treated as normalized (in ``[0, 1]``) and are scaled by the image dimensions. background_index: Fill value used for background pixels. - mask_info: Optional mapping with ``"width"`` and ``"height"`` keys - describing the output mask size. Defaults to ``img_info``. When it - differs from the image size, polygon coordinates are scaled to fit - the mask. + mask_size: Optional output mask size as a ``(width, height)`` tuple. + Defaults to ``img_size``. When it differs from the image size, + polygon coordinates are scaled to fit the mask. Returns: A tuple ``(flattened_mask, shape)`` where ``flattened_mask`` is the flattened ``(H * W,)`` mask array and ``shape`` is ``[height, width]``. """ # Extract image dimensions - image_width = img_info["width"] - image_height = img_info["height"] + image_width, image_height = img_size # Resolve target mask dimensions (default to the image dimensions) - mask_info = mask_info if mask_info is not None else img_info - out_width = mask_info["width"] - out_height = mask_info["height"] + out_width, out_height = mask_size if mask_size is not None else img_size # Scale factors to map image-space coordinates into mask-space scale_x = out_width / image_width @@ -115,11 +111,11 @@ def polygons_to_mask( def polygons_to_instance_masks( polygons_data: Sequence[PolygonCoords], - img_info: Mapping[str, int], + img_size: tuple[int, int], *, dtype: np.dtype = np.uint8, normalize: bool = False, - mask_info: Mapping[str, int] | None = None, + mask_size: tuple[int, int] | None = None, ) -> tuple[np.ndarray, list[int]]: """Rasterize polygons into binary instance masks using OpenCV contour filling. @@ -132,28 +128,24 @@ def polygons_to_instance_masks( converted with ``np.asarray`` into an ``(N, 2)`` array of ``(x, y)`` points, so ``numpy.ndarray``, ``polars.Series``, or nested ``list``/``tuple`` are all accepted. - img_info: Mapping with ``"width"`` and ``"height"`` keys. + img_size: Image size as a ``(width, height)`` tuple. dtype: Output numpy dtype for the masks. normalize: If True, polygon coordinates are treated as normalized (in ``[0, 1]``) and are scaled by the image dimensions. - mask_info: Optional mapping with ``"width"`` and ``"height"`` keys - describing the output mask size. Defaults to ``img_info``. When it - differs from the image size, polygon coordinates are scaled to fit - the mask. + mask_size: Optional output mask size as a ``(width, height)`` tuple. + Defaults to ``img_size``. When it differs from the image size, + polygon coordinates are scaled to fit the mask. Returns: A tuple ``(flattened_masks, shape)`` where ``flattened_masks`` is the flattened ``(N * H * W,)`` mask array and ``shape`` is ``[num_instances, height, width]``. """ - image_width = img_info["width"] - image_height = img_info["height"] + image_width, image_height = img_size n = len(polygons_data) # Resolve target mask dimensions (default to the image dimensions) - mask_info = mask_info if mask_info is not None else img_info - out_width = mask_info["width"] - out_height = mask_info["height"] + out_width, out_height = mask_size if mask_size is not None else img_size # Scale factors to map image-space coordinates into mask-space scale_x = out_width / image_width @@ -207,9 +199,6 @@ class PolygonToMaskConverter(Converter): # Configuration options background_index: int = 0 # Background value - # Optional output mask size as a mapping with "width"/"height" keys. - # When None, the mask matches the image size (img_info). - mask_info: Mapping[str, int] | None = None def filter_output_spec(self) -> bool: """ @@ -267,7 +256,6 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: _numpy_dtype = polars_to_numpy_dtype(self.output_mask.field.dtype) _normalize = self.input_polygon.field.normalize _background_index = self.background_index - _mask_info = self.mask_info # Apply conversion using map_batches def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: @@ -282,11 +270,10 @@ def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: mask_data, shape_data = polygons_to_mask( polygons, labels, - img_infos, + (img_infos["width"], img_infos["height"]), dtype=_numpy_dtype, normalize=_normalize, background_index=_background_index, - mask_info=_mask_info, ) results_batch_polygons.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) @@ -330,10 +317,6 @@ class PolygonToInstanceMaskConverter(Converter): input_image_info: AttributeSpec[ImageInfoField] output_instance_mask: AttributeSpec[InstanceMaskField] - # Optional output mask size as a mapping with "width"/"height" keys. - # When None, the mask matches the image size (img_info). - mask_info: Mapping[str, int] | None = None - def filter_output_spec(self) -> bool: """Configure output specification for instance mask format.""" # Configure output for instance mask format @@ -370,7 +353,6 @@ def convert(self, df: pl.DataFrame) -> pl.DataFrame: # For Boolean output, rasterize as uint8 for fast Polars List(UInt8) storage; # the cast to Boolean happens after collection. _storage_numpy_dtype = np.uint8 if _use_uint8_storage else _numpy_dtype - _mask_info = self.mask_info # Apply conversion using map_batches def apply_conversion_batch(batch_df: pl.DataFrame, **kwargs) -> pl.DataFrame: # noqa: ARG001 @@ -384,10 +366,9 @@ def apply_conversion_batch(batch_df: pl.DataFrame, **kwargs) -> pl.DataFrame: # for polygons, img_info in zip(batch_polygons, batch_img_infos): mask_data, shape_data = polygons_to_instance_masks( polygons, - img_info, + (img_info["width"], img_info["height"]), dtype=_storage_numpy_dtype, normalize=_normalize, - mask_info=_mask_info, ) results_batch_mask.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) diff --git a/tests/unit/experimental/converters/test_mask_converters.py b/tests/unit/experimental/converters/test_mask_converters.py index 4a859eb65c..060e5faf30 100644 --- a/tests/unit/experimental/converters/test_mask_converters.py +++ b/tests/unit/experimental/converters/test_mask_converters.py @@ -637,11 +637,11 @@ def test_polygon_to_instance_mask_converter_multi_sample_batch(): def test_polygons_to_mask_default_size_matches_image(): - """Without mask_info the mask matches the image dimensions.""" + """Without mask_size the mask matches the image dimensions.""" # Square polygon covering the left/top quadrant of a 20x20 image. polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] - mask_data, shape = polygons_to_mask(polygons, [0], {"width": 20, "height": 20}) + mask_data, shape = polygons_to_mask(polygons, [0], (20, 20)) mask = mask_data.reshape(shape) assert shape == [20, 20] @@ -653,9 +653,7 @@ def test_polygons_to_mask_smaller_than_image(): """A smaller mask scales the polygon coordinates down accordingly.""" polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] - mask_data, shape = polygons_to_mask( - polygons, [0], {"width": 20, "height": 20}, mask_info={"width": 10, "height": 10} - ) + mask_data, shape = polygons_to_mask(polygons, [0], (20, 20), mask_size=(10, 10)) mask = mask_data.reshape(shape) assert shape == [10, 10] @@ -668,9 +666,7 @@ def test_polygons_to_mask_larger_than_image(): """A larger mask scales the polygon coordinates up accordingly.""" polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] - mask_data, shape = polygons_to_mask( - polygons, [0], {"width": 20, "height": 20}, mask_info={"width": 40, "height": 40} - ) + mask_data, shape = polygons_to_mask(polygons, [0], (20, 20), mask_size=(40, 40)) mask = mask_data.reshape(shape) assert shape == [40, 40] @@ -683,9 +679,7 @@ def test_polygons_to_mask_non_uniform_size(): """Width and height can be scaled independently.""" polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] - mask_data, shape = polygons_to_mask( - polygons, [0], {"width": 20, "height": 20}, mask_info={"width": 40, "height": 10} - ) + mask_data, shape = polygons_to_mask(polygons, [0], (20, 20), mask_size=(40, 10)) mask = mask_data.reshape(shape) assert shape == [10, 40] # [height, width] @@ -698,9 +692,7 @@ def test_polygons_to_instance_masks_custom_size(): """Instance masks can also be rasterized at a different size than the image.""" polygons = [np.array([[0, 0], [10, 0], [10, 10], [0, 10]])] - mask_data, shape = polygons_to_instance_masks( - polygons, {"width": 20, "height": 20}, mask_info={"width": 40, "height": 40} - ) + mask_data, shape = polygons_to_instance_masks(polygons, (20, 20), mask_size=(40, 40)) masks = mask_data.reshape(shape) assert shape == [1, 40, 40] @@ -710,87 +702,7 @@ def test_polygons_to_instance_masks_custom_size(): def test_polygons_to_instance_masks_empty_custom_size(): """An empty polygon list still respects the requested mask size.""" - mask_data, shape = polygons_to_instance_masks( - [], {"width": 20, "height": 20}, mask_info={"width": 30, "height": 15} - ) + mask_data, shape = polygons_to_instance_masks([], (20, 20), mask_size=(30, 15)) assert shape == [0, 15, 30] assert mask_data.size == 0 - - -def _make_polygon_to_mask_converter(mask_info=None): - converter_instance = PolygonToMaskConverter(mask_info=mask_info) # type: ignore[call-arg] - setattr( - converter_instance, - "input_polygon", - AttributeSpec(name="polygons", field=PolygonField(dtype=pl.Float32(), format="xy", normalize=False)), - ) - setattr( - converter_instance, - "input_labels", - AttributeSpec(name="labels", field=LabelField(dtype=pl.UInt32(), multi_label=True)), - ) - setattr(converter_instance, "input_image_info", AttributeSpec(name="image_info", field=ImageInfoField())) - setattr(converter_instance, "output_mask", AttributeSpec(name="mask", field=MaskField(dtype=pl.UInt8()))) - converter_instance.filter_output_spec() - return converter_instance - - -def test_polygon_to_mask_converter_custom_size(): - """The PolygonToMaskConverter honors the mask_info option.""" - polygon_coords = [[0.0, 0.0], [10.0, 0.0], [10.0, 10.0], [0.0, 10.0]] - polygon_series = pl.Series([polygon_coords], dtype=pl.List(pl.Array(pl.Float32, 2))) - - df = pl.DataFrame( - { - "polygons": [polygon_series], - "labels": [[0]], - "image_info": [{"width": 20, "height": 20}], - } - ) - - converter_instance = _make_polygon_to_mask_converter(mask_info={"width": 10, "height": 10}) - result_df = converter_instance.convert(df) - - mask = np.array(result_df["mask"][0]).reshape(result_df["mask_shape"][0]) - - assert mask.shape == (10, 10) - assert mask[2, 2] == 1 # Inside the scaled square - assert mask[8, 8] == 0 # Background - - -def test_polygon_to_instance_mask_converter_custom_size(): - """The PolygonToInstanceMaskConverter honors the mask_info option.""" - polygon_coords = [[0.0, 0.0], [10.0, 0.0], [10.0, 10.0], [0.0, 10.0]] - polygon_series = pl.Series([polygon_coords], dtype=pl.List(pl.Array(pl.Float32, 2))) - - df = pl.DataFrame( - { - "polygons": [polygon_series], - "image_info": [{"width": 20, "height": 20}], - } - ) - - converter_instance = PolygonToInstanceMaskConverter( # type: ignore[call-arg] - mask_info={"width": 40, "height": 40} - ) - setattr( - converter_instance, - "input_polygon", - AttributeSpec(name="polygons", field=PolygonField(dtype=pl.Float32(), format="xy", normalize=False)), - ) - setattr(converter_instance, "input_image_info", AttributeSpec(name="image_info", field=ImageInfoField())) - setattr( - converter_instance, - "output_instance_mask", - AttributeSpec(name="instance_mask", field=InstanceMaskField(dtype=pl.Boolean())), - ) - converter_instance.filter_output_spec() - - result_df = converter_instance.convert(df) - - masks = np.array(result_df["instance_mask"][0]).reshape(result_df["instance_mask_shape"][0]) - - assert masks.shape == (1, 40, 40) - assert masks[0, 10, 10] # Inside the scaled-up instance - assert not masks[0, 35, 35] # Background