Masking one raster based on a values in a different raster

[scottstanie]

I'm wondering if there's a simple way to achieve this feature:

I have several sets of raster images that I'm viewing with titiler. Here's one (called "unwrapped phase"):

I'd like to be able to create/set a mask layer based on the values of another raster

(e.g. correlation < 0.1) so that I can render the above rasters transparent.

I've made a custom algorithm using parameters, but I'm not sure how it would work to pass another raster URI as a parameter.

Any thoughts/suggestions are appreciated. Thank you for all the great work on Titiler!

[vincentsarago]

@scottstanie this is a nice use case, sadly there is no direct ways to do this without either using a custom Reader and a custom path_dependency, or a custom dependency and full custom endpoint.

a custom reader could look like this

from typing import Any, Dict, List, Type, TypedDict, Optional
import contextlib

import attr
from rasterio.crs import CRS
from morecantile import TileMatrixSet
from rio_tiler.constants import WEB_MERCATOR_TMS, WGS84_CRS
from rio_tiler.io import BaseReader, Reader
from rio_tiler.types import BBox, Indexes
from rio_tiler.models import BandStatistics, ImageData, Info, PointData

class Input(TypedDict, total=True):
    """Reader Options."""

    data: str
    mask: str
    max_correlation: float


@attr.s
class CustomReader(BaseReader):

    input: Input = attr.ib()

    tms: TileMatrixSet = attr.ib(default=WEB_MERCATOR_TMS)
    geographic_crs: CRS = attr.ib(default=WGS84_CRS)

    dataset: Type[Reader] = attr.ib(
        init=False,
    )
    mask: Type[Reader] = attr.ib(
        init=False,
    )

    colormap: Dict = attr.ib(init=False, default=None)

    # Context Manager to handle rasterio open/close
    _ctx_stack = attr.ib(init=False, factory=contextlib.ExitStack)

    def __attrs_post_init__(self):
        """Define _kwargs, open dataset and get info."""
        self.dataset = self._ctx_stack.enter_context(Reader(self.input["data"]))
        self.mask = self._ctx_stack.enter_context(Reader(self.input["mask"]))

        self.bounds = self.dataset.bounds
        self.crs = self.dataset.crs
        self.colormap = self.dataset.colormap

    @property
    def minzoom(self):
        """Return dataset minzoom."""
        return self.dataset.minzoom

    @property
    def maxzoom(self):
        """Return dataset maxzoom."""
        return self.dataset.maxzoom

    def close(self):
        """Close rasterio dataset."""
        self._ctx_stack.close()

    def __exit__(self, exc_type, exc_value, traceback):
        """Support using with Context Managers."""
        self.close()

    def info(self) -> Info:
        # could return self.dataset.info()
        raise NotImplementedError

    def statistics(self, *args, **kwargs) -> Dict[str, BandStatistics]:
        # could return self.dataset.statistics(*args, **kwargs)
        raise NotImplementedError

    def tile(
        self,
        tile_x: int,
        tile_y: int,
        tile_z: int,
        tilesize: int = 256,
        indexes: Optional[Indexes] = None,
        expression: Optional[str] = None,
        **kwargs: Any,
    ) -> ImageData:
        img = self.dataset.tile(
            tile_x,
            tile_y,
            tile_z,
            tilesize,
            indexes=indexes,
            expression=expression,
            **kwargs,
        )

        mask = self.mask.tile(
            tile_x,
            tile_y,
            tile_z,
            tilesize,
            indexes=1,
            **kwargs,
        )

        # edited from original response
        # img.mask = mask.array > self.input["max_correlation"]
        img.array.mask = mask.array > self.input["max_correlation"]
        
        return mask

    def part(self, bbox: BBox) -> ImageData:
        raise NotImplementedError

    def preview(self) -> ImageData:
        raise NotImplementedError

    def point(self, lon: float, lat: float) -> PointData:
        raise NotImplementedError

    def feature(self, shape: Dict) -> ImageData:
        raise NotImplementedError

    def read(
        self,
        indexes: Optional[Indexes] = None,
        expression: Optional[str] = None,
        **kwargs: Any,
    ) -> ImageData:
        raise NotImplementedError

and the dependency

def InputDependency(
    url: Annotated[str, Query(description="Dataset URL")],
    mask: Annotated[str, Query(description="Mask URL")],
    max_correlation: Annotated[float, Query(description="Max Correlation")] = 0.3,
) -> Dict:
    """Create dataset path from args"""
    return {
        "data": url,
        "mask": mask,
        "max_correlation": max_correlation,
    }

[scottstanie]

I think the slowest part for me was realizing that your raise NotImplementedError was not "I leave this as an exercise to the reader", but rather that it's not needed haha

Basically the only functional change I had to make was creating a new ImageData object, since it yelled at me for doing img.mask = ... with ".mask has no setter".

thank you!

[vincentsarago]

Basically the only functional change I had to make was creating a new ImageData object, since it yelled at me for doing img.mask = ... with ".mask has no setter".

🤦 oh sorry about this, I fixed the example 🙏

[scottstanie]

Actually, while I can't yet just link to the whole thing, the real slowest part for me was

1. confusing the inverted numpy/rasterio mask formats
2. realizing that I should use both .masks from the image, as well as the new layer you're using to hide pixels:

            # Combine the invalid pixels of `.mask` and `.img`
            bad_image_pixels = img.mask == 0
            # img.mask has "valid" = 255, "invalid" = 0
            # need to convert it to numpy style
            filled_mask_layer = np.squeeze(mask_layer_img.array.filled(0))
            bad_mask_pixels = filled_mask_layer == 0

            # Also mask out the pixels which don't pass the threshold
            pixels_below_threshold = filled_mask_layer < self.input["min_correlation"]
            combined_mask = np.logical_or.reduce(
                [
                    bad_image_pixels,
                    bad_mask_pixels,
                    pixels_below_threshold,
                ]
            )
            img = ImageData(
                np.ma.MaskedArray(img.data, mask=combined_mask),
                assets=img.assets,
                crs=img.crs,
                bounds=img.bounds,
            )

(if this is doing something in an obviously dumb way please let me know)