OpenDroneMap-WebODM/app/raster_utils.py

import rasterio
import re
import logging
import os
import subprocess
import numpy as np
import numexpr as ne
from rasterio.enums import ColorInterp
from rio_tiler.utils import has_alpha_band, linear_rescale
from rio_tiler.colormap import cmap as colormap, apply_cmap
from rio_tiler.errors import InvalidColorMapName
from app.api.hsvblend import hsv_blend
from app.api.hillshade import LightSource
from rio_tiler.io import COGReader
from rasterio.warp import calculate_default_transform, reproject, Resampling

logger = logging.getLogger('app.logger')

ZOOM_EXTRA_LEVELS = 3

def extension_for_export_format(export_format):
    extensions = {
        'gtiff': 'tif',
        'gtiff-rgb': 'tif',
    }
    return extensions.get(export_format, export_format)

def export_raster(input, output, **opts):
    epsg = opts.get('epsg')
    expression = opts.get('expression')
    export_format = opts.get('format')
    rescale = opts.get('rescale')
    color_map = opts.get('color_map')
    hillshade = opts.get('hillshade')
    asset_type = opts.get('asset_type')
    name = opts.get('name', 'raster') # KMZ specific

    dem = asset_type in ['dsm', 'dtm']
    
    with COGReader(input) as ds_src:
        src = ds_src.dataset
        profile = src.meta.copy()

        # Output format
        driver = "GTiff"
        compress = None
        max_bands = 9999
        with_alpha = True
        rgb = False
        indexes = src.indexes
        output_raster = output
        jpg_background = 255 # white

        # KMZ is special, we just export it as PNG with EPSG:4326
        # and then call GDAL to tile/package it
        kmz = export_format == "kmz"
        if kmz:
            export_format = "png"
            epsg = 4326
            path_base, _ = os.path.splitext(output)
            output_raster = path_base + ".png"

        if export_format == "jpg":
            driver = "JPEG"
            profile.update(quality=90)
            band_count = 3
            with_alpha = False
            rgb = True
        elif export_format == "png":
            driver = "PNG"
            band_count = 4
            rgb = True
        elif export_format == "gtiff-rgb":
            compress = "JPEG"
            profile.update(jpeg_quality=90)
            band_count = 4
            rgb = True
        else:
            compress = "DEFLATE"
            band_count = src.count
        
        if compress is not None:
            profile.update(compress=compress)
            profile.update(predictor=2 if compress == "DEFLATE" else 1)

        if rgb and rescale is None:
            # Compute min max
            nodata = None
            if asset_type == 'orthophoto':
                nodata = 0
            md = ds_src.metadata(pmin=2.0, pmax=98.0, hist_options={"bins": 255}, nodata=nodata)
            rescale = [md['statistics']['1']['min'], md['statistics']['1']['max']]

        ci = src.colorinterp

        if rgb and expression is None:
            # More than 4 bands?
            if len(ci) > 4:
                # Try to find RGBA band order
                if ColorInterp.red in ci and \
                        ColorInterp.green in ci and \
                        ColorInterp.blue in ci and \
                        ColorInterp.alpha in ci:
                    indexes = (ci.index(ColorInterp.red) + 1,
                                ci.index(ColorInterp.green) + 1,
                                ci.index(ColorInterp.blue) + 1,
                                ci.index(ColorInterp.alpha) + 1)

        if ColorInterp.alpha in ci:
            mask = src.read(ci.index(ColorInterp.alpha) + 1)
        else:
            mask = src.dataset_mask()

        cmap = None
        if color_map:
            try:
                cmap = colormap.get(color_map)
            except InvalidColorMapName:
                logger.warning("Invalid colormap {}".format(color_map))


        def process(arr, skip_rescale=False, skip_alpha=False, skip_type=False):
            if not skip_rescale and rescale is not None:
                arr = linear_rescale(arr, in_range=rescale)
            if not skip_alpha and not with_alpha:
                arr[mask==0] = jpg_background
            if not skip_type and rgb and arr.dtype != np.uint8:
                arr = arr.astype(np.uint8)

            return arr

        def update_rgb_colorinterp(dst):
            if with_alpha:
                dst.colorinterp = [ColorInterp.red, ColorInterp.green, ColorInterp.blue, ColorInterp.alpha]
            else:
                dst.colorinterp = [ColorInterp.red, ColorInterp.green, ColorInterp.blue]

        profile.update(driver=driver, count=band_count)
        if rgb:
            profile.update(dtype=rasterio.uint8)

        if dem and rgb and profile.get('nodata') is not None:
            profile.update(nodata=None)

        # Define write band function
        # Reprojection needed?
        if src.crs is not None and epsg is not None and src.crs.to_epsg() != epsg:
            dst_crs = "EPSG:{}".format(epsg)

            transform, width, height = calculate_default_transform(
                src.crs, dst_crs, src.width, src.height, *src.bounds)

            profile.update(
                crs=dst_crs,
                transform=transform,
                width=width,
                height=height
            )

            def write_band(arr, dst, band_num):
                reproject(source=arr, 
                        destination=rasterio.band(dst, band_num),
                        src_transform=src.transform,
                        src_crs=src.crs,
                        dst_transform=transform,
                        dst_crs=dst_crs,
                        resampling=Resampling.nearest)

        else:
            # No reprojection needed
            def write_band(arr, dst, band_num):
                dst.write(arr, band_num)

        if expression is not None:
            # Apply band math
            if rgb:
                profile.update(dtype=rasterio.uint8, count=band_count)
            else:
                profile.update(dtype=rasterio.float32, count=1, nodata=-9999)

            bands_names = ["b{}".format(b) for b in tuple(sorted(set(re.findall(r"b(?P<bands>[0-9]{1,2})", expression))))]
            rgb_expr = expression.split(",")
            indexes = tuple([int(b.replace("b", "")) for b in bands_names])

            alpha_index = None
            if has_alpha_band(src):
                try:
                    alpha_index = src.colorinterp.index(ColorInterp.alpha) + 1
                    indexes += (alpha_index, )
                except ValueError:
                    pass

            data = src.read(indexes=indexes, out_dtype=np.float32)
            arr = dict(zip(bands_names, data))
            arr = np.array([np.nan_to_num(ne.evaluate(bloc.strip(), local_dict=arr)) for bloc in rgb_expr])

            # Set nodata values
            index_band = arr[0]
            if alpha_index is not None:
                # -1 is the last band = alpha
                index_band[data[-1] == 0] = -9999

            # Remove infinity values
            index_band[index_band>1e+30] = -9999
            index_band[index_band<-1e+30] = -9999

            # Make sure this is float32
            arr = arr.astype(np.float32)

            with rasterio.open(output_raster, 'w', **profile) as dst:
                # Apply colormap?
                if rgb and cmap is not None:
                    rgb_data, _ = apply_cmap(process(arr, skip_alpha=True), cmap)

                    band_num = 1
                    for b in rgb_data:
                        write_band(process(b, skip_rescale=True), dst, band_num)
                        band_num += 1

                    if with_alpha:
                        write_band(mask, dst, band_num)
                    
                    update_rgb_colorinterp(dst)
                else:
                    # Raw
                    write_band(process(arr)[0], dst, 1)
        elif dem:
            # Apply hillshading, colormaps to elevation
            with rasterio.open(output_raster, 'w', **profile) as dst:
                arr = src.read()

                intensity = None
                if hillshade is not None and hillshade > 0:
                    delta_scale = (ZOOM_EXTRA_LEVELS + 1) * 4
                    dx = src.meta["transform"][0] * delta_scale
                    dy = -src.meta["transform"][4] * delta_scale
                    ls = LightSource(azdeg=315, altdeg=45)
                    intensity = ls.hillshade(arr[0], dx=dx, dy=dy, vert_exag=hillshade)
                    intensity = intensity * 255.0

                # Apply colormap?
                if rgb and cmap is not None:
                    rgb_data, _ = apply_cmap(process(arr, skip_alpha=True), cmap)

                    if intensity is not None:
                        rgb_data = hsv_blend(rgb_data, intensity)
                        
                    band_num = 1
                    for b in rgb_data:
                        write_band(process(b, skip_rescale=True), dst, band_num)
                        band_num += 1
                    
                    if with_alpha:
                        write_band(mask, dst, band_num)

                    update_rgb_colorinterp(dst)
                else:
                    # Raw
                    write_band(process(arr)[0], dst, 1)
        else:
            # Copy bands as-is
            with rasterio.open(output_raster, 'w', **profile) as dst:
                band_num = 1
                for idx in indexes:
                    ci = src.colorinterp[idx - 1]
                    arr = src.read(idx)

                    if ci == ColorInterp.alpha:
                        if with_alpha:
                            write_band(arr, dst, band_num)
                            band_num += 1
                    else:
                        write_band(process(arr), dst, band_num)
                        band_num += 1
                
                new_ci = [src.colorinterp[idx - 1] for idx in indexes]
                if not with_alpha:
                    new_ci = [ci for ci in new_ci if ci != ColorInterp.alpha]
                    
                dst.colorinterp = new_ci
                
        if kmz:
            subprocess.check_output(["gdal_translate", "-of", "KMLSUPEROVERLAY", 
                                        "-co", "Name={}".format(name),
                                        "-co", "FORMAT=PNG", output_raster, output])