GeoData/scripts/ortho_build_master.py

#!/usr/bin/env python3
"""Build one global orthophoto master from tiled orthophoto JPGs."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from osgeo import gdal


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Create a global ortho master mosaic from ortho tile JPGs.")
    p.add_argument(
        "--input-dir",
        action="append",
        default=[],
        help="Input directory containing ortho JPG tiles. Can be specified multiple times.",
    )
    p.add_argument("--pattern", default="*.jpg", help="Input filename pattern.")
    p.add_argument("--out-dir", default="work/mask_master", help="Output directory.")
    p.add_argument("--vrt-name", default="ortho_master.vrt", help="Output VRT name.")
    p.add_argument("--tif-name", default="ortho_master.tif", help="Output GeoTIFF name.")
    p.add_argument("--preview-name", default="ortho_master.jpg", help="Output preview JPG name.")
    p.add_argument("--no-preview", action="store_true", help="Skip preview JPG generation.")
    p.add_argument(
        "--preview-max-size",
        type=int,
        default=8192,
        help="Longest preview edge in pixels (aspect preserved).",
    )
    p.add_argument(
        "--compress",
        default="LZW",
        help="GeoTIFF compression (e.g., LZW, DEFLATE, JPEG, NONE).",
    )
    p.add_argument("--resample", default="nearest", help="Resample algorithm for preview (nearest|bilinear|cubic...).")
    return p.parse_args()


def collect_inputs(input_dirs: list[str], pattern: str) -> list[Path]:
    dirs = [Path(d) for d in input_dirs] if input_dirs else [Path("export_unity/ortho_jpg")]
    files: list[Path] = []
    for d in dirs:
        if not d.exists():
            print(f"[ortho_build_master] Warning: input dir missing: {d}")
            continue
        for f in sorted(d.glob(pattern)):
            if f.suffix.lower() != ".jpg":
                continue
            files.append(f)
    return files


def preview_size(width: int, height: int, max_edge: int) -> tuple[int, int]:
    if width <= 0 or height <= 0:
        return width, height
    edge = max(width, height)
    if edge <= max_edge:
        return width, height
    scale = max_edge / float(edge)
    return max(1, int(round(width * scale))), max(1, int(round(height * scale)))


def main() -> int:
    args = parse_args()
    gdal.UseExceptions()

    inputs = collect_inputs(args.input_dir, args.pattern)
    if not inputs:
        raise SystemExit("[ortho_build_master] No input ortho tiles found.")

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    vrt_path = out_dir / args.vrt_name
    tif_path = out_dir / args.tif_name
    preview_path = out_dir / args.preview_name
    meta_path = out_dir / "ortho_master_meta.json"

    print(f"[ortho_build_master] Input tiles: {len(inputs)}")
    print(f"[ortho_build_master] Building VRT: {vrt_path}")
    vrt = gdal.BuildVRT(str(vrt_path), [str(p) for p in inputs])
    if vrt is None:
        raise SystemExit("[ortho_build_master] gdal.BuildVRT failed.")

    width = vrt.RasterXSize
    height = vrt.RasterYSize
    gt = vrt.GetGeoTransform(can_return_null=True)
    proj = vrt.GetProjectionRef()

    print(f"[ortho_build_master] Translating GeoTIFF: {tif_path}")
    tif_ds = gdal.Translate(
        str(tif_path),
        vrt,
        options=gdal.TranslateOptions(
            format="GTiff",
            creationOptions=[
                "TILED=YES",
                f"COMPRESS={args.compress}",
                "BIGTIFF=IF_SAFER",
            ],
        ),
    )
    if tif_ds is None:
        raise SystemExit("[ortho_build_master] gdal.Translate to GeoTIFF failed.")
    tif_ds = None

    if not args.no_preview:
        out_w, out_h = preview_size(width, height, args.preview_max_size)
        print(f"[ortho_build_master] Writing preview JPG: {preview_path} ({out_w}x{out_h})")
        jpg_ds = gdal.Translate(
            str(preview_path),
            vrt,
            options=gdal.TranslateOptions(
                format="JPEG",
                width=out_w,
                height=out_h,
                resampleAlg=args.resample,
                creationOptions=["QUALITY=92"],
            ),
        )
        if jpg_ds is None:
            raise SystemExit("[ortho_build_master] gdal.Translate to JPEG preview failed.")
        jpg_ds = None

    vrt = None

    meta = {
        "schema_version": 1,
        "inputs": [str(p) for p in inputs],
        "outputs": {
            "vrt": str(vrt_path),
            "tif": str(tif_path),
            "preview": None if args.no_preview else str(preview_path),
        },
        "raster": {
            "width": width,
            "height": height,
            "geotransform": list(gt) if gt else None,
            "projection": proj,
        },
        "settings": {
            "compress": args.compress,
            "preview_max_size": args.preview_max_size,
            "resample": args.resample,
            "pattern": args.pattern,
            "input_dirs": args.input_dir if args.input_dir else ["export_unity/ortho_jpg"],
        },
    }
    meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print(f"[ortho_build_master] Wrote metadata: {meta_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())