Python rasterio.Env() Examples

def test_tile_read_extmask():
    """Read masked area."""
    # non-boundless tile covering the masked part
    mercator_tile = mercantile.Tile(x=876431, y=1603669, z=22)
    bounds = mercantile.xy_bounds(mercator_tile)
    with rasterio.Env(GDAL_DISABLE_READDIR_ON_OPEN="TRUE"):
        with rasterio.open(S3_EXTMASK_PATH) as src_dst:
            arr, mask = reader.part(src_dst, bounds, 256, 256)
        assert arr.shape == (3, 256, 256)
        assert mask.shape == (256, 256)
        assert not mask.all()

    # boundless tile covering the masked part
    mercator_tile = mercantile.Tile(x=876431, y=1603668, z=22)
    bounds = mercantile.xy_bounds(mercator_tile)
    with rasterio.Env(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR"):
        with rasterio.open(S3_MASK_PATH) as src_dst:
            arr, mask = reader.part(src_dst, bounds, 256, 256)
        assert arr.shape == (3, 256, 256)
        assert not mask.all() 

def test_tile_read_mask():
    """Read masked area."""
    with rasterio.Env(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR"):
        # non-boundless tile covering the masked part
        with rasterio.open(S3_MASK_PATH) as src_dst:
            arr, mask = reader.tile(src_dst, 876431, 1603669, 22, tilesize=16)
        assert arr.shape == (3, 16, 16)
        assert mask.shape == (16, 16)
        assert not mask.all()

        # boundless tile covering the masked part
        with rasterio.open(S3_MASK_PATH) as src_dst:
            arr, mask = reader.tile(src_dst, 876431, 1603668, 22, tilesize=256)
        assert arr.shape == (3, 256, 256)
        assert not mask.all() 

def read_tile(src_path, tile):
    """Benchmark rio-tiler.utils._tile_read."""
    tile_bounds = mercantile.xy_bounds(tile)
    # We make sure to not store things in cache.
    with rasterio.Env(GDAL_CACHEMAX=0, NUM_THREADS="all"):
        with rasterio.open(src_path) as src_dst:
            return reader.part(
                src_dst,
                tile_bounds,
                256,
                256,
                resampling_method="nearest",
                dst_crs=constants.WEB_MERCATOR_CRS,
            ) 

def test_validate_external_overview(tmpdir):
    import os
    from terracotta import cog

    outfile = str(tmpdir / 'raster.tif')
    raster_data = 1000 * np.random.rand(512, 512).astype(np.uint16)

    profile = BASE_PROFILE.copy()
    profile.update(
        height=raster_data.shape[0],
        width=raster_data.shape[1],
        tiled=True,
        blockxsize=256,
        blockysize=256
    )

    with rasterio.Env(TIFF_USE_OVR=True):
        with rasterio.open(outfile, 'w', **profile) as dst:
            dst.write(raster_data, 1)

            overviews = [2 ** j for j in range(1, 4)]
            dst.build_overviews(overviews, Resampling.nearest)

        assert os.path.isfile(f'{outfile}.ovr')

    assert not cog.validate(outfile) 

def read_raster_no_crs(input_file, indexes=None, gdal_opts=None):
    """
    Wrapper function around rasterio.open().read().

    Parameters
    ----------
    input_file : str
        Path to file
    indexes : int or list
        Band index or list of band indexes to be read.
    gdal_opts : dict
        GDAL options passed on to rasterio.Env()

    Returns
    -------
    MaskedArray

    Raises
    ------
    FileNotFoundError if file cannot be found.
    """
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        try:
            with rasterio.Env(
                **get_gdal_options(
                    gdal_opts, is_remote=path_is_remote(input_file, s3=True)
                )
            ):
                with rasterio.open(input_file, "r") as src:
                    return src.read(indexes=indexes, masked=True)
        except RasterioIOError as e:
            try:
                if path_exists(input_file):
                    raise e
            except:
                raise e
            raise FileNotFoundError("%s not found" % input_file) 

def rasterio_write(path, array, profile={}, tags={}):
    """
    Write a numpy array in a tiff or png file with rasterio.

    Args:
        path (str): path to the output tiff/png file
        array (numpy array): 2D or 3D array containing the image to write.
        profile (dict): rasterio profile (ie dictionary of metadata)
        tags (dict): dictionary with additional geotiff tags
    """
    # determine the driver based on the file extension
    extension = os.path.splitext(path)[1].lower()
    if extension in ['.tif', '.tiff']:
        driver = 'GTiff'
    elif extension in ['.png']:
        driver = 'png'
    else:
        raise NotImplementedError('format {} not supported'.format(extension))

    # read image size and number of bands
    array = np.atleast_3d(array)
    height, width, nbands = array.shape

    # define image metadata dict
    profile.update(driver=driver, count=nbands, width=width, height=height,
                   dtype=array.dtype)

    # write to file
    with rasterio.Env():
        with rasterio.open(path, 'w', **profile) as dst:
            dst.write(np.transpose(array, (2, 0, 1)))
            dst.update_tags(**tags) 

def test_upsample(test_data):

    with rasterio.Env():

        pan, rgb, src_aff, src_crs, dst_aff, dst_crs = test_data
        up_rgb = utils._upsample(rgb, pan.shape, src_aff,
                                 src_crs, dst_aff, dst_crs)

        # test upsampled shape
        assert up_rgb.shape[0] == 3
        assert up_rgb.shape[1] / rgb.shape[1] == 2
        assert up_rgb.shape[2] / rgb.shape[2] == 2

        # test upsampled dtype
        assert up_rgb.dtype == np.uint8

        # test for seams. 1 px wide row/column with 0's
        assert np.all((up_rgb[0].max(axis=0)) != 0)
        assert np.all((up_rgb[0].max(axis=1)) != 0)
        assert np.all((up_rgb[1].max(axis=0)) != 0)
        assert np.all((up_rgb[1].max(axis=1)) != 0)
        assert np.all((up_rgb[2].max(axis=0)) != 0)
        assert np.all((up_rgb[2].max(axis=1)) != 0)

        # test upsampled values from reproject function
        assert up_rgb[0][0][0] == rgb[0][0][0]
        assert up_rgb[-1][-1][-1] == rgb[-1][-1][-1]


# Testing Brovey function from pansharp_methods 

def cloud_optimize(raster_file, outfile, create_mask=False, remove_nodata=False):
    import math
    import contextlib
    import rasterio
    import rasterio.io
    import rasterio.shutil

    COG_PROFILE = {
        'count': 1,
        'driver': 'GTiff',
        'interleave': 'pixel',
        'tiled': True,
        'blockxsize': 256,
        'blockysize': 256,
        'compress': 'DEFLATE',
        'photometric': 'MINISBLACK',
        'BIGTIFF': 'IF_SAFER'
    }

    with contextlib.ExitStack() as es:
        es.enter_context(rasterio.Env(
            GDAL_TIFF_INTERNAL_MASK=True,
            GDAL_TIFF_OVR_BLOCKSIZE=256,
        ))
        src = es.enter_context(rasterio.open(str(raster_file)))

        profile = src.profile.copy()
        profile.update(COG_PROFILE)

        if remove_nodata:
            profile['nodata'] = None

        memfile = es.enter_context(rasterio.io.MemoryFile())
        dst = es.enter_context(memfile.open(**profile))

        dst.write(src.read())

        if create_mask:
            dst.write_mask(src.dataset_mask().astype('uint8'))

        max_overview_level = math.ceil(math.log2(max(
            dst.height // profile['blockysize'],
            dst.width // profile['blockxsize']
        )))

        overviews = [2 ** j for j in range(1, max_overview_level + 1)]
        rs_method = rasterio.enums.Resampling.nearest
        dst.build_overviews(overviews, rs_method)
        rasterio.shutil.copy(dst, str(outfile), copy_src_overviews=True, **COG_PROFILE) 

def read_raster_window(
    input_files,
    tile,
    indexes=None,
    resampling="nearest",
    src_nodata=None,
    dst_nodata=None,
    gdal_opts=None
):
    """
    Return NumPy arrays from an input raster.

    NumPy arrays are reprojected and resampled to tile properties from input
    raster. If tile boundaries cross the antimeridian, data on the other side
    of the antimeridian will be read and concatenated to the numpy array
    accordingly.

    Parameters
    ----------
    input_files : string or list
        path to a raster file or list of paths to multiple raster files readable by
        rasterio.
    tile : Tile
        a Tile object
    indexes : list or int
        a list of band numbers; None will read all.
    resampling : string
        one of "nearest", "average", "bilinear" or "lanczos"
    src_nodata : int or float, optional
        if not set, the nodata value from the source dataset will be used
    dst_nodata : int or float, optional
        if not set, the nodata value from the source dataset will be used
    gdal_opts : dict
        GDAL options passed on to rasterio.Env()

    Returns
    -------
    raster : MaskedArray
    """
    with rasterio.Env(
        **get_gdal_options(
            gdal_opts,
            is_remote=path_is_remote(
                input_files[0] if isinstance(input_files, list) else input_files, s3=True
            ) if isinstance(input_files, str) else False
        )
    ) as env:
        logger.debug("reading %s with GDAL options %s", input_files, env.options)
        return _read_raster_window(
            input_files,
            tile,
            indexes=indexes,
            resampling=resampling,
            src_nodata=src_nodata,
            dst_nodata=dst_nodata
        ) 

def test_reproject():
    from rasterio.warp import reproject
    from rasterio.enums import Resampling

    with rasterio.Env():
        # As source: a 1024 x 1024 raster centered on 0 degrees E and 0
        # degrees N, each pixel covering 15".
        rows, cols = src_shape = (1024, 1024)
        # decimal degrees per pixel
        d = 1.0 / 240

        # The following is equivalent to
        # A(d, 0, -cols*d/2, 0, -d, rows*d/2).
        src_transform = rasterio.Affine.translation(
                    -cols*d/2,
                    rows*d/2) * rasterio.Affine.scale(d, -d)
        src_crs = {'init': 'EPSG:4326'}
        source = np.ones(src_shape, np.uint8) * 255

        # Destination: a 2048 x 2048 dataset in Web Mercator (EPSG:3857)
        # with origin at 0.0, 0.0.
        dst_shape = (2048, 2048)
        dst_transform = Affine.from_gdal(
            -237481.5, 425.0, 0.0, 237536.4, 0.0, -425.0)
        dst_crs = {'init': 'EPSG:3857'}
        destination = np.zeros(dst_shape, np.uint8)

        reproject(
            source,
            destination,
            src_transform=src_transform,
            src_crs=src_crs,
            dst_transform=dst_transform,
            dst_crs=dst_crs,
            resampling=Resampling.nearest)

        # Assert that the destination is only partly filled.
        assert destination.any()
        assert not destination.all()


# Testing upsample function