import cv2
import geopandas as gpd
import numpy as np
from rasterio.io import DatasetReader
from rasterio.windows import Window
from shapely.geometry import LineString
from skimage.morphology import remove_small_objects, skeletonize
from shapelib import find_lines_lang
from .transform_utils import get_transform_with_resampling, pad_mask_to_min_size
[docs]
def mask_to_lines(
mask: np.ndarray,
dataset_reader: DatasetReader,
window: Window = None,
dilation_kernel_size: int = 5,
min_distance_points: int = 10,
look_ahead_distance: int = 8,
min_line_size: int = 10,
min_object_size: int = 0,
target_resolution_cm_per_pixel: float | None = None,
) -> gpd.GeoDataFrame:
"""
Convert a binary mask to a GeoDataFrame of line geometries using optimized line
extraction.
Extracts line geometries from a 2D binary mask using morphological operations,
skeletonization, and line tracing.
Parameters
----------
mask : np.ndarray
A 2D binary array where nonzero values represent features to extract lines from.
Values should be 0 for background and any nonzero value for foreground.
dataset_reader : rasterio.io.DatasetReader
Rasterio dataset that the mask is derived from, used for coordinate transformation
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. Defaults to None.
dilation_kernel_size : int, optional
Size of the morphological dilation kernel. Larger values connect nearby features
but may merge distinct lines. Set to 0 to skip dilation. Defaults to 2.
min_distance_points : int, optional
Minimum distance between consecutive points in the line tracing algorithm.
Higher values produce smoother lines with fewer vertices. Defaults to 10.
look_ahead_distance : int, optional
Look-ahead distance for the line tracing algorithm.
Higher values produce smoother but potentially less accurate lines. Defaults to 8.
min_line_size : int, optional
Minimum number of pixels required for a line to be included in the output.
Shorter lines are filtered out to reduce noise. Defaults to 10.
min_object_size : int, optional
Minimum number of pixels required for a connected component to be included in
the output. Connected components smaller than this threshold are filtered out
to reduce noise. Defaults to 0.
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.
This works for both upsampling and downsampling.
Returns
-------
gpd.GeoDataFrame
GeoDataFrame containing LineString geometries extracted from the mask.
Each row represents a single line feature with proper coordinate transformation.
Examples
-------
>>> with rasterio.open('image.tif') as src:
... mask = model.predict(src.read(1)) # Get a binary mask
... gdf = mask_to_lines(mask, src)
... print(gdf)
...
... # Extract lines with resampling
... resampled_mask = model.predict_resampled(src.read(1))
... gdf_resampled = mask_to_lines(
... resampled_mask,
... src,
... target_resolution_cm_per_pixel=3.0
... )
...
... # Extract lines from a specific window
... window = rasterio.windows.Window(0, 0, 100, 100)
... mask_window = mask[:100, :100]
... gdf_window = mask_to_lines(mask_window, src, window=window)
... print(gdf_window)
"""
_validate_mask(mask)
transform = get_transform_with_resampling(
dataset_reader=dataset_reader,
window=window,
target_resolution_cm_per_pixel=target_resolution_cm_per_pixel,
)
processed_mask = _preprocess_mask(mask, dilation_kernel_size, min_object_size)
processed_mask = _skeletonize_and_normalize(processed_mask)
lines = _extract_lines(processed_mask, min_line_size, min_distance_points, look_ahead_distance)
geometries = _lines_to_geometries(lines, transform)
return _geometries_to_gdf(geometries, dataset_reader)
def _validate_mask(mask: np.ndarray) -> None:
if not isinstance(mask, np.ndarray):
raise TypeError("mask must be a numpy array")
if mask.ndim != 2:
raise ValueError("mask must be a 2D array")
def _preprocess_mask(mask: np.ndarray, dilation_kernel_size: int, min_object_size: int) -> np.ndarray:
processed_mask = mask.astype(np.uint8)
if dilation_kernel_size > 0:
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_kernel_size, dilation_kernel_size))
processed_mask = cv2.dilate(processed_mask, kernel)
if min_object_size is not None:
binary_mask = processed_mask > 0
cleaned_mask = remove_small_objects(binary_mask, min_object_size)
processed_mask = cleaned_mask.astype(np.uint8)
return processed_mask
def _skeletonize_and_normalize(processed_mask: np.ndarray) -> np.ndarray:
binary_mask = processed_mask > 0
skeleton = skeletonize(binary_mask, method="lee")
processed_mask = skeleton.astype(np.uint8)
if np.max(processed_mask) > 0:
processed_mask = (processed_mask > 0).astype(np.uint8) * 255
return processed_mask
def _extract_lines(
processed_mask: np.ndarray,
min_line_size: int,
min_distance_points: int,
look_ahead_distance: int,
):
_, _, lines = find_lines_lang.get_lines(
processed_mask,
min_size=min_line_size,
min_distance_points=min_distance_points,
look_ahead=look_ahead_distance,
)
return lines
def _lines_to_geometries(lines, transform):
geometries = []
for line_x, line_y, _ in lines:
if len(line_x) <= 1:
continue
coords = [transform * (line_y[i], line_x[i]) for i in range(len(line_x))]
if len(coords) >= 2:
line_geom = LineString(coords)
geometries.append(line_geom)
return geometries
def _geometries_to_gdf(geometries, dataset_reader):
if geometries:
return gpd.GeoDataFrame(geometry=geometries, crs=dataset_reader.crs)
return gpd.GeoDataFrame(geometry=[], crs=dataset_reader.crs)
[docs]
def lines_to_mask(
gdf: gpd.GeoDataFrame,
dataset_reader: DatasetReader,
fill_value: int = 1,
background_value: int = 0,
all_touched: bool = True,
window: Window | None = None,
min_size: tuple[int, int] | None = None,
) -> np.ndarray:
"""
Convert a GeoDataFrame of lines to a binary mask using rasterio.
This function rasterizes line 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, it will
be padded with background_value to match the exact minimum size.
Parameters
----------
gdf : gpd.GeoDataFrame
GeoDataFrame containing line 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 line geometries.
Defaults to 1.
background_value : int, optional
Value to assign to pixels that do not intersect with any line.
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 line or that
are selected by Bresenham's line algorithm will be burned in.
Defaults to True for better line visibility.
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
lines have the fill_value and background pixels have the background_value.
Examples
--------
>>> import rasterio
>>> import geopandas as gpd
>>> from shapelib.tools.converters import lines_to_mask
>>>
>>> # Load a raster and line data
>>> with rasterio.open('image.tif') as src:
... gdf = gpd.read_file('roads.shp')
... mask = lines_to_mask(gdf, src)
... print(mask.shape) # Should match src.shape
...
... # With window
... window = rasterio.windows.Window(0, 0, 50, 50)
... mask_window = lines_to_mask(gdf, src, window=window)
... print(mask_window.shape) # Should be (50, 50)
...
... # With minimum size padding
... mask_padded = lines_to_mask(gdf, src, window=window, min_size=(256, 256))
... print(mask_padded.shape) # Should be (256, 256) with zero padding
Notes
-----
The all_touched parameter is set to True by default for line geometries to
ensure better line visibility in the rasterized output. Lines are typically
thin features that may be missed if only center points are considered.
"""
from rasterio.features import rasterize
if gdf.crs != dataset_reader.crs:
gdf = gdf.to_crs(dataset_reader.crs)
# Keep only line-like geometries
gdf = gdf[gdf.geometry.notnull() & gdf.geometry.type.isin(["LineString", "MultiLineString"])]
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:
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