Coverage for lunar_toolkit/asp.py: 94%
190 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 02:58 +0000
1"""Ames Stereo Pipeline wrappers: positional inputs as arguments, other options as keywords."""
3from __future__ import annotations
5import json
6import re
7import time
8from collections.abc import Sequence
9from dataclasses import dataclass
10from pathlib import Path
11from typing import Any
13import numpy as np
15from .tools import ToolError, asp_options, gdal_options, is_dry_run, run
18def _paths(items: Sequence[Any]) -> list[str]:
19 return [str(p) for p in items]
22def _parent(prefix: Path) -> None:
23 if not is_dry_run():
24 Path(prefix).parent.mkdir(parents=True, exist_ok=True)
27def mapproject(
28 dem: Path, image: Path, out: Path, camera: Path | None = None, *, cwd: Path | None = None, **options: Any
29) -> Path:
30 """Orthorectify *image* onto *dem* (``tr=``, ``t_projwin=``, ``bundle_adjust_prefix=``)."""
31 _parent(out)
32 cams = [camera] if camera else []
33 run("mapproject", [*asp_options(options), dem, image, *cams, out], cwd=cwd)
34 return Path(out)
37def bundle_adjust(
38 images: Sequence[Path],
39 out_prefix: Path,
40 cameras: Sequence[Path] = (),
41 *,
42 mapprojected: Sequence[Path] | None = None,
43 dem: Path | None = None,
44 parallel: bool = False,
45 cwd: Path | None = None,
46 **options: Any,
47) -> Path:
48 """Run ``bundle_adjust`` (or ``parallel_bundle_adjust`` with *parallel*).
50 *mapprojected* (same order as *images*) and the *dem* they were projected
51 onto become ``--mapprojected-data``, which finds matches between images of
52 very different illumination. Returns *out_prefix*.
54 Raises [`ToolError`][lunar_toolkit.tools.ToolError] when the run writes no fresh
55 ``.adjust`` file for some image: ASP exits 0 even when it finds nothing to
56 optimize ("No points to optimize ... Cannot continue").
57 """
58 _parent(out_prefix)
59 if mapprojected is not None:
60 options["mapprojected_data"] = " ".join(_paths([*mapprojected, *([dem] if dem else [])]))
61 tool = "parallel_bundle_adjust" if parallel else "bundle_adjust"
62 start = time.time()
63 run(tool, [*_paths(images), *_paths(cameras), *asp_options(options), "-o", out_prefix], cwd=cwd)
64 if not is_dry_run():
65 prefix = Path(cwd or ".") / out_prefix
66 cams: Sequence[Path | None] = cameras if len(cameras) == len(images) else [None] * len(images)
67 missing = [
68 p
69 for p in (bundle_adjust_file_name(prefix, im, cam) for im, cam in zip(images, cams))
70 if not p.is_file() or p.stat().st_mtime < start - 2
71 ]
72 if missing:
73 raise ToolError(
74 f"{tool} wrote no adjustment for {len(missing)} of {len(images)} images "
75 f"(e.g. {missing[0].name}); see {out_prefix}-log-bundle_adjust-*.txt"
76 )
77 return Path(out_prefix)
80def apply_transform(
81 images: Sequence[Path],
82 transform: Path,
83 out_prefix: Path,
84 cameras: Sequence[Path] = (),
85 *,
86 input_adjustments_prefix: Path | None = None,
87 cwd: Path | None = None,
88 **options: Any,
89) -> Path:
90 """Apply a ``pc_align`` 4×4 transform to (adjusted) cameras, without optimizing."""
91 return bundle_adjust(
92 images,
93 out_prefix,
94 cameras,
95 initial_transform=transform,
96 input_adjustments_prefix=input_adjustments_prefix,
97 apply_initial_transform_only=True,
98 cwd=cwd,
99 **options,
100 )
103def sfs(
104 dem: Path,
105 images: Sequence[Path],
106 out_prefix: Path,
107 cameras: Sequence[Path] = (),
108 *,
109 parallel: bool = False,
110 cwd: Path | None = None,
111 **options: Any,
112) -> Path:
113 """Run ``sfs`` (or ``parallel_sfs``); returns ``<out_prefix>-DEM-final.tif``."""
114 _parent(out_prefix)
115 tool = "parallel_sfs" if parallel else "sfs"
116 run(
117 tool, ["-i", dem, *_paths(images), *_paths(cameras), *asp_options(options), "-o", out_prefix], cwd=cwd
118 )
119 return Path(f"{out_prefix}-DEM-final.tif")
122@dataclass(frozen=True)
123class SunAngles:
124 image: str
125 azimuth: float
126 elevation: float
129_SUN_RE = re.compile(r"Sun azimuth and elevation for: (?P<image>\S+) are (?P<az>\S+) and (?P<el>\S+) degrees")
132def parse_sun_angles(text: str) -> list[SunAngles]:
133 """Parse ``sfs --query`` output into per-image Sun azimuth / elevation."""
134 return [SunAngles(m["image"], float(m["az"]), float(m["el"])) for m in _SUN_RE.finditer(text)]
137def sun_angles(
138 dem: Path,
139 images: Sequence[Path],
140 cameras: Sequence[Path] = (),
141 *,
142 out_prefix: Path = Path("sfs_query/run"),
143 cwd: Path | None = None,
144 **options: Any,
145) -> list[SunAngles]:
146 """Sun azimuth (from north) and elevation per image at the DEM centre.
148 Runs ``sfs --query``; results come back in input order. Sort by azimuth
149 before bundle adjustment and SfS.
150 """
151 _parent(out_prefix)
152 res = run(
153 "sfs",
154 ["-i", dem, *_paths(images), *_paths(cameras), "--query", *asp_options(options), "-o", out_prefix],
155 cwd=cwd,
156 capture=True,
157 )
158 return parse_sun_angles(res.stdout or "")
161def pc_align(
162 reference: Path,
163 source: Path,
164 out_prefix: Path,
165 *,
166 max_displacement: float,
167 cwd: Path | None = None,
168 **options: Any,
169) -> Path:
170 """Align *source* to *reference*; returns ``<out_prefix>-transform.txt``."""
171 _parent(out_prefix)
172 run(
173 "pc_align",
174 ["--max-displacement", max_displacement, *asp_options(options), reference, source, "-o", out_prefix],
175 cwd=cwd,
176 )
177 return Path(f"{out_prefix}-transform.txt")
180def point2dem(cloud: Path, *, cwd: Path | None = None, **options: Any) -> Path:
181 """Grid a point cloud; returns ``<prefix>-DEM.tif`` (prefix from ``o=`` or the cloud)."""
182 run("point2dem", [*asp_options(options), cloud], cwd=cwd)
183 prefix = options.get("o") or options.get("output_prefix")
184 if prefix is None:
185 stem = str(cloud)
186 prefix = stem[: -len("-PC.tif")] if stem.endswith("-PC.tif") else str(Path(stem).with_suffix(""))
187 return Path(f"{prefix}-DEM.tif")
190def stereo(
191 left: Path,
192 right: Path,
193 out_prefix: Path,
194 left_camera: Path | None = None,
195 right_camera: Path | None = None,
196 *,
197 dem: Path | None = None,
198 parallel: bool = True,
199 cwd: Path | None = None,
200 **options: Any,
201) -> Path:
202 """Run ``parallel_stereo`` (or ``stereo``); returns ``<out_prefix>-PC.tif``.
204 Pass *dem* when the inputs are mapprojected onto it.
205 """
206 _parent(out_prefix)
207 cams = [c for c in (left_camera, right_camera) if c]
208 tail = [dem] if dem else []
209 tool = "parallel_stereo" if parallel else "stereo"
210 run(tool, [*asp_options(options), left, right, *cams, out_prefix, *tail], cwd=cwd)
211 return Path(f"{out_prefix}-PC.tif")
214def dem_mosaic(inputs: Sequence[Path], out: Path, *, cwd: Path | None = None, **options: Any) -> Path:
215 """Mosaic DEMs or images (``max=True`` for a maximally-lit mosaic)."""
216 _parent(out)
217 run("dem_mosaic", [*_paths(inputs), *asp_options(options), "-o", out], cwd=cwd)
218 return Path(out)
221def hillshade(
222 dem: Path,
223 out: Path | None = None,
224 *,
225 azimuth: float = 300,
226 elevation: float = 20,
227 cwd: Path | None = None,
228 **options: Any,
229) -> Path:
230 out = Path(out) if out else Path(dem).with_name(Path(dem).stem + "_hillshade.tif")
231 run(
232 "hillshade",
233 [dem, "--azimuth", azimuth, "--elevation", elevation, *asp_options(options), "-o", out],
234 cwd=cwd,
235 )
236 return out
239def geodiff(dem1: Path, dem2: Path, out_prefix: Path, *, cwd: Path | None = None, **options: Any) -> Path:
240 """``dem1 - dem2`` (``absolute=True`` for |diff|); returns ``<out_prefix>-diff.tif``."""
241 _parent(out_prefix)
242 run("geodiff", [*asp_options(options), dem1, dem2, "-o", out_prefix], cwd=cwd)
243 return Path(f"{out_prefix}-diff.tif")
246def gdal_translate(
247 src: Path,
248 dst: Path,
249 *,
250 projwin: Sequence[float] | None = None,
251 tr: float | None = None,
252 r: str | None = None,
253 cwd: Path | None = None,
254 **options: Any,
255) -> Path:
256 """Crop (``projwin`` = ulx uly lrx lry) and / or resample (``tr`` metres, ``r`` method)."""
257 _parent(dst)
258 opts = {
259 "projwin": projwin,
260 "tr": (tr, tr) if tr is not None else None,
261 "r": r,
262 "of": "GTiff",
263 "co": "COMPRESS=LZW",
264 **options,
265 }
266 run("gdal_translate", [*gdal_options(opts), src, dst], cwd=cwd)
267 return Path(dst)
270def bundle_adjust_file_name(prefix: Path, image: Path, camera: Path | None = None) -> Path:
271 """``<prefix>-<stem>.adjust`` as ASP names it (stem of *camera*, else *image*, up to the last dot)."""
272 name = Path(camera or image).name
273 stem = name[: name.rfind(".")] if "." in name else name
274 return Path(f"{prefix}-{stem}.adjust")
277@dataclass(frozen=True)
278class MatchOffsetStats:
279 """One row of ``<prefix>-mapproj_match_offset_stats.txt`` (metres; -1 when no matches)."""
281 image: str
282 p25: float
283 p50: float
284 p75: float
285 p85: float
286 p95: float
287 count: int
290def read_mapproj_offset_stats(path: Path) -> dict[str, MatchOffsetStats]:
291 """Parse ``bundle_adjust --mapproj-dem`` per-image offsets, keyed by image file name.
293 Format (ASP ``saveMapprojOffsets``): ``# image_name 25% 50% 75% 85% 95% count``.
294 """
295 out = {}
296 for line in Path(path).read_text().splitlines():
297 parts = line.split()
298 if not parts or parts[0].startswith("#") or len(parts) < 7:
299 continue
300 p25, p50, p75, p85, p95 = map(float, parts[1:6])
301 out[Path(parts[0]).name] = MatchOffsetStats(parts[0], p25, p50, p75, p85, p95, int(float(parts[6])))
302 return out
305@dataclass(frozen=True)
306class ResidualStats:
307 """One row of ``<prefix>-final_residuals_stats.txt`` (pixels)."""
309 image: str
310 mean: float
311 median: float
312 count: int
315def read_residual_stats(path: Path) -> dict[str, ResidualStats]:
316 """Parse per-camera reprojection errors (``# Image, mean, median, count``), keyed by file name."""
317 out = {}
318 for line in Path(path).read_text().splitlines():
319 parts = [p.strip() for p in line.split(",")]
320 if not parts[0] or parts[0].startswith("#") or len(parts) < 4:
321 continue
322 out[Path(parts[0]).name] = ResidualStats(parts[0], float(parts[1]), float(parts[2]), int(parts[3]))
323 return out
326def read_residual_norms(path: Path) -> dict[str, np.ndarray]:
327 """Per-image reprojection error norms (pixels) from ``<prefix>-final_residuals_raw_pixels.txt``.
329 The file has, per image, a ``name, count`` line followed by *count* ``ex, ey`` lines.
330 """
331 out: dict[str, np.ndarray] = {}
332 lines = Path(path).read_text().splitlines()
333 i = 0
334 while i < len(lines):
335 head = [p.strip() for p in lines[i].split(",")]
336 i += 1
337 if len(head) != 2 or not head[0] or head[0].startswith("#"): 337 ↛ 338line 337 didn't jump to line 338 because the condition on line 337 was never true
338 continue
339 n = int(head[1])
340 xy = np.array([[float(v) for v in ln.split(",")] for ln in lines[i : i + n]]).reshape(-1, 2)
341 out[Path(head[0]).name] = np.hypot(xy[:, 0], xy[:, 1])
342 i += n
343 return out
346def image_subset(
347 images: Sequence[Path], out: Path, *, threshold: float, cwd: Path | None = None, **options: Any
348) -> list[tuple[str, int]]:
349 """Run ``image_subset``; returns ``(image, contributing pixels)`` best first.
351 Writes the image list to ``<out>.in.txt``; ``t_projwin=(xmin, ymin, xmax, ymax)`` limits
352 the evaluated extent. Empty in a dry run.
353 """
354 _parent(out)
355 image_list = Path(f"{out}.in.txt")
356 if not is_dry_run(): 356 ↛ 358line 356 didn't jump to line 358 because the condition on line 356 was always true
357 image_list.write_text("".join(f"{p}\n" for p in images))
358 run(
359 "image_subset",
360 ["--image-list", image_list, "--threshold", threshold, *asp_options(options), "-o", out],
361 cwd=cwd,
362 )
363 if is_dry_run() or not Path(out).exists(): 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 return []
365 return [
366 (parts[0], int(float(parts[1])))
367 for parts in (ln.split() for ln in Path(out).read_text().splitlines())
368 if len(parts) >= 2 and not parts[0].startswith("#")
369 ]
372def raster_info(path: Path) -> dict:
373 """``gdalinfo -json`` of *path* (empty in a dry run)."""
374 res = run("gdalinfo", ["-json", path], capture=True)
375 return json.loads(res.stdout) if res.stdout else {}
378def raster_bounds(path: Path) -> tuple[float, float, float, float] | None:
379 """``(ulx, uly, lrx, lry)`` of a north-up raster, or None in a dry run."""
380 info = raster_info(path)
381 if not info: 381 ↛ 383line 381 didn't jump to line 383 because the condition on line 381 was always true
382 return None
383 x0, dx, _, y0, _, dy = info["geoTransform"]
384 nx, ny = info["size"]
385 return x0, y0, x0 + nx * dx, y0 + ny * dy
388def srs_proj4(path: Path) -> str:
389 """PROJ string of a raster's spatial reference (``gdalsrsinfo -o proj4``; empty in a dry run)."""
390 res = run("gdalsrsinfo", ["-o", "proj4", path], capture=True)
391 return (res.stdout or "").strip().strip("'")
394@dataclass(frozen=True)
395class Raster:
396 """Band 1 of a north-up raster: ``data[row, col]`` (NaN = nodata), top-left corner and pixel size."""
398 data: np.ndarray
399 x0: float
400 y0: float
401 dx: float
402 dy: float
405def read_raster(
406 path: Path,
407 *,
408 projwin: Sequence[float] | None = None,
409 outsize: tuple[int, int] | None = None,
410 r: str = "average",
411) -> Raster | None:
412 """Read (a resampled window of) a raster through ``gdal_translate -of AAIGrid /vsistdout/``.
414 *projwin* (``ulx uly lrx lry``) may extend past the raster (filled with nodata) and
415 *outsize* ``(columns, rows)`` resamples with *r*, so several images read with the
416 same window share one grid. Meant for small grids; returns None in a dry run.
417 """
418 opts = {
419 "q": True,
420 "of": "AAIGrid",
421 "ot": "Float32",
422 "b": 1,
423 "projwin": projwin,
424 "outsize": outsize,
425 "r": r if outsize else None,
426 }
427 res = run("gdal_translate", [*gdal_options(opts), path, "/vsistdout/"], capture=True)
428 return parse_aaigrid(res.stdout) if res.stdout else None
431def parse_aaigrid(text: str) -> Raster:
432 """Parse ESRI ASCII grid text (as GDAL's AAIGrid driver writes it)."""
433 lines = text.splitlines()
434 head: dict[str, float] = {}
435 k = 0
436 while k < len(lines) and lines[k][:1].isalpha():
437 key, value = lines[k].split()[:2]
438 head[key.lower()] = float(value)
439 k += 1
440 nx, ny = int(head["ncols"]), int(head["nrows"])
441 dx = head.get("dx", head.get("cellsize"))
442 dy = head.get("dy", head.get("cellsize"))
443 if dx is None or dy is None: 443 ↛ 444line 443 didn't jump to line 444 because the condition on line 443 was never true
444 raise ValueError("AAIGrid header lacks cellsize / dx / dy")
445 x0 = head["xllcorner"] if "xllcorner" in head else head["xllcenter"] - dx / 2
446 yll = head["yllcorner"] if "yllcorner" in head else head["yllcenter"] - dy / 2
447 # GDAL also writes the .prj (WKT) to /vsistdout/ after the cells.
448 data = np.array(" ".join(lines[k:]).split()[: nx * ny], dtype=float).reshape(ny, nx)
449 if "nodata_value" in head: 449 ↛ 451line 449 didn't jump to line 451 because the condition on line 449 was always true
450 data[data == head["nodata_value"]] = np.nan
451 return Raster(data, x0, yll + ny * dy, dx, -dy)