Create a water mask in EPSG:4326 using `esa_world_cover_2021` imagery or Natural Earth land polygons.

create-water-mask-natural-earth.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "geopandas",
#     "rasterio",
#     "tyro",
#     "requests",
# ]
# ///
"""Create a water mask in EPSG:4326 using Natural Earth land polygons.

This script downloads the Natural Earth 10m land shapefile and rasterizes it
to a 0/1 (water/land) mask for a specified bounding box.

Examples
--------
uv run --script create-water-mask.py --output land_mask.tif -120.9 34.85 -120.25 35.15

"""
from pathlib import Path

import geopandas as gpd
import rasterio as rio
import requests
import tyro
from rasterio import features
from rasterio.transform import from_bounds

NATURAL_EARTH_URL = "https://naciscdn.org/naturalearth/10m/physical/ne_10m_land.zip"


def download_ne_data(target_path: Path, source: str | Path = NATURAL_EARTH_URL) -> Path:
    """Download Natural Earth zip if it doesn't exist."""
    if not target_path.exists():
        print(f"Downloading Natural Earth data to {target_path}...")
        response = requests.get(source, stream=True)
        response.raise_for_status()
        with open(target_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
    return target_path


def create_land_mask(
    bounds: tuple[float, float, float, float],
    /,
    output: Path,
    resolution: float = 0.001,
) -> None:
    """Create a land mask from Natural Earth polygons.

    Parameters
    ----------
    bounds : tuple[float, float, float, float]
        (west, south, east, north) in degrees longitude/latitude.
    output : Path
        Path to save the output GeoTIFF.
    resolution : float
        Pixel size in degrees (default 0.001 is ~100m).
    """
    if output.exists():
        print(f"{output} exists, exiting.")
        return

    # 1. Setup Data
    cache_dir = Path("data_cache")
    cache_dir.mkdir(exist_ok=True)
    zip_path = cache_dir / "ne_10m_land.zip"
    download_ne_data(zip_path)

    print(f"Loading polygons and clipping to {bounds}...")
    # Load directly from zip using /vsizip/ or geopandas zip support
    gdf = gpd.read_file(zip_path)

    # Clip to bounds to speed up rasterization
    west, south, east, north = bounds
    gdf_clipped = gdf.clip([west, south, east, north])

    # 2. Define Raster Grid
    width = int((east - west) / resolution)
    height = int((north - south) / resolution)

    # Check for empty dimensions
    if width <= 0 or height <= 0:
        raise ValueError("Bounds or resolution result in invalid dimensions.")

    transform = from_bounds(west, south, east, north, width, height)

    print(f"Rasterizing to {width}x{height} grid: {output}")

    # 3. Rasterize
    # Land polygons = 1, Background (water) = 0
    mask_data = features.rasterize(
        [(geom, 1) for geom in gdf_clipped.geometry],
        out_shape=(height, width),
        transform=transform,
        fill=0,
        dtype="uint8",
    )

    # 4. Save
    output.parent.mkdir(exist_ok=True, parents=True)

    meta = {
        "driver": "GTiff",
        "dtype": "uint8",
        "nodata": 255, # Using 255 as nodata, 0/1 for data
        "width": width,
        "height": height,
        "count": 1,
        "crs": "EPSG:4326",
        "transform": transform,
        "compress": "lzw"
    }

    with rio.open(output, "w", **meta) as dst:
        dst.write(mask_data, 1)

    print("Done.")


if __name__ == "__main__":
    tyro.cli(create_land_mask)

create-water-mask-tile-mate.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "tile-mate",
#     "tyro",
# ]
# ///
"""Create a water mask in EPSG:4326 using `esa_world_cover_2021` imagery.

Examples
--------
uv run --script create-water-mask.py --output water_mask.tif -120.9 34.85 -120.25 35.15

"""
from pathlib import Path

import rasterio as rio
import tile_mate
import tyro


def download_water_mask(
    bounds: tuple[float, float, float, float], /, output: Path
) -> None:
    """Download a water mask for the given DISP dataset.

    Retrieves a water mask corresponding to the provided DISP dataset and saves it
    in the specified output directory.

    Parameters
    ----------
    bounds : tuple[int, int, int, int]
        (west, south, east, north) in degrees longitude/latitude of requested area.
    output_dir : str | Path
        Directory where the downloaded water mask file will be saved.

    """
    # Create output directory if does not exist
    if output.exists():
        print(f"{output} exists, exiting.")
        return

    print(f"Saving Water Mask: {output}")
    output_dir = output.parent
    output_dir.mkdir(exist_ok=True)

    mask_data, mask_attr = tile_mate.get_raster_from_tiles(
        bounds, tile_shortname="esa_world_cover_2021"
    )

    # Make byte mask
    # class 80 is water
    water_pixels = mask_data == 80
    # 0 was the old nodata
    nodata_pixels = mask_data == 0

    mask_data[:] = 1
    mask_data[water_pixels] = 0
    mask_data[nodata_pixels] = 255
    mask_data = mask_data.astype("byte")
    mask_attr["nodata"] = 255

    with rio.open(output, "w", **mask_attr) as dst:
        dst.write(mask_data)


if __name__ == "__main__":
    tyro.cli(download_water_mask)

create_water_mask.py

"""High-resolution water mask creation from ASF water mask tiles.

ASF's water mask mosaics OpenStreetMap and ESA WorldCover data at high resolution.
Tiles are publicly hosted at s3://asf-dem-west/WATER_MASK/TILES/ and use the naming
convention {lat_dir}{lat}{lon_dir}{lon}.tif (e.g., n60w165.tif) on a 5-degree grid.

In the source tiles: 1 = land, 0 = water.
Output masks can optionally invert this convention and apply shoreline buffering.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from enum import Enum
from pathlib import Path

import numpy as np
from numpy.typing import NDArray
from osgeo import gdal
from scipy import ndimage

gdal.UseExceptions()

logger = logging.getLogger(__name__)

TILE_URL_BASE = "https://asf-dem-west.s3.amazonaws.com/WATER_MASK/TILES/"


class WaterValue(Enum):
    """Convention for water mask pixel values."""

    ZERO = 0  # water=0, land=1 (ASF native)
    ONE = 1  # water=1, land=0 (inverted)


@dataclass
class GeoTransform:
    """Affine geotransform parameters."""

    x_origin: float
    pixel_width: float
    row_rotation: float
    y_origin: float
    col_rotation: float
    pixel_height: float  # typically negative

    @classmethod
    def from_gdal(cls, gt: tuple[float, ...]) -> GeoTransform:
        return cls(*gt)

    def to_gdal(self) -> tuple[float, ...]:
        return (
            self.x_origin,
            self.pixel_width,
            self.row_rotation,
            self.y_origin,
            self.col_rotation,
            self.pixel_height,
        )

    @property
    def bounds(self) -> tuple[float, float, float, float]:
        """Return (x_min, y_min, x_max, y_max) - requires shape to compute."""
        raise NotImplementedError("Need raster dimensions to compute bounds")

    def get_bounds(self, width: int, height: int) -> tuple[float, float, float, float]:
        """Return (west, south, east, north) given raster dimensions."""
        west = self.x_origin
        north = self.y_origin
        east = west + self.pixel_width * width
        south = north + self.pixel_height * height  # pixel_height is negative
        return (west, south, east, north)


def _coord_to_tile_name(lon: float, lat: float) -> str:
    """Convert a coordinate to ASF tile filename.

    Tiles are on a 5-degree grid. E.g., (lon=-163.5, lat=59.0) -> n55w165.tif
    """
    lat_floor = int(np.floor(lat / 5) * 5)
    lon_floor = int(np.floor(lon / 5) * 5)

    lat_dir = "n" if lat_floor >= 0 else "s"
    lon_dir = "e" if lon_floor >= 0 else "w"

    return f"{lat_dir}{abs(lat_floor):02d}{lon_dir}{abs(lon_floor):03d}.tif"


def _get_tile_urls(west: float, south: float, east: float, north: float) -> list[str]:
    """Get URLs for all tiles covering the given bounds."""
    tiles: list[str] = []
    corners = [
        (west, north),  # upper-left
        (west, south),  # lower-left
        (east, north),  # upper-right
        (east, south),  # lower-right
    ]
    # Handle wide extents (high latitudes) that span multiple tiles
    width = east - west
    if width > 5:
        mid_lon = west + width / 2
        corners.extend([(mid_lon, north), (mid_lon, south)])

    for lon, lat in corners:
        tile_name = _coord_to_tile_name(lon, lat)
        url = f"/vsicurl/{TILE_URL_BASE}{tile_name}"
        if url not in tiles:
            tiles.append(url)

    return tiles


def _get_raster_info(path: str | Path) -> tuple[GeoTransform, int, int, str]:
    """Get geotransform, dimensions, and projection from a raster."""
    ds = gdal.Open(str(path))
    gt = GeoTransform.from_gdal(ds.GetGeoTransform())
    width, height = ds.RasterXSize, ds.RasterYSize
    proj = ds.GetProjection()
    ds = None
    return gt, width, height, proj


def _buffer_mask(
    mask: NDArray[np.uint8],
    buffer_pixels: int,
    water_value: WaterValue,
) -> NDArray[np.uint8]:
    """Expand water regions by buffer_pixels using morphological dilation."""
    if buffer_pixels <= 0:
        return mask

    # Create circular structuring element
    diameter = 2 * buffer_pixels + 1
    y, x = np.ogrid[:diameter, :diameter]
    center = buffer_pixels
    struct = ((x - center) ** 2 + (y - center) ** 2) <= buffer_pixels**2

    if water_value == WaterValue.ONE:
        # Water is 1, dilate it
        return ndimage.binary_dilation(mask, structure=struct).astype(np.uint8)
    else:
        # Water is 0, erode land (dilate inverted, then invert back)
        land = mask == 1
        eroded_land = ~ndimage.binary_dilation(~land, structure=struct)
        return eroded_land.astype(np.uint8)


def create_water_mask(
    bounds: tuple[float, float, float, float],
    output: Path,
    resolution: float | None = None,
    buffer_meters: float = 0.0,
    water_value: WaterValue = WaterValue.ONE,
    overwrite: bool = False,
) -> Path:
    """Create a water mask GeoTIFF for the given bounds.

    Parameters
    ----------
    bounds
        (west, south, east, north) in EPSG:4326 degrees.
    output
        Output GeoTIFF path.
    resolution
        Pixel size in degrees. If None, uses native tile resolution (~0.0001 deg).
    buffer_meters
        Buffer distance to expand water mask into land. Approximate conversion
        to pixels uses 111 km/degree at equator.
    water_value
        Convention for output: ONE means water=1/land=0, ZERO means water=0/land=1.
    overwrite
        If False and output exists, skip creation.

    Returns
    -------
    Path to the output raster.
    """
    if output.exists() and not overwrite:
        logger.info("Output %s exists, skipping.", output)
        return output

    west, south, east, north = bounds
    tiles = _get_tile_urls(west, south, east, north)

    if not tiles:
        raise ValueError(f"No water mask tiles found for bounds {bounds}")

    logger.info("Using %d tile(s): %s", len(tiles), [t.split("/")[-1] for t in tiles])

    # Build VRT from tiles
    vrt_options = gdal.BuildVRTOptions(resampleAlg="nearest")
    vrt_ds = gdal.BuildVRT("", tiles, options=vrt_options)

    if vrt_ds is None:
        raise RuntimeError(f"Failed to build VRT from tiles: {tiles}")

    # Set up output bounds and resolution
    if resolution is None:
        # Use native resolution from first tile
        tile_ds = gdal.Open(tiles[0])
        resolution = tile_ds.GetGeoTransform()[1]
        tile_ds = None

    width = int(np.ceil((east - west) / resolution))
    height = int(np.ceil((north - south) / resolution))

    # Warp to target extent
    warp_options = gdal.WarpOptions(
        format="MEM",
        outputBounds=(west, south, east, north),
        width=width,
        height=height,
        resampleAlg="nearest",
    )
    mem_ds = gdal.Warp("", vrt_ds, options=warp_options)
    vrt_ds = None

    if mem_ds is None:
        raise RuntimeError("Failed to warp water mask to target extent")

    # Read data
    data = mem_ds.GetRasterBand(1).ReadAsArray()
    gt = mem_ds.GetGeoTransform()
    proj = mem_ds.GetProjection()
    mem_ds = None

    # ASF convention: 1=land, 0=water. Invert if requested.
    if water_value == WaterValue.ONE:
        data = 1 - data

    # Apply buffer if requested
    if buffer_meters > 0:
        # Approximate degrees per meter at this latitude
        center_lat = (south + north) / 2
        meters_per_deg = 111_000 * np.cos(np.radians(center_lat))
        buffer_pixels = int(np.ceil(buffer_meters / (resolution * meters_per_deg)))
        logger.info(
            "Buffering water by %d pixels (~%.0f m)", buffer_pixels, buffer_meters
        )
        data = _buffer_mask(data, buffer_pixels, water_value)

    # Write output
    output.parent.mkdir(parents=True, exist_ok=True)
    driver = gdal.GetDriverByName("GTiff")
    out_ds = driver.Create(
        str(output),
        width,
        height,
        1,
        gdal.GDT_Byte,
        options=["COMPRESS=LZW", "TILED=YES"],
    )
    out_ds.SetGeoTransform(gt)
    out_ds.SetProjection(proj)
    band = out_ds.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(255)
    out_ds.FlushCache()
    out_ds = None

    logger.info("Created water mask: %s", output)
    return output


def create_water_mask_like(
    reference: Path,
    output: Path,
    buffer_meters: float = 0.0,
    water_value: WaterValue = WaterValue.ONE,
    overwrite: bool = False,
) -> Path:
    """Create a water mask matching the extent and resolution of a reference raster.

    Parameters
    ----------
    reference
        Reference GeoTIFF to match (must be EPSG:4326).
    output
        Output GeoTIFF path.
    buffer_meters
        Buffer distance to expand water mask into land.
    water_value
        Convention for output values.
    overwrite
        If False and output exists, skip creation.

    Returns
    -------
    Path to the output raster.
    """
    gt, width, height, _ = _get_raster_info(reference)
    bounds = gt.get_bounds(width, height)
    resolution = gt.pixel_width

    return create_water_mask(
        bounds=bounds,
        output=output,
        resolution=resolution,
        buffer_meters=buffer_meters,
        water_value=water_value,
        overwrite=overwrite,
    )


def cli(
    output: Path,
    bounds: tuple[float, float, float, float] | None = None,
    reference: Path | None = None,
    resolution: float | None = None,
    buffer_meters: float = 0.0,
    water_is_one: bool = True,
    overwrite: bool = False,
    verbose: bool = False,
) -> None:
    """Create a high-resolution water mask from ASF water mask tiles.

    Provide either --reference to match an existing raster, or all four
    bounds (west, south, east, north).

    Args:
        output: Output GeoTIFF path.
        bounds: (west, south, east, north) in latitude/ longitude.
        reference: Reference raster to match extent/resolution.
        resolution: Pixel size in degrees (default: native ~0.0001 deg).
        buffer_meters: Expand water mask into land by this distance.
        water_is_one: If True, water=1/land=0; if False, water=0/land=1.
        overwrite: Overwrite existing output.
        verbose: Enable debug logging.
    """
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.INFO,
        format="%(levelname)s: %(message)s",
    )

    water_value = WaterValue.ONE if water_is_one else WaterValue.ZERO

    if reference is None and bounds is None:
        raise ValueError(
            "Provide either --reference or all bounds (--west, --south, --east, --north)"
        )
    if reference is not None:
        create_water_mask_like(
            reference=reference,
            output=output,
            buffer_meters=buffer_meters,
            water_value=water_value,
            overwrite=overwrite,
        )
    else:
        create_water_mask(
            bounds=bounds,
            output=output,
            resolution=resolution,
            buffer_meters=buffer_meters,
            water_value=water_value,
            overwrite=overwrite,
        )


if __name__ == "__main__":
    import tyro

    tyro.cli(cli)