import geopandas as gpd
import numpy as np
import rasterio
from rasterio.features import rasterize, shapes
from rasterio.io import DatasetReader
from shapely.geometry import shape
from .transform_utils import get_transform_with_resampling, pad_mask_to_min_size
[docs]
def mask_to_polygons(
mask: np.ndarray,
dataset_reader: DatasetReader,
window: rasterio.windows.Window | None = None,
target_resolution_cm_per_pixel: float | None = None,
) -> gpd.GeoDataFrame:
"""
Convert a binary or labeled mask to a GeoDataFrame of polygons using rasterio.
This function extracts polygon geometries from a 2D NumPy array (mask), where nonzero
values represent features of interest (e.g., 1 for foreground, 0 for background, or
any labeled regions).
When the mask has been resampled to a different resolution than the original dataset,
the `target_resolution_cm_per_pixel` parameter ensures correct spatial positioning by
automatically adjusting the affine transform. This works for both upsampling (finer
resolution) and downsampling (coarser resolution).
Parameters
----------
mask : np.ndarray
A 2D array with 0 for background and any nonzero value for features.
dataset_reader : rasterio.io.DatasetReader
Rasterio dataset that the mask is derived from, used for transform and CRS.
window : rasterio.windows.Window, optional
If provided, the transform will be adjusted for the window region.
The mask should correspond to the windowed region.
target_resolution_cm_per_pixel : float, optional
Target resolution in centimeters per pixel if the mask was resampled.
If provided and different from the dataset's original resolution, the
transform will be automatically adjusted to account for resampling.
Examples:
- Original 6 cm/pixel downsampled to 3 cm/pixel: pass 3.0
- Original 3 cm/pixel upsampled to 6 cm/pixel: pass 6.0
- No resampling: omit this parameter or pass None
Returns
-------
gpd.GeoDataFrame
GeoDataFrame with the extracted polygons and their associated values.
Raises
------
ValueError
If target_resolution_cm_per_pixel results in an invalid scale factor.
Examples
--------
>>> import rasterio
>>> import numpy as np
>>> from shapelib.tools.converters import mask_to_polygons
>>>
>>> # Example 1: Basic usage without resampling
>>> mask = np.zeros((100, 100), dtype=np.uint8)
>>> mask[20:80, 20:80] = 1
>>> with rasterio.open('image.tif') as src:
... polygons = mask_to_polygons(mask, src)
... print(len(polygons))
>>>
>>> # Example 2: With resampling (original 6cm -> resampled to 3cm)
>>> resampled_mask = np.zeros((200, 200), dtype=np.uint8)
>>> resampled_mask[40:160, 40:160] = 1
>>> with rasterio.open('image_6cm.tif') as src:
... polygons = mask_to_polygons(
... resampled_mask,
... src,
... target_resolution_cm_per_pixel=3.0
... )
... print(len(polygons))
>>>
>>> # Example 3: With window
>>> window = rasterio.windows.Window(0, 0, 50, 50)
>>> mask_window = mask[:50, :50]
>>> with rasterio.open('image.tif') as src:
... polygons = mask_to_polygons(mask_window, src, window=window)
"""
transform = get_transform_with_resampling(
dataset_reader=dataset_reader,
window=window,
target_resolution_cm_per_pixel=target_resolution_cm_per_pixel,
)
geometries = []
values = []
for geom, value in shapes(mask, mask > 0, transform=transform):
if value == 0:
continue
geometries.append(shape(geom))
values.append(value)
if geometries:
return gpd.GeoDataFrame({"value": values}, geometry=geometries, crs=dataset_reader.crs)
return gpd.GeoDataFrame(geometry=[], crs=dataset_reader.crs)
[docs]
def polygons_to_mask(
gdf: gpd.GeoDataFrame,
dataset_reader: DatasetReader,
fill_value: int = 1,
background_value: int = 0,
all_touched: bool = False,
window=None,
min_size: tuple[int, int] | None = None,
) -> np.ndarray:
"""
Convert a GeoDataFrame of polygons to a binary mask using rasterio.
This function rasterizes polygon geometries from a GeoDataFrame into a binary mask
with the same dimensions and coordinate reference system as the provided rasterio
dataset reader. If a window is provided, only the region inside the window is
rasterized. If min_size is specified and the resulting mask is smaller than the
minimum size, it will be padded with background_value to match the exact minimum size.
Parameters
----------
gdf : gpd.GeoDataFrame
GeoDataFrame containing polygon geometries to be rasterized.
dataset_reader : rasterio.io.DatasetReader
Rasterio dataset reader that defines the target raster properties
(transform, CRS, height, width) for the output mask.
fill_value : int, optional
Value to assign to pixels that intersect with polygon geometries.
Defaults to 1.
background_value : int, optional
Value to assign to pixels that do not intersect with any polygon.
Defaults to 0.
all_touched : bool, optional
If True, all pixels touched by geometries will be burned in.
If False, only pixels whose center is within the polygon or that
are selected by Bresenham's line algorithm will be burned in.
Defaults to False for more precise boundary adherence.
window : rasterio.windows.Window, optional
If provided, only rasterize the region inside the window. The output mask
will have shape (window.height, window.width).
min_size : tuple[int, int], optional
Minimum size (height, width) for the output mask. If the generated mask is
smaller than this size, it will be padded with background_value to match
the exact minimum size. If None or if the mask is already equal to or larger
than min_size, no padding is applied. Defaults to None.
Returns
-------
np.ndarray
A 2D numpy array (mask) with the same dimensions as the input raster or window
(or min_size if padding was applied), where pixels intersecting with polygons
have the fill_value and background pixels have the background_value.
Examples
--------
>>> import rasterio
>>> import geopandas as gpd
>>> from shapelib.tools.converters import polygons_to_mask
>>>
>>> # Load a raster and polygon data
>>> with rasterio.open('image.tif') as src:
... gdf = gpd.read_file('polygons.shp')
... mask = polygons_to_mask(gdf, src)
... print(mask.shape) # Should match src.shape
...
... # With window
... window = rasterio.windows.Window(0, 0, 50, 50)
... mask_window = polygons_to_mask(gdf, src, window=window)
... print(mask_window.shape) # Should be (50, 50)
...
... # With minimum size padding
... mask_padded = polygons_to_mask(gdf, src, window=window, min_size=(256, 256))
... print(mask_padded.shape) # Should be (256, 256) with zero padding
"""
if gdf.crs != dataset_reader.crs:
gdf = gdf.to_crs(dataset_reader.crs)
gdf = gdf[gdf.geometry.notnull() & gdf.geometry.type.isin(["Polygon", "MultiPolygon"])]
if window is not None:
raster_shape = (window.height, window.width)
raster_transform = dataset_reader.window_transform(window)
else:
raster_shape = (dataset_reader.height, dataset_reader.width)
raster_transform = dataset_reader.transform
if len(gdf) == 0: # If no valid geometries, return background mask
mask = np.full(raster_shape, background_value, dtype="uint8")
else:
mask = rasterize(
[(geom, fill_value) for geom in gdf.geometry],
out_shape=raster_shape,
transform=raster_transform,
fill=background_value,
dtype="uint8",
all_touched=all_touched,
)
# Apply padding if min_size is specified and mask is smaller
if min_size is not None:
mask = pad_mask_to_min_size(mask, min_size, background_value)
return mask