diff --git a/CLAUDE.md b/CLAUDE.md index 76e8e2c..048f387 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -56,10 +56,11 @@ python firemapper_gui.py CLI: ``` -python embed_metadata.py [SESSION] [--out DIR] [--gps-source position|gps] [--dry-run] +python embed_metadata.py [SESSION] [--out DIR] [--gps-source position|gps] + [--workers N] [--dry-run] python stretch_thermal.py [SESSION] [--colormap inferno|ironbow|gray] [--lo-pct P --hi-pct P | --absolute | --lo N --hi N] - [--out DIR] [--no-embed-exif] [--sample] + [--out DIR] [--no-embed-exif] [--workers N] [--sample] ``` With no `SESSION` argument they auto-pick the newest session near the script. diff --git a/embed_metadata.py b/embed_metadata.py index 3f18dee..9dd10cd 100644 --- a/embed_metadata.py +++ b/embed_metadata.py @@ -38,6 +38,8 @@ import os import shutil import subprocess import sys +import threading +from concurrent.futures import ThreadPoolExecutor from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent)) @@ -46,6 +48,66 @@ import session_map as smap # noqa: E402 (shared camera geometry + session disc IMAGE_EXTS = {".tif", ".tiff", ".png", ".jpg", ".jpeg"} GPS_EPOCH = dt.datetime(1980, 1, 6, tzinfo=dt.timezone.utc) GPS_UTC_LEAP_SECONDS = 18 # GPS-UTC offset as of 2017-2026 +WORKER_CAP = 16 # default upper bound on parallel workers + + +# --------------------------------------------------------------------------- # +# parallelism +# --------------------------------------------------------------------------- # +def default_workers(cap: int = WORKER_CAP) -> int: + """Sensible default parallelism: most cores, capped (NVMe handles it well).""" + return max(1, min(os.cpu_count() or 4, cap)) + + +def run_exiftool_parallel(et, segments, scratch_dir: Path, workers: int, total: int, + progress=None, log=None, base=0.0, span=1.0, msg="Tagging"): + """Run exiftool over per-file argument segments using several exiftool processes. + + `segments` is a list of arg-lists, one per file, each already ending with the + target file path and ``-execute``. The segments are split round-robin across + `workers` exiftool processes that run concurrently. Progress is reported as + files complete (counter is shared + locked). Returns + ``(returncode, log_lines, updated)`` - returncode is the first non-zero seen. + """ + progress = progress or (lambda *_: None) + log = log or (lambda *_: None) + workers = max(1, min(workers, len(segments))) + chunks: list[list[str]] = [[] for _ in range(workers)] + for i, seg in enumerate(segments): + chunks[i % workers].extend(seg) + + lock = threading.Lock() + state = {"done": 0, "updated": 0} + lines: list[str] = [] + rcs: list[int] = [] + + def run_chunk(k: int, arg_lines: list[str]): + argfile = scratch_dir / f"_exif_args_{k}.txt" + argfile.write_text("\n".join(arg_lines) + "\n", encoding="utf-8") + proc = subprocess.Popen( + [et, "-m", "-charset", "UTF8", "-charset", "filename=UTF8", "-@", str(argfile)], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + for line in proc.stdout: # type: ignore[union-attr] + line = line.rstrip() + if "image files" in line and ("updated" in line or "unchanged" in line): + with lock: + state["done"] += 1 + if "updated" in line: + state["updated"] += 1 + d = state["done"] + progress(base + span * d / max(total, 1), f"{msg} {d}/{total}") + elif line.strip(): + with lock: + lines.append(line) + proc.wait() + argfile.unlink(missing_ok=True) + with lock: + rcs.append(proc.returncode) + + with ThreadPoolExecutor(max_workers=workers) as ex: + list(ex.map(lambda kv: run_chunk(*kv), list(enumerate(chunks)))) + rc = next((c for c in rcs if c != 0), 0) + return rc, lines, state["updated"] # --------------------------------------------------------------------------- # @@ -211,19 +273,21 @@ def load_manifest_cams(session: Path) -> dict: def embed_session(session, out_root=None, gps_source="position", exiftool=None, - progress=None, log=None, cancel=None) -> dict: + progress=None, log=None, cancel=None, workers=None) -> dict: """ Write the JSON sidecars into EXIF/XMP of tagged copies of every image. progress(fraction_0_to_1, message) and log(message) are optional callbacks so a GUI (or the CLI) can show what is happening; cancel() may return True - to abort. Returns a small summary dict. + to abort. `workers` parallel threads copy the originals and several exiftool + processes tag them concurrently. Returns a small summary dict. """ session = Path(session) out_root = Path(out_root) if out_root else session.with_name(session.name + "_exif") log = log or (lambda *_: None) progress = progress or (lambda *_: None) cancel = cancel or (lambda: False) + workers = workers or default_workers() pairs = collect_pairs(session) if not pairs: @@ -236,52 +300,49 @@ def embed_session(session, out_root=None, gps_source="position", exiftool=None, log(f"Output : {out_root}") log(f"GPS source : {gps_source} (cam frames; thermal always uses raw gps)") log(f"exiftool : {et}") + log(f"Workers : {workers}") log(f"Found {total} image + JSON pairs.\nCopying originals into the output folder ...") - arg_lines: list[str] = [] - copied = 0 - for i, (img, sidecar) in enumerate(pairs, 1): + # pre-create the output directory tree (avoids mkdir races between threads) + for d in {(out_root / img.relative_to(session)).parent for img, _ in pairs}: + d.mkdir(parents=True, exist_ok=True) + + lock = threading.Lock() + state = {"done": 0} + + def copy_one(pair): if cancel(): raise RuntimeError("Cancelled by user.") + img, sidecar = pair rel = img.relative_to(session) cam_id = rel.parts[0] if rel.parts[0].startswith("cam") else None try: data = json.loads(sidecar.read_text(encoding="utf-8")) except Exception as e: # noqa: BLE001 log(f" ! skip {rel}: bad json ({e})") - continue + return None dst = out_root / rel - dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(img, dst) - copied += 1 - arg_lines += ["-overwrite_original", *build_tags(data, cams, cam_id, gps_source), - str(dst), "-execute"] - progress(0.5 * i / total, f"Copying {i}/{total}") + with lock: + state["done"] += 1 + d = state["done"] + progress(0.5 * d / total, f"Copying {d}/{total}") + return ["-overwrite_original", *build_tags(data, cams, cam_id, gps_source), + str(dst), "-execute"] - argfile = out_root / "_exiftool_args.txt" - argfile.write_text("\n".join(arg_lines) + "\n", encoding="utf-8") - log(f"Copied {copied} images. Writing EXIF/XMP tags with exiftool ...") + with ThreadPoolExecutor(max_workers=workers) as ex: + segments = [seg for seg in ex.map(copy_one, pairs) if seg] + copied = len(segments) + log(f"Copied {copied} images. Writing EXIF/XMP tags with {workers} exiftool workers ...") - # stream exiftool so progress advances per file - proc = subprocess.Popen( - [et, "-m", "-charset", "UTF8", "-charset", "filename=UTF8", "-@", str(argfile)], - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, - ) - done = updated = 0 - for line in proc.stdout: # type: ignore[union-attr] - line = line.rstrip() - if "image files" in line and ("updated" in line or "unchanged" in line): - done += 1 - if "updated" in line: - updated += 1 - progress(0.5 + 0.5 * done / total, f"Tagging {done}/{total}") - elif line.strip(): - log(line) - proc.wait() - argfile.unlink(missing_ok=True) - if proc.returncode != 0: + rc, lines, updated = run_exiftool_parallel( + et, segments, out_root, workers, copied, progress, log, + base=0.5, span=0.5, msg="Tagging") + for line in lines: + log(line) + if rc != 0: raise RuntimeError( - f"exiftool exited with code {proc.returncode}. The usual cause is a FULL DISK " + f"exiftool exited with code {rc}. The usual cause is a FULL DISK " "(metadata is written via a temp copy) - free up space and re-run.") progress(1.0, "Done") @@ -299,6 +360,8 @@ def main() -> int: ap.add_argument("--gps-source", choices=["position", "gps"], default="position", help="GPS source for cam frames (default: position = fused INS)") ap.add_argument("--exiftool", help="path to exiftool (default: auto-detect/install)") + ap.add_argument("--workers", type=int, default=None, + help=f"parallel copy/exiftool workers (default: {default_workers()})") ap.add_argument("--dry-run", action="store_true", help="report only, write nothing") args = ap.parse_args() @@ -321,7 +384,7 @@ def main() -> int: return 0 embed_session(session, args.out, args.gps_source, args.exiftool, - progress=lambda f, m: None, log=print) + progress=lambda f, m: None, log=print, workers=args.workers) return 0 diff --git a/firemapper_gui.py b/firemapper_gui.py index b8d6594..ad70cc3 100644 --- a/firemapper_gui.py +++ b/firemapper_gui.py @@ -88,6 +88,7 @@ class FireMapperGUI(tk.Tk): self.bars: dict[str, ttk.Progressbar] = {} self.status_vars: dict[str, tk.StringVar] = {} self._preview_imgtk = None # keep a reference so Tk doesn't GC it + self.workers = tk.IntVar(self, value=embed.default_workers()) self._build_header() nb = ttk.Notebook(self) @@ -135,6 +136,16 @@ class FireMapperGUI(tk.Tk): ttk.Button(row, text="Browse...", command=browse).pack(side="left", padx=(6, 0)) return row + def _workers_row(self, parent): + row = ttk.Frame(parent) + ttk.Label(row, text="Parallel workers", width=16).pack(side="left") + ttk.Spinbox(row, from_=1, to=(os.cpu_count() or 32), increment=1, width=6, + textvariable=self.workers).pack(side="left") + ttk.Label(row, foreground="#777", + text=f" threads / exiftool processes (you have {os.cpu_count()} cores)" + ).pack(side="left", padx=6) + return row + def _progress_block(self, parent, which): frame = ttk.Frame(parent) bar = ttk.Progressbar(frame, mode="determinate", maximum=1000) @@ -190,6 +201,8 @@ class FireMapperGUI(tk.Tk): ttk.Entry(adv, textvariable=self.embed_exif).pack(side="left", fill="x", expand=True) ttk.Label(adv, text="(optional - auto)", foreground="#777").pack(side="left", padx=6) + self._workers_row(tab).pack(fill="x", pady=(8, 2)) + btn = ttk.Button(tab, text="Embed metadata", command=self._start_embed) btn.pack(anchor="w", pady=10) self.action_buttons.append(btn) @@ -242,6 +255,8 @@ class FireMapperGUI(tk.Tk): text="Embed GPS / orientation metadata into the stretched PNGs " "(needs exiftool)").pack(anchor="w", pady=(6, 0)) + self._workers_row(tab).pack(fill="x", pady=(6, 2)) + bar = ttk.Frame(tab) bar.pack(fill="x", pady=10) pv = ttk.Button(bar, text="Preview palettes", command=self._start_preview) @@ -490,11 +505,13 @@ class FireMapperGUI(tk.Tk): out = self.embed_out.get().strip() or None gps = self.embed_gps.get() et = self.embed_exif.get().strip() or None + workers = max(1, int(self.workers.get())) progress, log = self._callbacks("embed") def work(): try: - res = embed.embed_session(session, out, gps, et, progress=progress, log=log) + res = embed.embed_session(session, out, gps, et, progress=progress, log=log, + workers=workers) self.q.put(("done", "embed", res)) except Exception as e: # noqa: BLE001 self.q.put(("error", "embed", str(e))) @@ -509,12 +526,13 @@ class FireMapperGUI(tk.Tk): if not self._guard(session, "thermal"): return params = self._read_window_params() + workers = max(1, int(self.workers.get())) progress, log = self._callbacks("thermal") def work(): try: pngs = stretch.list_frames(Path(session)) - win = stretch.compute_window(pngs, progress=progress, **params) + win = stretch.compute_window(pngs, progress=progress, workers=workers, **params) log(f"Frames: {len(pngs)} | absolute signal {win['abs_lo']}-{win['abs_hi']}") log(f"Window: {win['lo']}-{win['hi']} ({win['how']})") img = stretch.make_sample_image(pngs, win["lo"], win["hi"]) @@ -531,16 +549,18 @@ class FireMapperGUI(tk.Tk): out = self.th_out.get().strip() or None cmap = self.th_cmap.get() params = self._read_window_params() + workers = max(1, int(self.workers.get())) + embed_exif = bool(self.th_embed.get()) progress, log = self._callbacks("thermal") def work(): try: pngs = stretch.list_frames(Path(session)) - win = stretch.compute_window(pngs, progress=progress, **params) + win = stretch.compute_window(pngs, progress=progress, workers=workers, **params) log(f"Window: {win['lo']}-{win['hi']} ({win['how']})") res = stretch.stretch_session(session, out, cmap, win["lo"], win["hi"], - embed_exif=self.th_embed.get(), - progress=progress, log=log) + embed_exif=embed_exif, + progress=progress, log=log, workers=workers) self.q.put(("done", "thermal", {"kind": "stretch", **res})) except Exception as e: # noqa: BLE001 self.q.put(("error", "thermal", str(e))) diff --git a/stretch_thermal.py b/stretch_thermal.py index 9df73a3..385d2c8 100644 --- a/stretch_thermal.py +++ b/stretch_thermal.py @@ -31,6 +31,8 @@ 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)) @@ -90,15 +92,27 @@ def list_frames(session: Path) -> list[Path]: return pngs -def session_histogram(pngs, progress=None) -> np.ndarray: - """Pooled value histogram over all frames (exact for integer signal data).""" +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) - for i, p in enumerate(pngs, 1): - arr = np.asarray(Image.open(p)).astype(np.uint16, copy=False) - hist += np.bincount(arr.ravel(), minlength=U16_MAX) - if progress: - progress(0.5 * i / n, f"Analysing frame {i}/{n}") + 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 @@ -109,9 +123,9 @@ def percentile_from_hist(hist: np.ndarray, pct: float) -> int: def compute_window(pngs, lo_pct=1.0, hi_pct=99.0, lo=None, hi=None, - absolute=False, progress=None) -> dict: + 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) + 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: @@ -155,16 +169,18 @@ def make_sample_image(pngs, lo: int, hi: int, gap: int = 8) -> Image.Image: return canvas -def _embed_exif(pngs, out_dir: Path, exiftool=None, log=None): +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.""" + 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 - arg_lines, dsts, srcs = [], [], [] + segments, dsts, srcs = [], [], [] have_sidecars = False for p in pngs: dst = out_dir / (p.stem + ".png") @@ -181,15 +197,12 @@ def _embed_exif(pngs, out_dir: Path, exiftool=None, log=None): except Exception: # noqa: BLE001 pass seg += [str(dst), "-execute"] - arg_lines += seg + segments.append(seg) dsts.append(dst) srcs.append(p) - if not arg_lines: + if not segments: return - argfile = out_dir / "_exif_args.txt" - argfile.write_text("\n".join(arg_lines) + "\n", encoding="utf-8") - cmd = [et, "-m", "-charset", "UTF8", "-charset", "filename=UTF8", "-@", str(argfile)] - idx = sorted({0, len(dsts) // 2, len(dsts) - 1}) # spot-check first/mid/last + 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)], @@ -198,22 +211,22 @@ def _embed_exif(pngs, out_dir: Path, exiftool=None, log=None): # 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) - log("Embedding GPS / orientation EXIF into the stretched frames ...") - proc = subprocess.run(cmd, capture_output=True, text=True) - if expect_gps and proc.returncode == 0 and not all(has_gps(dsts[i]) for i in idx): - log(" ! GPS not present after first pass - retrying once ...") - proc = subprocess.run(cmd, capture_output=True, text=True) - argfile.unlink(missing_ok=True) + def run(): + return embed.run_exiftool_parallel(et, segments, out_dir, workers, len(segments)) - err = (proc.stderr or "").strip() - if err: - log(err) - if proc.returncode != 0 or (expect_gps and not all(has_gps(dsts[i]) for i in idx)): + 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"{proc.returncode}). The usual cause is a FULL DISK - free up space and re-run. " - "The stretched images were written; only the embedded metadata is missing." - + (f"\nexiftool: {err}" if err else "")) + 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), " @@ -223,27 +236,39 @@ def _embed_exif(pngs, out_dir: Path, exiftool=None, log=None): def stretch_session(session, out_dir=None, colormap="inferno", lo=None, hi=None, - embed_exif=True, exiftool=None, progress=None, log=None) -> dict: + 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. Callbacks optional. + 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} ...") - for i, p in enumerate(pngs, 1): + 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")) - progress(0.5 + 0.45 * i / n, f"Rendering {i}/{n}") + 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) + _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} @@ -266,11 +291,14 @@ def main() -> int: 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) + 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']}") @@ -284,7 +312,7 @@ def main() -> int: 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) + 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