Handling Out-of-Memory Errors in Large Raster Mosaics

pip install rasterio numpy
# GDAL Python bindings are usually installed alongside GDAL:
pip install gdal==$(gdal-config --version)

# mosaic_stream.py
import argparse
import gc
import os
import tempfile
from pathlib import Path

import numpy as np
import rasterio
from osgeo import gdal
from rasterio.windows import Window


def set_gdal_cache(fraction: float = 0.25) -> None:
    """Cap GDAL's block cache to a fraction of available RAM.

    gdal.SetCacheMax() takes bytes. Must be called before opening any dataset.
    """
    import psutil  # optional but recommended; falls back to a 256 MB hard cap

    try:
        available = psutil.virtual_memory().available
        gdal.SetCacheMax(int(available * fraction))
    except ImportError:
        gdal.SetCacheMax(256 * 1024 * 1024)  # 256 MB fallback


def build_vrt(inputs: list[str], vrt_path: str) -> None:
    """Unify source rasters into a GDAL VRT.

    resolution="highest" selects the finest GSD among all inputs.
    resampleAlg="nearest" avoids blurring discrete raster classes (e.g. land-cover).
    FlushCache() + assigning None closes the C-level dataset handle immediately.
    """
    vrt = gdal.BuildVRT(
        vrt_path,
        inputs,
        resolution="highest",
        resampleAlg="nearest",
    )
    if vrt is None:
        raise RuntimeError(f"gdal.BuildVRT failed; check inputs exist and share a CRS: {inputs}")
    vrt.FlushCache()
    vrt = None  # Release C handle; do NOT rely on Python's GC timing here


def stream_mosaic(
    vrt_path: str,
    output_path: str,
    chunk_size: int = 1024,
) -> None:
    """Write the mosaic to disk one tile at a time.

    Peak RAM at any point = chunk_size * chunk_size * bands * dtype_bytes.
    For a 1024-pixel, 3-band uint16 tile: 1024 * 1024 * 3 * 2 = ~6 MB.
    """
    with rasterio.open(vrt_path) as src:
        meta = src.meta.copy()
        meta.update(
            driver="GTiff",
            compress="LZW",
            tiled=True,
            blockxsize=chunk_size,   # Output block size matches read chunk_size —
            blockysize=chunk_size,   # misalignment causes read-amplification on re-access.
            interleave="band",
            dtype=src.dtypes[0],
        )

        with rasterio.open(output_path, "w", **meta) as dst:
            for col_off in range(0, src.width, chunk_size):
                for row_off in range(0, src.height, chunk_size):
                    # Clamp window to dataset boundary; avoids out-of-bounds reads
                    # at the right and bottom edges where the extent may not divide evenly.
                    w = min(chunk_size, src.width - col_off)
                    h = min(chunk_size, src.height - row_off)
                    window = Window(col_off, row_off, w, h)

                    data = src.read(window=window)   # lazy read: only this tile in RAM
                    dst.write(data, window=window)   # direct disk write; no accumulation

                    del data          # Explicit dereference
                    gc.collect()      # Force immediate reclamation; prevents heap fragmentation


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Stream-merge large rasters to a GeoTIFF without OOM."
    )
    parser.add_argument("inputs", nargs="+", help="Input raster paths (GeoTIFF or any GDAL source)")
    parser.add_argument("-o", "--output", required=True, help="Output GeoTIFF path")
    parser.add_argument(
        "-c", "--chunk",
        type=int,
        default=1024,
        help="Tile size in pixels (default: 1024; match to storage block size)",
    )
    parser.add_argument(
        "--cache-fraction",
        type=float,
        default=0.25,
        help="Fraction of available RAM to allocate for GDAL block cache (default: 0.25)",
    )
    args = parser.parse_args()

    set_gdal_cache(args.cache_fraction)

    # tempfile.mkstemp() avoids the TOCTOU race of the deprecated mktemp().
    # os.close() shuts the file descriptor immediately; build_vrt will overwrite the file.
    fd, vrt_path = tempfile.mkstemp(suffix=".vrt")
    os.close(fd)
    try:
        build_vrt(args.inputs, vrt_path)
        stream_mosaic(vrt_path, args.output, args.chunk)
    finally:
        Path(vrt_path).unlink(missing_ok=True)


if __name__ == "__main__":
    main()

vrt = gdal.BuildVRT(
    vrt_path,
    inputs,
    outputSRS="EPSG:32633",      # target CRS: WGS 84 / UTM zone 33N
    resolution="highest",
    resampleAlg="bilinear",       # bilinear for continuous data; nearest for categorical
    VRTNodata=0,
)

# 1. Check CRS, extent, and band count match expectations
gdalinfo -stats output_mosaic.tif | grep -E "Size|SRS|Band|Min|Max"

# 2. Verify nodata regions are at the edges, not scattered through valid pixels
# (scattered nodata often indicates CRS mismatch or misaligned block writes)
gdal_translate -of GTiff -b 1 output_mosaic.tif /vsistdout/ \
  | python3 -c "
import sys, numpy as np, rasterio
with rasterio.open('/vsistdin/') as src:
    band = src.read(1)
    assert band.max() > 0, 'Band 1 is all nodata — check CRS alignment'
    print(f'Band 1 min={band.min()} max={band.max()} nodata_fraction={np.isnan(band.astype(float)).mean():.4f}')
"

# 3. Spot-check a known coordinate: pixel at EPSG:32633 easting=500000, northing=5000000
python3 - <<'EOF'
import rasterio
with rasterio.open("output_mosaic.tif") as src:
    row, col = src.index(500000, 5000000)  # (easting, northing) → (row, col)
    data = src.read(window=rasterio.windows.Window(col, row, 1, 1))
    print("Sample pixel values:", data.flatten())
EOF