Source code for rio_vrt.vrt

"""Rasterio based vrt creation."""

import xml.etree.cElementTree as ET
from os.path import relpath
from pathlib import Path
from statistics import mean
from typing import List, Tuple, Union
from xml.dom import minidom

import rasterio as rio
from rasterio.enums import ColorInterp

from .enums import resolutions, types


def _add_source_content(
    Source: ET.Element, src: rio.DatasetReader, type: str, xoff: str, yoff: str
) -> None:
    """Add the content of a sourcefile in xml."""
    width, height = str(src.width), str(src.height)
    blockx = str(src.profile.get("blockxsize", ""))
    blocky = str(src.profile.get("blockysize", ""))

    attr = {
        "RasterXSize": width,
        "RasterYSize": height,
        "DataType": type,
    }

    # optional attributes
    if blockx and blocky:
        attr["BlockXSize"], attr["BlockYSize"] = blockx, blocky

    ET.SubElement(Source, "SourceProperties", attr)

    attr = {"xOff": "0", "yOff": "0", "xSize": width, "ySize": height}
    ET.SubElement(Source, "SrcRect", attr)

    attr = {"xOff": xoff, "yOff": yoff, "xSize": width, "ySize": height}
    ET.SubElement(Source, "DstRect", attr)


[docs] def build_vrt( vrt_path: Union[str, Path], files: List[Union[str, Path]], relative: bool = False, mosaic: bool = True, res: Union[str, Tuple[float, float]] = "average", ) -> Path: """Create a vrt file from multiple files. Arguments: vrt_path: the final vrt file files: a list of rasterio readable files relative: use a path relative to the vrt file. The files path must be relative to the vrt. mosaic: The method to use to gather images in the vrt. ``MOSAIC`` (True) will mosaic each band of each image together. ``STACK`` (False) will create one band for each file using the first band of each file.* res: The resolution to use in the vrt geotransform. You can use a string (average, highest or lowest) or use a defined tuple of values (xres, yres). Returns: the path to the vrt file """ # transform the final file in Path vrt_path = Path(vrt_path).resolve() # transform all the file path into Path objects files = [Path(f).resolve() for f in files] # cannot do anything if there are no files if len(files) == 0: raise ValueError("There should be at least 1 file to create a vrt.") # check the res value if isinstance(res, str) and res not in resolutions: raise ValueError( 'the provided resolution cannot be use: "{res}", please use one of the existing keywords' ) # read global informations from the first file with rio.open(files[0]) as f: crs = f.crs count = f.count dtypes = f.dtypes colorinterps = f.colorinterp indexes = f.indexes nodatavals = f.nodatavals # for stacks replace count and indexes as we only take the first band if not mosaic: count, indexes = len(files), [1] # sanity checks for file in files: with rio.open(file) as f: if f.crs != crs: raise ValueError( f'the crs ({f.crs}) from file "{file}" is not corresponding to the global one ({crs})' ) if mosaic and f.count != count: raise ValueError( f'the crs ({f.count}) from file "{file}" is not corresponding to the global one ({count})' ) # read all files to extract information on the spatial extend of the vrt left_, bottom_, right_, top_, xres_, yres_ = [], [], [], [], [], [] for file in files: with rio.open(file) as f: xres_.append(f.res[0]) yres_.append(f.res[0]) left_.append(f.bounds.left) right_.append(f.bounds.right) top_.append(f.bounds.top) bottom_.append(f.bounds.bottom) # get the spatial extend of the dataset left = min(*left_) bottom = min(*bottom_) right = max(*right_) top = max(*top_) # get the resolution if res == "highest": xres, yres = max(*xres_), max(*yres_) elif res == "lowest": xres, yres = min(*xres_), min(*yres_) elif res == "average": xres, yres = mean(xres_), mean(yres_) elif isinstance(res, tuple): xres, yres = res # rebuild the affine transformation from gathered information along with total bounds # negative y_res as we start from the top-left corner transform = rio.Affine.from_gdal(left, xres, 0, top, 0, -yres) total_width = round((right - left) / xres) total_height = round((top - bottom) / yres) # start the tree attr = {"rasterXSize": str(total_width), "rasterYSize": str(total_height)} VRTDataset = ET.Element("VRTDataset", attr) # don't know how to extract dataAxisToSRSAxisMapping # https://gis.stackexchange.com/questions/458781/how-to-get-dataaxistosrsaxismapping-from-an-image # revert to OAMS_TRADITIONAL_GIS_ORDER until then ET.SubElement(VRTDataset, "SRS").text = crs.wkt text = ", ".join([str(i) for i in transform.to_gdal()]) ET.SubElement(VRTDataset, "GeoTransform").text = text ET.SubElement(VRTDataset, "OverviewList", {"resampling": "nearest"}).text = "2 4 8" # add the rasterbands # mosaicking create 1 band for each band of the images and add al the fils as # simple sources along with color informations if mosaic: VRTRasterBands_dict = {} for i in indexes: attr = {"dataType": types[dtypes[i - 1]], "band": str(i)} VRTRasterBands_dict[i] = ET.SubElement(VRTDataset, "VRTRasterBand", attr) ET.SubElement(VRTRasterBands_dict[i], "Offset").text = "0.0" ET.SubElement(VRTRasterBands_dict[i], "Scale").text = "1.0" if colorinterps[i - 1] != ColorInterp.undefined: color = colorinterps[i - 1].name.capitalize() ET.SubElement(VRTRasterBands_dict[i], "ColorInterp").text = color if nodatavals[i - 1] is not None: text = str(nodatavals[i - 1]) ET.SubElement(VRTRasterBands_dict[i], "NoDataValue").text = text # add the files for f in files: relativeToVRT = "1" if relative is True else "0" with rio.open(f) as src: for i in indexes: is_alpha = colorinterps[i - 1] == ColorInterp.alpha has_nodata = nodatavals[i - 1] is not None source_type = ( "ComplexSource" if is_alpha or has_nodata else "SimpleSource" ) Source = ET.SubElement(VRTRasterBands_dict[i], source_type) attr = {"relativeToVRT": relativeToVRT} text = str(f) if not relative else relpath(f, vrt_path.parent) ET.SubElement(Source, "SourceFilename", attr).text = text ET.SubElement(Source, "SourceBand").text = str(i) _add_source_content( Source=Source, src=src, type=types[dtypes[i - 1]], xoff=str(abs(round((src.bounds.left - left) / xres))), yoff=str(abs(round((src.bounds.top - top) / yres))), ) if nodatavals[i - 1] is not None: text = str(nodatavals[i - 1]) ET.SubElement(Source, "NODATA").text = text if colorinterps[i - 1] == ColorInterp.alpha: ET.SubElement(Source, "UseMaskBand").text = "true" # in stacked vrt, each file is added as a single band and only the first band is # considered. They are all complex sources to make sure GIS softwares don't do funny # display upon reading elif not mosaic: for i, f in enumerate(files): attr = {"dataType": types[dtypes[0]], "band": str(i)} VRTRasterBands = ET.SubElement(VRTDataset, "VRTRasterBand", attr) ComplexSource = ET.SubElement(VRTRasterBands, "ComplexSource") relativeToVRT = "1" if relative is True else "0" attr = {"relativeToVRT": relativeToVRT} text = str(f) if not relative else relpath(f, vrt_path.parent) ET.SubElement(ComplexSource, "SourceFilename", attr).text = text ET.SubElement(ComplexSource, "SourceBand").text = "1" with rio.open(f) as src: _add_source_content( Source=ComplexSource, src=src, type=types[dtypes[0]], xoff=str(abs(round((src.bounds.left - left) / xres))), yoff=str(abs(round((src.bounds.top - top) / yres))), ) # write the file vrt_path.resolve().write_text( minidom.parseString(ET.tostring(VRTDataset).decode("utf-8")) .toprettyxml(indent=" ") .replace(""", '"') ) return vrt_path