Files
insar-management-system-v2/scripts/prepare_landsar_dem_int16.py
T

292 lines
10 KiB
Python

#!/usr/bin/env python3
"""
Prepare reusable LandSAR DEM GeoTIFFs.
The script converts large DEM rasters to uncompressed Int16 GeoTIFFs with a
stable nodata value. It streams data by windows, so it can process the 10 m
Heilongjiang DEM and the COPDEM China DEM without loading them into memory.
"""
from __future__ import annotations
import argparse
import math
import os
import sys
from pathlib import Path
from typing import Iterable, Optional
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
def _configure_proj_lib() -> None:
candidates = [
Path(sys.prefix) / "Library" / "share" / "proj",
Path(sys.prefix) / "lib" / "site-packages" / "rasterio" / "proj_data",
Path(sys.prefix) / "Lib" / "site-packages" / "rasterio" / "proj_data",
]
for data_dir in candidates:
if Path(data_dir, "proj.db").is_file():
os.environ["PROJ_LIB"] = str(data_dir)
os.environ["PROJ_DATA"] = str(data_dir)
return
try:
import pyproj
data_dir = pyproj.datadir.get_data_dir()
if data_dir and Path(data_dir, "proj.db").is_file():
os.environ["PROJ_LIB"] = str(data_dir)
os.environ["PROJ_DATA"] = str(data_dir)
except Exception:
return
_configure_proj_lib()
try:
import rasterio
from rasterio.crs import CRS
from rasterio.transform import array_bounds
from rasterio.windows import Window, from_bounds
except Exception as exc: # pragma: no cover - CLI dependency guard
raise SystemExit(
"rasterio/numpy are required. Run this with the project Python, for example:\n"
r" C:\ProgramData\anaconda3\envs\InSAR\python.exe scripts\prepare_landsar_dem_int16.py"
) from exc
DEFAULT_DEM_ROOT = Path(r"D:\DEM")
DEFAULT_OUTPUT_ROOT = DEFAULT_DEM_ROOT / "landsar_prepared"
DEFAULT_NODATA = -32768
DEFAULT_CRS = CRS.from_epsg(4326)
HEILONGJIANG_10M_DEM = "HeiLongJiang10M_DEM.tif"
HEILONGJIANG_10M_ALIAS = "\u9ed1\u9f99\u6c5f\u770110M_DEM"
SOURCE_ALIASES = {
"HeiLongJiang10M_DEM": HEILONGJIANG_10M_DEM,
"Heilongjiang10M_DEM": HEILONGJIANG_10M_DEM,
HEILONGJIANG_10M_ALIAS: HEILONGJIANG_10M_DEM,
"COPDEM_GLO30_China_4326_DEM": "COPDEM_GLO30_China_4326_DEM",
}
def _source_path(alias_or_path: str, dem_root: Path) -> Path:
text = str(alias_or_path or "").strip().strip('"')
if not text:
raise ValueError("source must not be empty")
mapped = SOURCE_ALIASES.get(text, text)
candidate = Path(mapped)
if not candidate.is_absolute():
candidate = dem_root / mapped
if candidate.exists():
return candidate
for suffix in (".tif", ".tiff", ".vrt", ".jp2"):
with_suffix = candidate.with_suffix(suffix)
if with_suffix.exists():
return with_suffix
raise FileNotFoundError(f"DEM source not found: {alias_or_path} -> {candidate}")
def _safe_stem(alias_or_path: str, source: Path) -> str:
text = str(alias_or_path or "").strip()
if text in SOURCE_ALIASES:
return text
return source.stem or source.name
def _parse_bbox(value: str | None) -> Optional[tuple[float, float, float, float]]:
if not value:
return None
parts = [part.strip() for part in value.replace(";", ",").split(",") if part.strip()]
if len(parts) != 4:
raise ValueError("--bbox must be xmin,ymin,xmax,ymax")
xmin, ymin, xmax, ymax = (float(part) for part in parts)
if xmin >= xmax or ymin >= ymax:
raise ValueError("--bbox requires xmin < xmax and ymin < ymax")
return xmin, ymin, xmax, ymax
def _align_window(window: Window, width: int, height: int) -> Window:
col_off = max(0, int(math.floor(window.col_off)))
row_off = max(0, int(math.floor(window.row_off)))
col_stop = min(width, int(math.ceil(window.col_off + window.width)))
row_stop = min(height, int(math.ceil(window.row_off + window.height)))
if col_stop <= col_off or row_stop <= row_off:
raise ValueError("requested bbox does not overlap source raster")
return Window(col_off, row_off, col_stop - col_off, row_stop - row_off)
def _iter_windows(width: int, height: int, block_size: int) -> Iterable[Window]:
step = max(64, int(block_size))
for row in range(0, height, step):
h = min(step, height - row)
for col in range(0, width, step):
w = min(step, width - col)
yield Window(col, row, w, h)
def _convert_array(data: np.ma.MaskedArray | np.ndarray, nodata: int) -> np.ndarray:
if isinstance(data, np.ma.MaskedArray):
mask = np.ma.getmaskarray(data)
array = np.asarray(data.filled(np.nan), dtype="float32")
else:
array = np.asarray(data, dtype="float32")
mask = np.zeros(array.shape, dtype=bool)
invalid = mask | ~np.isfinite(array)
rounded = np.rint(array)
rounded = np.clip(rounded, nodata + 1, 32767)
out = rounded.astype("int16", copy=False)
if invalid.any():
out = out.copy()
out[invalid] = nodata
return out
def _format_gib(byte_count: int) -> str:
return f"{byte_count / (1024 ** 3):.3f} GiB"
def convert_dem(
source_text: str,
*,
dem_root: Path,
output_root: Path,
bbox: Optional[tuple[float, float, float, float]],
suffix: str,
nodata: int,
block_size: int,
overwrite: bool,
dry_run: bool,
) -> Path:
source = _source_path(source_text, dem_root)
stem = _safe_stem(source_text, source)
if suffix:
stem = f"{stem}_{suffix.strip('_')}"
target = output_root / f"{stem}_int16.tif"
with rasterio.open(source) as src:
src_crs = src.crs or DEFAULT_CRS
if bbox:
window = _align_window(from_bounds(*bbox, transform=src.transform), src.width, src.height)
else:
window = Window(0, 0, src.width, src.height)
window = Window(int(window.col_off), int(window.row_off), int(window.width), int(window.height))
transform = src.window_transform(window)
bounds = array_bounds(int(window.height), int(window.width), transform)
estimated_bytes = int(window.width) * int(window.height) * np.dtype("int16").itemsize
print(f"Source: {source}")
print(f" driver={src.driver} dtype={src.dtypes[0]} size={src.width}x{src.height} crs={src.crs or 'EPSG:4326 assumed'}")
print(f" output window={int(window.width)}x{int(window.height)} bounds={tuple(round(v, 8) for v in bounds)}")
print(f" target={target}")
print(f" estimated raw int16 size={_format_gib(estimated_bytes)}")
if dry_run:
return target
output_root.mkdir(parents=True, exist_ok=True)
if target.exists() and not overwrite:
raise FileExistsError(f"target exists; pass --overwrite to replace it: {target}")
profile = src.profile.copy()
profile.update(
driver="GTiff",
height=int(window.height),
width=int(window.width),
count=1,
dtype="int16",
crs=src_crs,
transform=transform,
nodata=nodata,
compress="NONE",
tiled=True,
blockxsize=512,
blockysize=512,
BIGTIFF="YES",
interleave="band",
)
profile.pop("photometric", None)
profile.pop("predictor", None)
if target.exists():
target.unlink()
with rasterio.open(target, "w", **profile) as dst:
total_pixels = int(window.width) * int(window.height)
done_pixels = 0
last_percent = -1
for rel_window in _iter_windows(int(window.width), int(window.height), block_size):
src_window = Window(
window.col_off + rel_window.col_off,
window.row_off + rel_window.row_off,
rel_window.width,
rel_window.height,
)
data = src.read(1, window=src_window, masked=True)
dst.write(_convert_array(data, nodata), 1, window=rel_window)
done_pixels += int(rel_window.width) * int(rel_window.height)
percent = int(done_pixels * 100 / max(1, total_pixels))
if percent != last_percent and (percent % 5 == 0 or percent == 100):
print(f" progress={percent}%")
last_percent = percent
actual_size = target.stat().st_size if target.exists() else 0
print(f"Done: {target} ({_format_gib(actual_size)})")
return target
def main() -> int:
parser = argparse.ArgumentParser(
description="Convert large DEMs to reusable LandSAR Int16 GeoTIFFs."
)
parser.add_argument(
"--source",
action="append",
default=[],
help=(
"Source alias/path. Can be repeated. Defaults to HeiLongJiang10M_DEM "
"and COPDEM_GLO30_China_4326_DEM."
),
)
parser.add_argument("--dem-root", default=str(DEFAULT_DEM_ROOT))
parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT))
parser.add_argument(
"--bbox",
default="",
help="Optional crop bounds as xmin,ymin,xmax,ymax in EPSG:4326. Omit to convert full raster.",
)
parser.add_argument("--suffix", default="landsar")
parser.add_argument("--nodata", type=int, default=DEFAULT_NODATA)
parser.add_argument("--block-size", type=int, default=2048)
parser.add_argument("--overwrite", action="store_true")
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
sources = args.source or ["HeiLongJiang10M_DEM", "COPDEM_GLO30_China_4326_DEM"]
bbox = _parse_bbox(args.bbox)
for source in sources:
convert_dem(
source,
dem_root=Path(args.dem_root),
output_root=Path(args.output_root),
bbox=bbox,
suffix=args.suffix,
nodata=args.nodata,
block_size=args.block_size,
overwrite=args.overwrite,
dry_run=args.dry_run,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())