322 lines
14 KiB
Python
322 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
stretch_thermal.py - Contrast-stretch a session's radiometric thermal frames
|
|
to one common signal window, for consistent, comparable visualisation.
|
|
|
|
The thermal PNGs are 16-bit radiometric (raw signal counts, monotonic with
|
|
temperature but not calibrated to degC). Each frame on its own spans only a
|
|
narrow part of the 16-bit range and uses a different sub-range, so viewing the
|
|
raw files is near-black and frame-to-frame brightness flickers.
|
|
|
|
This computes ONE window [lo, hi] pooled across every frame in the session
|
|
(default: 1st-99th percentile, robust to outlier pixels), then maps that window
|
|
to 8-bit and writes the result with a chosen palette. Because all frames share
|
|
the same window, brightness/colour is directly comparable across the session.
|
|
|
|
Usage:
|
|
python stretch_thermal.py [SESSION_DIR] [options]
|
|
|
|
--colormap {inferno,ironbow,gray} palette (default: inferno)
|
|
--out DIR output folder (default: <session>/thermal_stretched)
|
|
--lo-pct P low percentile (default: 1.0)
|
|
--hi-pct P high percentile (default: 99.0)
|
|
--lo N / --hi N hard-set the signal window, overriding percentiles
|
|
--absolute use the true session min/max instead of percentiles
|
|
--sample write one side-by-side comparison of all palettes and exit
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
import embed_metadata as embed # noqa: E402 (reuse exiftool + tag building)
|
|
import session_map as smap # noqa: E402 (session discovery)
|
|
|
|
import numpy as np # noqa: E402
|
|
from PIL import Image # noqa: E402
|
|
|
|
U16_MAX = 65536
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# palettes - built-in 256-entry LUTs, no matplotlib dependency
|
|
# --------------------------------------------------------------------------- #
|
|
def _lut_from_anchors(anchors) -> np.ndarray:
|
|
"""Build a 256x3 uint8 LUT by linear interpolation between (pos, r,g,b) anchors."""
|
|
pos = np.array([a[0] for a in anchors])
|
|
rgb = np.array([a[1:] for a in anchors], dtype=float)
|
|
x = np.linspace(0.0, 1.0, 256)
|
|
lut = np.stack([np.interp(x, pos, rgb[:, c]) for c in range(3)], axis=1)
|
|
return np.clip(lut, 0, 255).astype(np.uint8)
|
|
|
|
|
|
INFERNO = _lut_from_anchors([
|
|
(0.00, 0, 0, 4), (0.13, 31, 12, 72), (0.25, 85, 15, 109), (0.38, 136, 34, 106),
|
|
(0.50, 186, 54, 85), (0.63, 227, 89, 51), (0.75, 249, 140, 10),
|
|
(0.88, 249, 201, 50), (1.00, 252, 255, 164),
|
|
])
|
|
IRONBOW = _lut_from_anchors([
|
|
(0.00, 0, 0, 0), (0.12, 0, 0, 70), (0.25, 60, 0, 130), (0.40, 160, 0, 120),
|
|
(0.55, 220, 40, 60), (0.70, 250, 110, 0), (0.85, 255, 200, 40), (1.00, 255, 255, 255),
|
|
])
|
|
PALETTES = {"inferno": INFERNO, "ironbow": IRONBOW}
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# helpers
|
|
# --------------------------------------------------------------------------- #
|
|
def find_session(arg: str | None) -> Path:
|
|
if arg:
|
|
return Path(arg).resolve()
|
|
sessions = smap.list_sessions(Path(__file__).resolve().parent)
|
|
if not sessions:
|
|
sys.exit("No session folder found nearby; pass SESSION_DIR explicitly.")
|
|
return max(sessions, key=lambda p: p.stat().st_mtime)
|
|
|
|
|
|
def list_frames(session: Path) -> list[Path]:
|
|
"""Sorted thermal PNG frames in a session (raises if none)."""
|
|
thermal_dir = Path(session) / "thermal"
|
|
if not thermal_dir.is_dir():
|
|
raise RuntimeError(f"No thermal/ folder in {session}")
|
|
pngs = sorted(thermal_dir.glob("*.png"))
|
|
if not pngs:
|
|
raise RuntimeError(f"No PNG frames in {thermal_dir}")
|
|
return pngs
|
|
|
|
|
|
def _frame_hist(p) -> np.ndarray:
|
|
arr = np.asarray(Image.open(p)).astype(np.uint16, copy=False)
|
|
return np.bincount(arr.ravel(), minlength=U16_MAX)
|
|
|
|
|
|
def session_histogram(pngs, progress=None, workers=None) -> np.ndarray:
|
|
"""Pooled value histogram over all frames (exact for integer signal data).
|
|
|
|
Frames are read + binned in parallel (PNG decode and numpy release the GIL),
|
|
then the per-frame histograms are summed.
|
|
"""
|
|
workers = workers or embed.default_workers()
|
|
hist = np.zeros(U16_MAX, dtype=np.int64)
|
|
n = len(pngs)
|
|
done = 0
|
|
with ThreadPoolExecutor(max_workers=workers) as ex:
|
|
for h in ex.map(_frame_hist, pngs):
|
|
hist += h
|
|
done += 1
|
|
if progress:
|
|
progress(0.5 * done / n, f"Analysing frame {done}/{n}")
|
|
return hist
|
|
|
|
|
|
def percentile_from_hist(hist: np.ndarray, pct: float) -> int:
|
|
total = hist.sum()
|
|
cdf = np.cumsum(hist)
|
|
return int(np.searchsorted(cdf, total * pct / 100.0))
|
|
|
|
|
|
def compute_window(pngs, lo_pct=1.0, hi_pct=99.0, lo=None, hi=None,
|
|
absolute=False, progress=None, workers=None) -> dict:
|
|
"""Resolve the common [lo, hi] signal window for a set of frames."""
|
|
hist = session_histogram(pngs, progress=progress, workers=workers)
|
|
nz = np.nonzero(hist)[0]
|
|
abs_lo, abs_hi = int(nz[0]), int(nz[-1])
|
|
if lo is not None or hi is not None:
|
|
rlo = lo if lo is not None else abs_lo
|
|
rhi = hi if hi is not None else abs_hi
|
|
how = "manual"
|
|
elif absolute:
|
|
rlo, rhi, how = abs_lo, abs_hi, "absolute min/max"
|
|
else:
|
|
rlo = percentile_from_hist(hist, lo_pct)
|
|
rhi = percentile_from_hist(hist, hi_pct)
|
|
how = f"{lo_pct:g}-{hi_pct:g} percentile"
|
|
if rhi <= rlo:
|
|
rhi = rlo + 1
|
|
return {"lo": rlo, "hi": rhi, "abs_lo": abs_lo, "abs_hi": abs_hi, "how": how}
|
|
|
|
|
|
def render(arr: np.ndarray, lo: int, hi: int, colormap: str) -> Image.Image:
|
|
"""Stretch one 16-bit frame through [lo, hi] -> 8-bit, apply palette."""
|
|
span = max(hi - lo, 1)
|
|
norm = np.clip((arr.astype(np.float32) - lo) / span, 0.0, 1.0)
|
|
idx = (norm * 255.0 + 0.5).astype(np.uint8)
|
|
if colormap == "gray":
|
|
return Image.fromarray(idx, mode="L")
|
|
return Image.fromarray(PALETTES[colormap][idx], mode="RGB")
|
|
|
|
|
|
def render_frame(path, lo: int, hi: int, colormap: str) -> Image.Image:
|
|
"""Open one PNG and render it with the given window/palette (for previews)."""
|
|
return render(np.asarray(Image.open(path)).astype(np.uint16), lo, hi, colormap)
|
|
|
|
|
|
def make_sample_image(pngs, lo: int, hi: int, gap: int = 8) -> Image.Image:
|
|
"""Side-by-side comparison of all palettes for the middle frame."""
|
|
mid = np.asarray(Image.open(pngs[len(pngs) // 2])).astype(np.uint16)
|
|
h, w = mid.shape
|
|
tiles = [render(mid, lo, hi, c).convert("RGB") for c in ("gray", "inferno", "ironbow")]
|
|
canvas = Image.new("RGB", (w * 3 + gap * 2, h), (255, 255, 255))
|
|
for i, t in enumerate(tiles):
|
|
canvas.paste(t, (i * (w + gap), 0))
|
|
return canvas
|
|
|
|
|
|
def _embed_exif(pngs, out_dir: Path, exiftool=None, log=None, workers=None):
|
|
"""Write each frame's JSON sidecar into the matching stretched PNG (GPS, true
|
|
camera orientation, time, full JSON) so the 8-bit outputs are self-contained.
|
|
Tagging runs across several exiftool processes."""
|
|
log = log or (lambda *_: None)
|
|
workers = workers or embed.default_workers()
|
|
try:
|
|
et = exiftool or embed.ensure_exiftool(None, Path(__file__).resolve().parent / "tools")
|
|
except BaseException as e: # noqa: BLE001 (ensure_exiftool may sys.exit)
|
|
log(f" ! skipping EXIF embed - exiftool unavailable ({e})")
|
|
return
|
|
segments, dsts, srcs = [], [], []
|
|
have_sidecars = False
|
|
for p in pngs:
|
|
dst = out_dir / (p.stem + ".png")
|
|
# 1) copy any metadata already on the SOURCE image (covers stretching
|
|
# already-embedded PNGs that have no JSON sidecar next to them);
|
|
# 2) overlay/refresh from the JSON sidecar when it exists (raw session).
|
|
seg = ["-overwrite_original", "-tagsFromFile", str(p), "-EXIF:all", "-XMP:all"]
|
|
sidecar = p.with_suffix(".json")
|
|
if sidecar.exists():
|
|
try:
|
|
data = json.loads(sidecar.read_text(encoding="utf-8"))
|
|
seg += embed.build_tags(data, {}, None, "gps")
|
|
have_sidecars = True
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
seg += [str(dst), "-execute"]
|
|
segments.append(seg)
|
|
dsts.append(dst)
|
|
srcs.append(p)
|
|
if not segments:
|
|
return
|
|
idx = sorted({0, len(dsts) // 2, len(dsts) - 1}) # spot-check first/mid/last
|
|
|
|
def has_gps(path):
|
|
return bool(subprocess.run([et, "-s", "-s", "-s", "-GPSLatitude", str(path)],
|
|
capture_output=True, text=True).stdout.strip())
|
|
|
|
# only expect GPS in the output if there was any to embed (sidecar or source EXIF)
|
|
expect_gps = have_sidecars or any(has_gps(srcs[i]) for i in idx)
|
|
|
|
def run():
|
|
return embed.run_exiftool_parallel(et, segments, out_dir, workers, len(segments))
|
|
|
|
log(f"Embedding GPS / orientation EXIF into the stretched frames ({workers} workers) ...")
|
|
rc, lines, _ = run()
|
|
if expect_gps and rc == 0 and not all(has_gps(dsts[i]) for i in idx):
|
|
log(" ! GPS not present after first pass - retrying once ...")
|
|
rc, lines, _ = run()
|
|
|
|
for line in lines:
|
|
log(line)
|
|
if rc != 0 or (expect_gps and not all(has_gps(dsts[i]) for i in idx)):
|
|
raise RuntimeError(
|
|
"Could not embed GPS/EXIF into the stretched PNGs (exiftool exit "
|
|
f"{rc}). The usual cause is a FULL DISK - free up space and re-run. "
|
|
"The stretched images were written; only the embedded metadata is missing.")
|
|
if not expect_gps:
|
|
log(" note: no JSON sidecars and the source images carry no GPS - the stretched "
|
|
"PNGs have no location metadata. Stretch the raw session (with .json sidecars), "
|
|
"or run Embed first.")
|
|
else:
|
|
log(f"EXIF embedded into {len(dsts)} stretched frames.")
|
|
|
|
|
|
def stretch_session(session, out_dir=None, colormap="inferno", lo=None, hi=None,
|
|
embed_exif=True, exiftool=None, progress=None, log=None,
|
|
workers=None) -> dict:
|
|
"""
|
|
Stretch every thermal frame of a session through the [lo, hi] window and
|
|
write 8-bit images with the chosen palette. lo/hi must already be resolved
|
|
(see compute_window). If embed_exif, the JSON sidecar metadata (GPS + true
|
|
camera orientation) is written into each output PNG. Rendering and tagging
|
|
run across `workers` parallel workers. Callbacks optional.
|
|
"""
|
|
session = Path(session)
|
|
log = log or (lambda *_: None)
|
|
progress = progress or (lambda *_: None)
|
|
workers = workers or embed.default_workers()
|
|
pngs = list_frames(session)
|
|
out_dir = Path(out_dir) if out_dir else session / "thermal_stretched"
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
n = len(pngs)
|
|
log(f"Writing {n} {colormap} frames to {out_dir} ({workers} workers) ...")
|
|
lock = threading.Lock()
|
|
state = {"done": 0}
|
|
|
|
def render_one(p):
|
|
render_frame(p, lo, hi, colormap).save(out_dir / (p.stem + ".png"))
|
|
with lock:
|
|
state["done"] += 1
|
|
d = state["done"]
|
|
progress(0.5 + 0.45 * d / n, f"Rendering {d}/{n}")
|
|
|
|
with ThreadPoolExecutor(max_workers=workers) as ex:
|
|
list(ex.map(render_one, pngs))
|
|
if embed_exif:
|
|
_embed_exif(pngs, out_dir, exiftool=exiftool, log=log, workers=workers)
|
|
progress(1.0, "Done")
|
|
log(f"Done. {n} frames written.\nOutput: {out_dir}")
|
|
return {"out_dir": out_dir, "frames": n}
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# CLI
|
|
# --------------------------------------------------------------------------- #
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description="Stretch session thermal frames to a common window.")
|
|
ap.add_argument("session", nargs="?", help="session folder (default: newest session_* here)")
|
|
ap.add_argument("--colormap", choices=["inferno", "ironbow", "gray"], default="inferno")
|
|
ap.add_argument("--out", help="output folder (default: <session>/thermal_stretched)")
|
|
ap.add_argument("--lo-pct", type=float, default=1.0, help="low percentile (default 1.0)")
|
|
ap.add_argument("--hi-pct", type=float, default=99.0, help="high percentile (default 99.0)")
|
|
ap.add_argument("--lo", type=int, help="hard low signal value (overrides percentile)")
|
|
ap.add_argument("--hi", type=int, help="hard high signal value (overrides percentile)")
|
|
ap.add_argument("--absolute", action="store_true", help="use true session min/max")
|
|
ap.add_argument("--no-embed-exif", action="store_true",
|
|
help="do not write GPS/orientation EXIF into the stretched PNGs")
|
|
ap.add_argument("--sample", action="store_true",
|
|
help="write one side-by-side palette comparison and exit")
|
|
ap.add_argument("--workers", type=int, default=None,
|
|
help=f"parallel render/exiftool workers (default: {embed.default_workers()})")
|
|
args = ap.parse_args()
|
|
|
|
session = find_session(args.session)
|
|
pngs = list_frames(session)
|
|
win = compute_window(pngs, args.lo_pct, args.hi_pct, args.lo, args.hi, args.absolute,
|
|
workers=args.workers)
|
|
lo, hi = win["lo"], win["hi"]
|
|
print(f"Session : {session}")
|
|
print(f"Frames : {len(pngs)} | absolute signal range {win['abs_lo']}-{win['abs_hi']}")
|
|
print(f"Window : {lo}-{hi} ({win['how']})")
|
|
|
|
if args.sample:
|
|
out = session / "thermal_palette_sample.png"
|
|
make_sample_image(pngs, lo, hi).save(out)
|
|
print(f"\nSample written (left->right: gray | inferno | ironbow):\n {out}")
|
|
return 0
|
|
|
|
out_dir = Path(args.out).resolve() if args.out else None
|
|
res = stretch_session(session, out_dir, args.colormap, lo, hi,
|
|
embed_exif=not args.no_embed_exif, log=print, workers=args.workers)
|
|
print(f"\nWrote {res['frames']} {args.colormap} frames to:\n {res['out_dir']}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|