create_mask.py

#!/usr/bin/env python

import click
import numpy as np
import rasterio as rio


@click.command('append-mask')
@click.argument('srcpath', type=click.Path(exists=True))
@click.argument('dstpath', type=click.Path(exists=False))
@click.option('--nodata', default=0, type=click.INT)
def append_mask(srcpath, dstpath, nodata):
    """
    Create a mask from a specified nodata value.
    """

    n_tile_x = 10
    n_tile_y = 10

    with rio.drivers():
        with rio.open(srcpath, 'r') as src:

            profile = src.profile
            profile.update(count=4)

            width = src.width
            height = src.height
            dtype = src.dtypes[0]
            dtype_max = np.iinfo(dtype).max

            w = int(np.round(width / n_tile_x))
            h = int(np.round(height / n_tile_y))

            with rio.open(dstpath, 'w', **profile) as dst:

                for j in range(0, n_tile_y):
                    for i in range(0, n_tile_x):

                        y0, y1 = j * h, (j + 1) * h
                        x0, x1 = i * w, (i + 1) * w

                        y1 = np.minimum(height, y1)
                        x1 = np.minimum(width, x1)

                        window = ((y0, y1), (x0, x1))

                        r = src.read(1, window=window)
                        g = src.read(2, window=window)
                        b = src.read(3, window=window)

                        indices = np.where((r == 0) & (g == 0) & (b == 0))

                        mask = np.ones_like(r, dtype=dtype)
                        mask[indices] = 0
                        mask *= dtype_max
                        
                        dst.write_band(1, r, window=window)
                        dst.write_band(2, g, window=window)
                        dst.write_band(3, b, window=window)
                        dst.write_band(4, mask, window=window)
                        print window

if __name__ == '__main__':
    append_mask()