Parallelize embed and thermal stretch (configurable workers)
Both conversion steps ran their hot loops serially. Now they fan out across a configurable worker pool (default min(cpu, 16); the rig has 20 cores + NVMe): - embed_metadata: parallel image copy (thread pool) + tagging split across several concurrent exiftool processes via new run_exiftool_parallel(). - stretch_thermal: parallel session histogram, parallel render/save, and parallel exiftool embedding (reuses run_exiftool_parallel); GPS verify + retry preserved. - CLI: --workers N on both tools. GUI: shared "Parallel workers" spinbox. Measured on a 135-frame session (270 embed pairs), workers 1 -> 16: render+embed 38.9s -> 6.0s (6.5x), embed copy+tag 73.5s -> 13.0s (5.7x), histogram 1.7s -> 0.7s. 0 frames missing GPS in every run. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
parent
6e6b675140
commit
4f1ae7f7c4
|
|
@ -56,10 +56,11 @@ python firemapper_gui.py
|
||||||
|
|
||||||
CLI:
|
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]
|
python stretch_thermal.py [SESSION] [--colormap inferno|ironbow|gray]
|
||||||
[--lo-pct P --hi-pct P | --absolute | --lo N --hi N]
|
[--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.
|
With no `SESSION` argument they auto-pick the newest session near the script.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -38,6 +38,8 @@ import os
|
||||||
import shutil
|
import shutil
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
|
import threading
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
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"}
|
IMAGE_EXTS = {".tif", ".tiff", ".png", ".jpg", ".jpeg"}
|
||||||
GPS_EPOCH = dt.datetime(1980, 1, 6, tzinfo=dt.timezone.utc)
|
GPS_EPOCH = dt.datetime(1980, 1, 6, tzinfo=dt.timezone.utc)
|
||||||
GPS_UTC_LEAP_SECONDS = 18 # GPS-UTC offset as of 2017-2026
|
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,
|
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.
|
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
|
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
|
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)
|
session = Path(session)
|
||||||
out_root = Path(out_root) if out_root else session.with_name(session.name + "_exif")
|
out_root = Path(out_root) if out_root else session.with_name(session.name + "_exif")
|
||||||
log = log or (lambda *_: None)
|
log = log or (lambda *_: None)
|
||||||
progress = progress or (lambda *_: None)
|
progress = progress or (lambda *_: None)
|
||||||
cancel = cancel or (lambda: False)
|
cancel = cancel or (lambda: False)
|
||||||
|
workers = workers or default_workers()
|
||||||
|
|
||||||
pairs = collect_pairs(session)
|
pairs = collect_pairs(session)
|
||||||
if not pairs:
|
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"Output : {out_root}")
|
||||||
log(f"GPS source : {gps_source} (cam frames; thermal always uses raw gps)")
|
log(f"GPS source : {gps_source} (cam frames; thermal always uses raw gps)")
|
||||||
log(f"exiftool : {et}")
|
log(f"exiftool : {et}")
|
||||||
|
log(f"Workers : {workers}")
|
||||||
log(f"Found {total} image + JSON pairs.\nCopying originals into the output folder ...")
|
log(f"Found {total} image + JSON pairs.\nCopying originals into the output folder ...")
|
||||||
|
|
||||||
arg_lines: list[str] = []
|
# pre-create the output directory tree (avoids mkdir races between threads)
|
||||||
copied = 0
|
for d in {(out_root / img.relative_to(session)).parent for img, _ in pairs}:
|
||||||
for i, (img, sidecar) in enumerate(pairs, 1):
|
d.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
lock = threading.Lock()
|
||||||
|
state = {"done": 0}
|
||||||
|
|
||||||
|
def copy_one(pair):
|
||||||
if cancel():
|
if cancel():
|
||||||
raise RuntimeError("Cancelled by user.")
|
raise RuntimeError("Cancelled by user.")
|
||||||
|
img, sidecar = pair
|
||||||
rel = img.relative_to(session)
|
rel = img.relative_to(session)
|
||||||
cam_id = rel.parts[0] if rel.parts[0].startswith("cam") else None
|
cam_id = rel.parts[0] if rel.parts[0].startswith("cam") else None
|
||||||
try:
|
try:
|
||||||
data = json.loads(sidecar.read_text(encoding="utf-8"))
|
data = json.loads(sidecar.read_text(encoding="utf-8"))
|
||||||
except Exception as e: # noqa: BLE001
|
except Exception as e: # noqa: BLE001
|
||||||
log(f" ! skip {rel}: bad json ({e})")
|
log(f" ! skip {rel}: bad json ({e})")
|
||||||
continue
|
return None
|
||||||
dst = out_root / rel
|
dst = out_root / rel
|
||||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
shutil.copy2(img, dst)
|
shutil.copy2(img, dst)
|
||||||
copied += 1
|
with lock:
|
||||||
arg_lines += ["-overwrite_original", *build_tags(data, cams, cam_id, gps_source),
|
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"]
|
str(dst), "-execute"]
|
||||||
progress(0.5 * i / total, f"Copying {i}/{total}")
|
|
||||||
|
|
||||||
argfile = out_root / "_exiftool_args.txt"
|
with ThreadPoolExecutor(max_workers=workers) as ex:
|
||||||
argfile.write_text("\n".join(arg_lines) + "\n", encoding="utf-8")
|
segments = [seg for seg in ex.map(copy_one, pairs) if seg]
|
||||||
log(f"Copied {copied} images. Writing EXIF/XMP tags with exiftool ...")
|
copied = len(segments)
|
||||||
|
log(f"Copied {copied} images. Writing EXIF/XMP tags with {workers} exiftool workers ...")
|
||||||
|
|
||||||
# stream exiftool so progress advances per file
|
rc, lines, updated = run_exiftool_parallel(
|
||||||
proc = subprocess.Popen(
|
et, segments, out_root, workers, copied, progress, log,
|
||||||
[et, "-m", "-charset", "UTF8", "-charset", "filename=UTF8", "-@", str(argfile)],
|
base=0.5, span=0.5, msg="Tagging")
|
||||||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
|
for line in lines:
|
||||||
)
|
|
||||||
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)
|
log(line)
|
||||||
proc.wait()
|
if rc != 0:
|
||||||
argfile.unlink(missing_ok=True)
|
|
||||||
if proc.returncode != 0:
|
|
||||||
raise RuntimeError(
|
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.")
|
"(metadata is written via a temp copy) - free up space and re-run.")
|
||||||
|
|
||||||
progress(1.0, "Done")
|
progress(1.0, "Done")
|
||||||
|
|
@ -299,6 +360,8 @@ def main() -> int:
|
||||||
ap.add_argument("--gps-source", choices=["position", "gps"], default="position",
|
ap.add_argument("--gps-source", choices=["position", "gps"], default="position",
|
||||||
help="GPS source for cam frames (default: position = fused INS)")
|
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("--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")
|
ap.add_argument("--dry-run", action="store_true", help="report only, write nothing")
|
||||||
args = ap.parse_args()
|
args = ap.parse_args()
|
||||||
|
|
||||||
|
|
@ -321,7 +384,7 @@ def main() -> int:
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
embed_session(session, args.out, args.gps_source, args.exiftool,
|
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
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -88,6 +88,7 @@ class FireMapperGUI(tk.Tk):
|
||||||
self.bars: dict[str, ttk.Progressbar] = {}
|
self.bars: dict[str, ttk.Progressbar] = {}
|
||||||
self.status_vars: dict[str, tk.StringVar] = {}
|
self.status_vars: dict[str, tk.StringVar] = {}
|
||||||
self._preview_imgtk = None # keep a reference so Tk doesn't GC it
|
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()
|
self._build_header()
|
||||||
nb = ttk.Notebook(self)
|
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))
|
ttk.Button(row, text="Browse...", command=browse).pack(side="left", padx=(6, 0))
|
||||||
return row
|
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):
|
def _progress_block(self, parent, which):
|
||||||
frame = ttk.Frame(parent)
|
frame = ttk.Frame(parent)
|
||||||
bar = ttk.Progressbar(frame, mode="determinate", maximum=1000)
|
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.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)
|
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 = ttk.Button(tab, text="Embed metadata", command=self._start_embed)
|
||||||
btn.pack(anchor="w", pady=10)
|
btn.pack(anchor="w", pady=10)
|
||||||
self.action_buttons.append(btn)
|
self.action_buttons.append(btn)
|
||||||
|
|
@ -242,6 +255,8 @@ class FireMapperGUI(tk.Tk):
|
||||||
text="Embed GPS / orientation metadata into the stretched PNGs "
|
text="Embed GPS / orientation metadata into the stretched PNGs "
|
||||||
"(needs exiftool)").pack(anchor="w", pady=(6, 0))
|
"(needs exiftool)").pack(anchor="w", pady=(6, 0))
|
||||||
|
|
||||||
|
self._workers_row(tab).pack(fill="x", pady=(6, 2))
|
||||||
|
|
||||||
bar = ttk.Frame(tab)
|
bar = ttk.Frame(tab)
|
||||||
bar.pack(fill="x", pady=10)
|
bar.pack(fill="x", pady=10)
|
||||||
pv = ttk.Button(bar, text="Preview palettes", command=self._start_preview)
|
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
|
out = self.embed_out.get().strip() or None
|
||||||
gps = self.embed_gps.get()
|
gps = self.embed_gps.get()
|
||||||
et = self.embed_exif.get().strip() or None
|
et = self.embed_exif.get().strip() or None
|
||||||
|
workers = max(1, int(self.workers.get()))
|
||||||
progress, log = self._callbacks("embed")
|
progress, log = self._callbacks("embed")
|
||||||
|
|
||||||
def work():
|
def work():
|
||||||
try:
|
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))
|
self.q.put(("done", "embed", res))
|
||||||
except Exception as e: # noqa: BLE001
|
except Exception as e: # noqa: BLE001
|
||||||
self.q.put(("error", "embed", str(e)))
|
self.q.put(("error", "embed", str(e)))
|
||||||
|
|
@ -509,12 +526,13 @@ class FireMapperGUI(tk.Tk):
|
||||||
if not self._guard(session, "thermal"):
|
if not self._guard(session, "thermal"):
|
||||||
return
|
return
|
||||||
params = self._read_window_params()
|
params = self._read_window_params()
|
||||||
|
workers = max(1, int(self.workers.get()))
|
||||||
progress, log = self._callbacks("thermal")
|
progress, log = self._callbacks("thermal")
|
||||||
|
|
||||||
def work():
|
def work():
|
||||||
try:
|
try:
|
||||||
pngs = stretch.list_frames(Path(session))
|
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"Frames: {len(pngs)} | absolute signal {win['abs_lo']}-{win['abs_hi']}")
|
||||||
log(f"Window: {win['lo']}-{win['hi']} ({win['how']})")
|
log(f"Window: {win['lo']}-{win['hi']} ({win['how']})")
|
||||||
img = stretch.make_sample_image(pngs, win["lo"], win["hi"])
|
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
|
out = self.th_out.get().strip() or None
|
||||||
cmap = self.th_cmap.get()
|
cmap = self.th_cmap.get()
|
||||||
params = self._read_window_params()
|
params = self._read_window_params()
|
||||||
|
workers = max(1, int(self.workers.get()))
|
||||||
|
embed_exif = bool(self.th_embed.get())
|
||||||
progress, log = self._callbacks("thermal")
|
progress, log = self._callbacks("thermal")
|
||||||
|
|
||||||
def work():
|
def work():
|
||||||
try:
|
try:
|
||||||
pngs = stretch.list_frames(Path(session))
|
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']})")
|
log(f"Window: {win['lo']}-{win['hi']} ({win['how']})")
|
||||||
res = stretch.stretch_session(session, out, cmap, win["lo"], win["hi"],
|
res = stretch.stretch_session(session, out, cmap, win["lo"], win["hi"],
|
||||||
embed_exif=self.th_embed.get(),
|
embed_exif=embed_exif,
|
||||||
progress=progress, log=log)
|
progress=progress, log=log, workers=workers)
|
||||||
self.q.put(("done", "thermal", {"kind": "stretch", **res}))
|
self.q.put(("done", "thermal", {"kind": "stretch", **res}))
|
||||||
except Exception as e: # noqa: BLE001
|
except Exception as e: # noqa: BLE001
|
||||||
self.q.put(("error", "thermal", str(e)))
|
self.q.put(("error", "thermal", str(e)))
|
||||||
|
|
|
||||||
|
|
@ -31,6 +31,8 @@ import argparse
|
||||||
import json
|
import json
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
|
import threading
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||||
|
|
@ -90,15 +92,27 @@ def list_frames(session: Path) -> list[Path]:
|
||||||
return pngs
|
return pngs
|
||||||
|
|
||||||
|
|
||||||
def session_histogram(pngs, progress=None) -> np.ndarray:
|
def _frame_hist(p) -> np.ndarray:
|
||||||
"""Pooled value histogram over all frames (exact for integer signal data)."""
|
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)
|
hist = np.zeros(U16_MAX, dtype=np.int64)
|
||||||
n = len(pngs)
|
n = len(pngs)
|
||||||
for i, p in enumerate(pngs, 1):
|
done = 0
|
||||||
arr = np.asarray(Image.open(p)).astype(np.uint16, copy=False)
|
with ThreadPoolExecutor(max_workers=workers) as ex:
|
||||||
hist += np.bincount(arr.ravel(), minlength=U16_MAX)
|
for h in ex.map(_frame_hist, pngs):
|
||||||
|
hist += h
|
||||||
|
done += 1
|
||||||
if progress:
|
if progress:
|
||||||
progress(0.5 * i / n, f"Analysing frame {i}/{n}")
|
progress(0.5 * done / n, f"Analysing frame {done}/{n}")
|
||||||
return hist
|
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,
|
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."""
|
"""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]
|
nz = np.nonzero(hist)[0]
|
||||||
abs_lo, abs_hi = int(nz[0]), int(nz[-1])
|
abs_lo, abs_hi = int(nz[0]), int(nz[-1])
|
||||||
if lo is not None or hi is not None:
|
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
|
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
|
"""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)
|
log = log or (lambda *_: None)
|
||||||
|
workers = workers or embed.default_workers()
|
||||||
try:
|
try:
|
||||||
et = exiftool or embed.ensure_exiftool(None, Path(__file__).resolve().parent / "tools")
|
et = exiftool or embed.ensure_exiftool(None, Path(__file__).resolve().parent / "tools")
|
||||||
except BaseException as e: # noqa: BLE001 (ensure_exiftool may sys.exit)
|
except BaseException as e: # noqa: BLE001 (ensure_exiftool may sys.exit)
|
||||||
log(f" ! skipping EXIF embed - exiftool unavailable ({e})")
|
log(f" ! skipping EXIF embed - exiftool unavailable ({e})")
|
||||||
return
|
return
|
||||||
arg_lines, dsts, srcs = [], [], []
|
segments, dsts, srcs = [], [], []
|
||||||
have_sidecars = False
|
have_sidecars = False
|
||||||
for p in pngs:
|
for p in pngs:
|
||||||
dst = out_dir / (p.stem + ".png")
|
dst = out_dir / (p.stem + ".png")
|
||||||
|
|
@ -181,14 +197,11 @@ def _embed_exif(pngs, out_dir: Path, exiftool=None, log=None):
|
||||||
except Exception: # noqa: BLE001
|
except Exception: # noqa: BLE001
|
||||||
pass
|
pass
|
||||||
seg += [str(dst), "-execute"]
|
seg += [str(dst), "-execute"]
|
||||||
arg_lines += seg
|
segments.append(seg)
|
||||||
dsts.append(dst)
|
dsts.append(dst)
|
||||||
srcs.append(p)
|
srcs.append(p)
|
||||||
if not arg_lines:
|
if not segments:
|
||||||
return
|
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):
|
def has_gps(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)
|
# 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)
|
expect_gps = have_sidecars or any(has_gps(srcs[i]) for i in idx)
|
||||||
|
|
||||||
log("Embedding GPS / orientation EXIF into the stretched frames ...")
|
def run():
|
||||||
proc = subprocess.run(cmd, capture_output=True, text=True)
|
return embed.run_exiftool_parallel(et, segments, out_dir, workers, len(segments))
|
||||||
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)
|
|
||||||
|
|
||||||
err = (proc.stderr or "").strip()
|
log(f"Embedding GPS / orientation EXIF into the stretched frames ({workers} workers) ...")
|
||||||
if err:
|
rc, lines, _ = run()
|
||||||
log(err)
|
if expect_gps and rc == 0 and not all(has_gps(dsts[i]) for i in idx):
|
||||||
if proc.returncode != 0 or (expect_gps 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(
|
raise RuntimeError(
|
||||||
"Could not embed GPS/EXIF into the stretched PNGs (exiftool exit "
|
"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. "
|
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."
|
"The stretched images were written; only the embedded metadata is missing.")
|
||||||
+ (f"\nexiftool: {err}" if err else ""))
|
|
||||||
if not expect_gps:
|
if not expect_gps:
|
||||||
log(" note: no JSON sidecars and the source images carry no GPS - the stretched "
|
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), "
|
"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,
|
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
|
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
|
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
|
(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)
|
session = Path(session)
|
||||||
log = log or (lambda *_: None)
|
log = log or (lambda *_: None)
|
||||||
progress = progress or (lambda *_: None)
|
progress = progress or (lambda *_: None)
|
||||||
|
workers = workers or embed.default_workers()
|
||||||
pngs = list_frames(session)
|
pngs = list_frames(session)
|
||||||
out_dir = Path(out_dir) if out_dir else session / "thermal_stretched"
|
out_dir = Path(out_dir) if out_dir else session / "thermal_stretched"
|
||||||
out_dir.mkdir(parents=True, exist_ok=True)
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
n = len(pngs)
|
n = len(pngs)
|
||||||
log(f"Writing {n} {colormap} frames to {out_dir} ...")
|
log(f"Writing {n} {colormap} frames to {out_dir} ({workers} workers) ...")
|
||||||
for i, p in enumerate(pngs, 1):
|
lock = threading.Lock()
|
||||||
|
state = {"done": 0}
|
||||||
|
|
||||||
|
def render_one(p):
|
||||||
render_frame(p, lo, hi, colormap).save(out_dir / (p.stem + ".png"))
|
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:
|
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")
|
progress(1.0, "Done")
|
||||||
log(f"Done. {n} frames written.\nOutput: {out_dir}")
|
log(f"Done. {n} frames written.\nOutput: {out_dir}")
|
||||||
return {"out_dir": out_dir, "frames": n}
|
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")
|
help="do not write GPS/orientation EXIF into the stretched PNGs")
|
||||||
ap.add_argument("--sample", action="store_true",
|
ap.add_argument("--sample", action="store_true",
|
||||||
help="write one side-by-side palette comparison and exit")
|
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()
|
args = ap.parse_args()
|
||||||
|
|
||||||
session = find_session(args.session)
|
session = find_session(args.session)
|
||||||
pngs = list_frames(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"]
|
lo, hi = win["lo"], win["hi"]
|
||||||
print(f"Session : {session}")
|
print(f"Session : {session}")
|
||||||
print(f"Frames : {len(pngs)} | absolute signal range {win['abs_lo']}-{win['abs_hi']}")
|
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
|
out_dir = Path(args.out).resolve() if args.out else None
|
||||||
res = stretch_session(session, out_dir, args.colormap, lo, hi,
|
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']}")
|
print(f"\nWrote {res['frames']} {args.colormap} frames to:\n {res['out_dir']}")
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue