Finish mosaicking when all pixels within a cutline are populated

[yellowcap]

The current implementation of the MosaicMethodBase will exit_when_filled only when all pixels of the array are populated.

When using a cutline, this is never the case. So for methods such as the FirstMethod, the algorithm will run through all the scenes or assets that are passed on.

This is not efficient and could be improved.

The following example will run through all the S2 scenes that are returned from the STAC search (limited to 3 here for demo purposes)

import datetime

import httpx
import rasterio
from matplotlib import pyplot as plt
from rasterio.crs import CRS
from rasterio.features import bounds as featureBounds
from rasterio.warp import transform_geom

from rio_tiler.io import STACReader
from rio_tiler.mosaic import mosaic_reader
# from ipyleaflet import Map, basemaps, GeoJSON
from rio_tiler.utils import create_cutline

stac_endpoint = "https://earth-search.aws.element84.com/v0/search"

# Date filter
date_min = "2022-08-01"
date_max = "2022-08-31"

start = datetime.datetime.strptime(date_min, "%Y-%m-%d").strftime("%Y-%m-%dT00:00:00Z")
end = datetime.datetime.strptime(date_max, "%Y-%m-%d").strftime("%Y-%m-%dT23:59:59Z")

# AOI over Bern
aoi = {
    "type": "Feature",
    "properties": {},
    "geometry": {
        "type": "Polygon",
        "coordinates": [
            [
                [-151.98538128444676, 137.785308042979636],
                [427.064224021652308, 334.495939510601829],
                [727.671315771469153, 78.218004147854629],
                [1410.617381359762931, 64.365142776895368],
                [1750.012484948266319, -33.990172956915444],
                [1675.207033545085778, -581.178197109808139],
                [1435.552531827489929, -773.732970166142536],
                [873.126360166541872, -869.317713625761712],
                [722.130171223085426, -885.941147270912893],
                [590.527988198972253, -1478.843613947971562],
                [222.041875731454525, -1501.00819214150647],
                [-250.34069701825797, -1215.63924789974476],
                [-530.168496711635839, -297.19453900514236],
                [-151.98538128444676, 137.785308042979636],
            ]
        ],
    },
}
epsg = 21780

query = {
    "collections": [
        "sentinel-s2-l2a-cogs"
    ],  # Make sure to query only sentinel-2 COGs collection
    "datetime": f"{start}/{end}",
    "query": {"eo:cloud_cover": {"lt": 90}},
    "intersects": transform_geom(f"epsg:{epsg}", "epsg:4326", aoi["geometry"]),
    "limit": 3,
    "fields": {
        "include": [
            "id",
            "properties.datetime",
            "properties.eo:cloud_cover",
        ],  # Make returned response ligth
        "exclude": ["links"],
    },
}

headers = {
    "Content-Type": "application/json",
    "Accept-Encoding": "gzip",
    "Accept": "application/geo+json",
}

data = httpx.post(stac_endpoint, headers=headers, json=query).json()

# Prepare Stac Item dicts for the STACReader.
items = [feat | {"stac_version": data["stac_version"]} for feat in data["features"]]


# Define custom tiler
def tiler(asset, *args, **kwargs):
    with STACReader(None, asset) as src:
        # Create WTT Cutline
        with rasterio.open(src.item.assets[kwargs["assets"][0]].href) as itm:
            cutline = create_cutline(
                itm, kwargs["geometry"], geometry_crs=kwargs["geometry_crs"]
            )

        # Get BBOX of the polygon
        bbox = featureBounds(kwargs["geometry"])

        # Read part of the data (bbox) and use the cutline to mask the data
        data, mask = src.part(
            bbox,
            bounds_crs=kwargs["geometry_crs"],
            dst_crs=kwargs["geometry_crs"],
            vrt_options={"cutline": cutline},
            assets=kwargs["assets"],
            max_size=2000,
        )

        print("min/max/shape", data.min(), data.max(), data.shape)

        return data, mask


# Get false color image data
img, assets_used = mosaic_reader(
    items,
    tiler,
    geometry=aoi,
    geometry_crs=CRS.from_epsg(21780),
    assets=["B04", "B03", "B08"],
    threads=1,
)

# Plot false color map
fig = plt.figure(figsize=(30, 10))

ax = fig.add_subplot(1, 2, 1)
img.rescale(in_range=((0, 5000),))
ax.imshow(img.data_as_image())

ax = fig.add_subplot(1, 2, 2)
ax.imshow(img.mask)

plt.show()

And produce the following plot:

The inspiration for this was an attempt of combining the cutline feature with the mosaic mechanism.

[yellowcap]

Here is a solution that does the mosaicking with early stop without changes on the repository. We could consider adding a version of the CutlineMethod into the package. It could also simply go into the docs as an advanced use case. Open to suggestions.

import datetime
from typing import Dict, Optional

import httpx
import numpy
from affine import Affine
from matplotlib import pyplot as plt
from rasterio.crs import CRS
from rasterio.features import bounds as featureBounds
from rasterio.features import is_valid_geom, rasterize
from rasterio.warp import transform_geom

from rio_tiler.errors import RioTilerError
from rio_tiler.io import STACReader
from rio_tiler.models import ImageData
from rio_tiler.mosaic import mosaic_reader
from rio_tiler.mosaic.methods.defaults import FirstMethod


class CutlineMethod(FirstMethod):
    """Feed the mosaic tile with the first pixel available over a cutline"""

    def __init__(self, cutline_mask: numpy.ndarray):
        """Overwrite base and init Cutline method."""
        super().__init__()
        self.cutline_mask = cutline_mask

    def feed(self, tile: Optional[numpy.ma.MaskedArray]):
        """Add data to tile."""
        super().feed(tile)
        self.tile.mask = numpy.where(~self.cutline_mask, self.tile.mask, True)

    @property
    def is_done(self) -> bool:
        """Check if the tile filling is done.

        Returns:
            bool

        """
        if self.tile is None:
            return False

        if numpy.sum(numpy.where(~self.cutline_mask, self.tile.mask, False)) == 0:
            return True

        return False

# Define helper functions for creating the cutline mask
def compute_mask(
    geometry: Dict,
    height: int,
    width: int,
    transform: Affine,
    all_touched: bool = False,
) -> numpy.ndarray:
    """Rasterize the input geometry using the raster definition in transform and size"""

    if "geometry" in geometry:
        geometry = geometry["geometry"]

    if not is_valid_geom(geometry):
        raise RioTilerError("Invalid geometry")

    rasterized = rasterize(
        [geometry],
        out_shape=(height, width),
        transform=transform,
        all_touched=all_touched,
        default_value=0,
        fill=1,
        dtype="uint8",
    )

    return rasterized.astype("bool")


def get_cutline_args(
    geometry: Optional[Dict] = None,
    target_resolution: Optional[float] = None,
) -> tuple:
    """Compute rasterization parameters for cutline mask array"""

    bbox = featureBounds(geometry)

    # Compute width and height in target resolution
    width = int(numpy.ceil((bbox[2] - bbox[0]) / target_resolution))
    height = int(numpy.ceil((bbox[3] - bbox[1]) / target_resolution))

    # Ensure bbox is a multiple of the target resolution.
    bbox = (
        bbox[0],
        bbox[1],
        bbox[0] + width * target_resolution,
        bbox[1]
        + height * target_resolution,  # Assuming Y min is at the top of the image
    )

    transform = Affine(
        target_resolution,
        0,
        bbox[0],
        0,
        -target_resolution,  # Assuming Y min is at the top of the image
        bbox[3],
    )
    cutline_mask = compute_mask(geometry, height, width, transform)

    return bbox, width, height, cutline_mask

# Define custom tiler
def tiler(asset, bbox, bounds_crs, width, height, *args, **kwargs) -> ImageData:
    with STACReader(None, asset) as src:
        img = src.part(
            bbox,
            bounds_crs=bounds_crs,
            dst_crs=bounds_crs,
            width=width,
            height=height,
            assets=kwargs["assets"],
        )

    print("min/max/shape", img.data.min(), img.data.max(), img.data.shape)

    return img

# Set stac endpoint
stac_endpoint = "https://earth-search.aws.element84.com/v0/search"

# Date filter
date_min = "2022-05-01"
date_max = "2022-05-31"

start = datetime.datetime.strptime(date_min, "%Y-%m-%d").strftime("%Y-%m-%dT00:00:00Z")
end = datetime.datetime.strptime(date_max, "%Y-%m-%d").strftime("%Y-%m-%dT23:59:59Z")

# AOI over Bern
aoi = {
    "type": "Feature",
    "properties": {},
    "geometry": {
        "type": "Polygon",
        "coordinates": [
            [
                [-151.98538128444676, 137.785308042979636],
                [427.064224021652308, 334.495939510601829],
                [727.671315771469153, 78.218004147854629],
                [1410.617381359762931, 64.365142776895368],
                [1750.012484948266319, -33.990172956915444],
                [1675.207033545085778, -581.178197109808139],
                [1435.552531827489929, -773.732970166142536],
                [873.126360166541872, -869.317713625761712],
                [722.130171223085426, -885.941147270912893],
                [590.527988198972253, -1478.843613947971562],
                [222.041875731454525, -1501.00819214150647],
                [-250.34069701825797, -1215.63924789974476],
                [-530.168496711635839, -297.19453900514236],
                [-151.98538128444676, 137.785308042979636],
            ]
        ],
    },
}
epsg = 21780

query = {
    "collections": [
        "sentinel-s2-l2a-cogs"
    ],  # Make sure to query only sentinel-2 COGs collection
    "datetime": f"{start}/{end}",
    "query": {"eo:cloud_cover": {"lt": 90}},
    "intersects": transform_geom(f"epsg:{epsg}", "epsg:4326", aoi["geometry"]),
    "limit": 2,
    "fields": {
        "include": [
            "id",
            "properties.datetime",
            "properties.eo:cloud_cover",
            "properties.proj:epsg",
            "properties.proj:transform",
        ],  # Make returned response ligth
        "exclude": ["links"],
    },
}

headers = {
    "Content-Type": "application/json",
    "Accept-Encoding": "gzip",
    "Accept": "application/geo+json",
}

data = httpx.post(stac_endpoint, headers=headers, json=query).json()

# Prepare Stac Item dicts for the STACReader.
items = [feat | {"stac_version": data["stac_version"]} for feat in data["features"]]

# Compute bbox and size of target raster. Also compute cutline mask as numpy array.
bbox, width, height, cutline_mask = get_cutline_args(
    geometry=aoi,
    target_resolution=10,
)

# Get false color image data
img, assets_used = mosaic_reader(
    items,
    tiler,
    assets=["B08", "B03", "B02"],
    threads=1,
    pixel_selection=CutlineMethod(cutline_mask=cutline_mask),
    bbox=bbox,
    bounds_crs=CRS.from_epsg(epsg),
    width=width,
    height=height,
)

# Plot false color map
fig = plt.figure(figsize=(30, 10))

ax = fig.add_subplot(1, 2, 1)
img.rescale(in_range=((0, 5000),))
ax.imshow(img.data_as_image())

ax = fig.add_subplot(1, 2, 2)
ax.imshow(img.mask)

plt.show()