GeoData/export_heightmaps.py

#!/usr/bin/env python3
"""Export DGM1 tiles to Unity-ready 16-bit PNG heightmaps and a manifest."""
from __future__ import annotations

import argparse
import glob
import os
from typing import Iterable

from osgeo import gdal

from gdal_utils import (
    build_vrt,
    cleanup_aux_files,
    ensure_dir,
    ensure_parent,
    open_dataset,
    resolve_first_existing,
    safe_remove,
)

gdal.UseExceptions()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export heightmaps and manifest from DGM tiles.")
    parser.add_argument("--raw-dir", default="raw_dgm1", help="Directory containing input DGM GeoTIFFs.")
    parser.add_argument(
        "--raw-dir-new",
        dest="raw_dir_new",
        default="raw/dgm1",
        help="Preferred directory containing input DGM GeoTIFFs (new layout).",
    )
    parser.add_argument("--vrt-path", default="work/dgm.vrt", help="Path to build/read the DGM VRT.")
    parser.add_argument("--out-dir", default="export_unity/height_png16", help="Output directory for PNG heightmaps.")
    parser.add_argument(
        "--manifest-path",
        default=os.path.join("export_unity", "tile_index.csv"),
        help="Output CSV manifest path.",
    )
    parser.add_argument("--out-res", type=int, default=1025, help="Output resolution per tile (2^n + 1 for Unity).")
    parser.add_argument("--resample", default="bilinear", help="GDAL resampling algorithm used during warp.")
    parser.add_argument(
        "--tile-size-m",
        type=int,
        default=1000,
        help="Real-world tile size in meters (informational; input footprints drive bounds).",
    )
    parser.add_argument(
        "--skip-cleanup",
        action="store_true",
        help="Leave temp GDAL files instead of deleting aux XML and tmp rasters.",
    )
    return parser.parse_args()


def build_patterns(raw_dir: str) -> Iterable[str]:
    return [
        os.path.join("work", "*_tmp.tif"),
        os.path.join("work", "*_tmp.tif.aux.xml"),
        os.path.join("work", "*.aux.xml"),
        os.path.join(raw_dir, "*.aux.xml"),
    ]


def export_heightmaps(args: argparse.Namespace) -> int:
    ensure_dir("work")
    ensure_dir(args.out_dir)
    ensure_parent(args.manifest_path)

    raw_dir = (
        resolve_first_existing([args.raw_dir_new, args.raw_dir], "DGM input directory")
        if args.raw_dir_new
        else args.raw_dir
    )

    tif_paths = sorted(glob.glob(os.path.join(raw_dir, "*.tif")))
    build_vrt(args.vrt_path, tif_paths)
    ds = open_dataset(args.vrt_path, f"Could not open {args.vrt_path} after attempting to build it.")

    band = ds.GetRasterBand(1)
    gmin, gmax = band.ComputeRasterMinMax(False)
    print(f"GLOBAL_MIN={gmin}, GLOBAL_MAX={gmax}")

    with open(args.manifest_path, "w", encoding="utf-8") as f:
        f.write("tile_id,xmin,ymin,xmax,ymax,global_min,global_max,out_res\n")

        skipped = 0
        written = 0

        for tif in tif_paths:
            try:
                tds = open_dataset(tif, f"Skipping unreadable {tif}")
            except SystemExit as exc:
                print(exc)
                skipped += 1
                continue

            gt = tds.GetGeoTransform()
            ulx, xres, _, uly, _, yres = gt  # yres typically negative in north-up rasters

            # Use the source tile footprint directly to avoid shifting during export.
            xmax = ulx + xres * tds.RasterXSize
            ymin = uly + yres * tds.RasterYSize
            xmin = ulx
            ymax = uly

            base = os.path.splitext(os.path.basename(tif))[0]
            tile_id = base  # keep stable naming = easy re-export + reimport

            tmp_path = os.path.join("work", f"{tile_id}_tmp.tif")
            out_path = os.path.join(args.out_dir, f"{tile_id}.png")

            warp_opts = gdal.WarpOptions(
                outputBounds=(xmin, ymin, xmax, ymax),
                width=args.out_res,
                height=args.out_res,
                resampleAlg=args.resample,
                srcNodata=-9999,
                dstNodata=gmin,  # fill nodata with global min to avoid deep pits
            )
            try:
                gdal.Warp(tmp_path, ds, options=warp_opts)
            except RuntimeError as exc:
                print(f"Warp failed for {tile_id}: {exc}")
                skipped += 1
                continue

            # Scale to UInt16 (0..65535) using Strategy-B global min/max
            trans_opts = gdal.TranslateOptions(
                outputType=gdal.GDT_UInt16,
                scaleParams=[(gmin, gmax, 0, 65535)],
                format="PNG",
                creationOptions=["WORLDFILE=YES"],  # emit .wld so GIS tools place tiles correctly
            )
            try:
                gdal.Translate(out_path, tmp_path, options=trans_opts)
            except RuntimeError as exc:
                print(f"Translate failed for {tile_id}: {exc}")
                skipped += 1
                continue
            safe_remove(tmp_path)
            safe_remove(f"{tmp_path}.aux.xml")

            f.write(f"{tile_id},{xmin},{ymin},{xmax},{ymax},{gmin},{gmax},{args.out_res}\n")
            print(f"Wrote {out_path}")
            written += 1

    print(f"Manifest: {args.manifest_path}")
    print(f"Summary: wrote {written} tiles; skipped {skipped}.")
    if not args.skip_cleanup:
        removed = cleanup_aux_files(build_patterns(raw_dir))
        print(f"Cleanup removed {removed} temporary files/sidecars.")

    return 1 if skipped else 0


def main() -> None:
    args = parse_args()
    raise SystemExit(export_heightmaps(args))


if __name__ == "__main__":
    main()