""" Utilities for extracting training data from region of interests (ROI)
"""
import numpy as np
from rasterio.features import rasterize
from shapely.geometry import shape as geom_shape


def extract_roi(src, features, feature_prop=None, all_touched=False, fill=0):
    """ Yield pixel data from ``src`` for ROIs in ``features``

    Args:
        src (rasterio.RasterReader): The ``rasterio`` dataset used to extract
            training data values from
        features (list[dict]): A list of features from a polygon vector file
            given in the format used by fiona
        feature_prop (str): The name of the attribute from ``features``
            containing the ROI labels
        all_touched (bool): Rasterization option that decides if all pixels
            touching the ROI should be included, or just pixels from within
            the ROI
        fill (int, float): A fill value for the ROI rasterization. This fill
            value should not be the same as any of the labels in the set of
            ROIs

    Returns:
        tuple (np.ndarray, np.ndarray, np.ndarray, np.ndarray): A tuple
            containing an array of ROI data from ``src`` (``band x n``), the
            ROI data label (``n``), and the X and Y coordinates of each data
            point (``n`` and ``n`` sized)

    """
    if not feature_prop:
        feature_prop = list(features[0]['properties'].keys())[0]

    for feat, label in features:
        geom = geom_shape(feat)
        bounds = tuple(geom.bounds)

        window = src.window(*bounds, boundless=True)
        data = src.read(window=window, boundless=True)
        shape = data.shape
        transform = src.window_transform(window)

        roi = rasterize(
            [feat],
            out_shape=shape[1:],
            transform=transform,
            fill=fill,
            all_touched=all_touched
        )

        mask = np.logical_or(
            (data == src.nodata).any(axis=0),
            roi == fill
        )
        masked = np.ma.MaskedArray(
            data,
            mask=np.ones_like(data) * mask
        )

        ys, xs = np.where(~mask)
        coord_xs, coord_ys = transform * (xs, ys)

        masked = masked.compressed()
        npix = masked.size / shape[0]
        masked = masked.reshape((shape[0], npix))

        label = label * np.ones(coord_ys.size, dtype=np.uint8)

        yield (masked, label, coord_xs, coord_ys, )

Add ROI extraction code for use w/ xarray
Add ROI extraction code for use w/ xarray
