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 40e83dd835..129538dbc9 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 Sequence from typing import Any import cv2 @@ -22,6 +23,165 @@ 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_size: tuple[int, int], + *, + dtype: np.dtype = np.uint8, + normalize: bool = False, + background_index: int = 0, + mask_size: tuple[int, int] | None = None, +) -> 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_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_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, image_height = img_size + + # Resolve target mask dimensions (default to the image dimensions) + 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 + 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=(out_height, out_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 + + # 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), [out_height, out_width] + + +def polygons_to_instance_masks( + polygons_data: Sequence[PolygonCoords], + img_size: tuple[int, int], + *, + dtype: np.dtype = np.uint8, + normalize: bool = False, + mask_size: tuple[int, int] | None = None, +) -> 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_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_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, image_height = img_size + n = len(polygons_data) + + # Resolve target mask dimensions (default to the image dimensions) + 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 + 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, 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, 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) + + if normalize: + coords = coords.astype(np.float64, copy=True) + 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) + + 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): @@ -92,50 +252,10 @@ 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]]: - """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 - 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, - ) - - # 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 self.input_polygon.field.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) - # 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, - ) - - return mask.reshape(-1), [image_height, image_width] + # Pre-compute dtype and normalize flag once + _numpy_dtype = polars_to_numpy_dtype(self.output_mask.field.dtype) + _normalize = self.input_polygon.field.normalize + _background_index = self.background_index # Apply conversion using map_batches def apply_conversion_batch(batch_df: pl.DataFrame) -> pl.DataFrame: @@ -147,13 +267,20 @@ 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["width"], img_infos["height"]), + dtype=_numpy_dtype, + normalize=_normalize, + background_index=_background_index, + ) results_batch_polygons.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) 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,54 +344,15 @@ 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 - image_width = img_info["width"] - image_height = img_info["height"] - - # 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] - - # Create instance masks for each polygon - instance_masks = [] - - for polygon_data in 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: - coords = coords.copy() - coords[:, 0] *= image_width - coords[:, 1] *= image_height - - # Convert to OpenCV contour format - contour = coords.astype(np.int32) - - # 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) - return stacked_masks.reshape(-1), list(stacked_masks.shape) + # 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 + # 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 @@ -276,7 +364,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["width"], img_info["height"]), + dtype=_storage_numpy_dtype, + normalize=_normalize, + ) results_batch_mask.append(pl.Series(mask_data)) results_batch_shape.append(shape_data) @@ -293,14 +386,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..060e5faf30 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, @@ -501,3 +502,207 @@ 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_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( + [[[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 + + +# --------------------------------------------------------------------------- +# Tests for generating masks at a different size than the image +# --------------------------------------------------------------------------- + + +def test_polygons_to_mask_default_size_matches_image(): + """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], (20, 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], (20, 20), mask_size=(10, 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], (20, 20), mask_size=(40, 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], (20, 20), mask_size=(40, 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, (20, 20), mask_size=(40, 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([], (20, 20), mask_size=(30, 15)) + + assert shape == [0, 15, 30] + assert mask_data.size == 0