FireMapper_Postprocess/embed_metadata.py

393 lines
17 KiB
Python

#!/usr/bin/env python3
"""
embed_metadata.py - Write FireMapper per-frame JSON sidecars into image EXIF/XMP.
For every image in a capture session that has a same-basename `.json` sidecar
(cam*/<ts>_step_NN.tiff and thermal/<ts>_step_NN.png), this:
* maps the structured fields to real EXIF/XMP tags so the images are
georeferenced and usable in mapping / photogrammetry software:
- GPS lat/lon/alt -> EXIF GPS* tags
- platform yaw -> EXIF GPSImgDirection + XMP-Camera:PoseHeadingDegrees
- pitch / roll -> XMP-Camera:PosePitch/PoseRollDegrees
- GPS week/tow -> DateTimeOriginal + GPSDateStamp/GPSTimeStamp (UTC)
- lens / exposure / camera -> FocalLength, ExposureTime, Make/Model/Serial
* stores the COMPLETE original JSON in EXIF:UserComment so nothing is lost.
Originals are never modified: tagged copies are written to an output folder
that mirrors the session layout.
Usage:
python embed_metadata.py [SESSION_DIR] [options]
SESSION_DIR session folder (default: the single session_* dir here)
--out DIR output folder (default: <session>_exif next to it)
--gps-source {position,gps}
which block feeds GPS coords for cam frames
(default: position = fused INS, falls back to gps)
--exiftool PATH path to exiftool.exe (default: auto-detect / download)
--dry-run report what would be written, touch nothing
"""
from __future__ import annotations
import argparse
import datetime as dt
import json
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))
import session_map as smap # noqa: E402 (shared camera geometry + session discovery)
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"]
# --------------------------------------------------------------------------- #
# exiftool acquisition
# --------------------------------------------------------------------------- #
def _exiftool_candidates(cache_dir: Path):
"""Yield likely exiftool locations in priority order."""
yield shutil.which("exiftool")
yield str(cache_dir / "exiftool.exe")
local = os.environ.get("LOCALAPPDATA", "")
pf = os.environ.get("ProgramFiles", "")
for base in (local and Path(local) / "Programs" / "ExifTool", pf and Path(pf) / "ExifTool"):
if base:
yield str(base / "ExifTool.exe")
yield str(base / "exiftool.exe")
def ensure_exiftool(explicit: str | None, cache_dir: Path) -> str:
"""Return a path to a working exiftool, installing it via winget if needed."""
if explicit:
return explicit
for cand in _exiftool_candidates(cache_dir):
if cand and Path(cand).exists():
return cand
# not found anywhere - try a one-time install via winget (Windows)
if shutil.which("winget"):
print("exiftool not found - installing via winget (OliverBetz.ExifTool)...")
subprocess.run(
["winget", "install", "--id", "OliverBetz.ExifTool",
"--accept-package-agreements", "--accept-source-agreements", "--silent"],
check=False,
)
for cand in _exiftool_candidates(cache_dir):
if cand and Path(cand).exists():
return cand
sys.exit(
"exiftool could not be found or installed automatically.\n"
"Install it (e.g. `winget install OliverBetz.ExifTool` or `choco install exiftool`),\n"
"or pass its path with --exiftool C:\\path\\to\\exiftool.exe"
)
# --------------------------------------------------------------------------- #
# JSON -> tag mapping
# --------------------------------------------------------------------------- #
def gps_time_to_utc(week, tow) -> dt.datetime | None:
if week is None or tow is None:
return None
return GPS_EPOCH + dt.timedelta(seconds=week * 604800 + tow - GPS_UTC_LEAP_SECONDS)
def signed_ref(value, pos, neg):
"""Return (abs_value, ref_letter) for a signed coordinate."""
return (abs(value), pos if value >= 0 else neg)
def build_tags(data: dict, manifest_cams: dict, cam_id: str | None, gps_source: str) -> list[str]:
"""Return a list of '-Tag=value' exiftool arguments for one frame."""
tags: list[str] = []
# --- coordinates: prefer the requested source, fall back to raw gps ----
pos = None
if cam_id and gps_source == "position":
pos = data.get("position")
if not pos:
pos = data.get("gps")
if pos and pos.get("lat") is not None and pos.get("lon") is not None:
lat, lat_ref = signed_ref(pos["lat"], "N", "S")
lon, lon_ref = signed_ref(pos["lon"], "E", "W")
tags += [f"-GPSLatitude={lat}", f"-GPSLatitudeRef={lat_ref}",
f"-GPSLongitude={lon}", f"-GPSLongitudeRef={lon_ref}"]
if pos.get("alt") is not None:
alt = pos["alt"]
tags += [f"-GPSAltitude={abs(alt)}", f"-GPSAltitudeRef={0 if alt >= 0 else 1}"]
# --- orientation: TRUE camera pointing ---------------------------------
# aircraft attitude + platform roll (platform_angle_deg) + camera mounting offset,
# not the bare IMU. Computed by the shared geometry so EXIF matches the map.
imu = data.get("imu") or {}
plat = data.get("platform") or {}
ori = None
if None not in (imu.get("yaw"), imu.get("pitch"), imu.get("roll")):
ori = smap.camera_orientation({
"kind": "cam" if cam_id else "thermal",
"off_nadir": smap.off_nadir_from_name(cam_id) if cam_id else 0.0,
"yaw": imu["yaw"], "pitch": imu["pitch"], "roll": imu["roll"],
"platform_angle": plat.get("platform_angle_deg"),
})
if ori:
heading, pitch, roll = ori
tags += [f"-GPSImgDirection={heading % 360:.4f}", "-GPSImgDirectionRef=T",
f"-XMP-GPano:PoseHeadingDegrees={heading % 360:.4f}",
f"-XMP-GPano:PosePitchDegrees={pitch:.4f}",
f"-XMP-GPano:PoseRollDegrees={roll:.4f}"]
# --- timestamp (UTC, derived from GPS week/tow) ------------------------
g = data.get("gps") or {}
when = gps_time_to_utc(g.get("week"), g.get("tow"))
if when:
tags += [f"-DateTimeOriginal={when:%Y:%m:%d %H:%M:%S}",
f"-SubSecTimeOriginal={when.microsecond // 1000:03d}",
"-OffsetTimeOriginal=+00:00",
f"-GPSDateStamp={when:%Y:%m:%d}",
f"-GPSTimeStamp={when:%H:%M:%S}"]
# --- camera / lens -----------------------------------------------------
lens = data.get("lens") or {}
if lens.get("focal_length_mm") is not None:
tags.append(f"-FocalLength={lens['focal_length_mm']}")
if data.get("exposure_time_us") is not None:
tags.append(f"-ExposureTime={data['exposure_time_us'] / 1_000_000.0:.9f}")
if cam_id and cam_id in manifest_cams:
cam = manifest_cams[cam_id]
tags += ["-Make=Allied Vision", f"-Model={cam.get('model', cam_id)}",
f"-SerialNumber={cam.get('serial', '')}"]
# --- full fidelity: entire JSON in UserComment + a short description ----
compact = json.dumps(data, separators=(",", ":"), sort_keys=True)
label_bits = [p for p in (cam_id or ("thermal" if "thermal" in data else None),) if p]
if g.get("week") is not None:
label_bits.append(f"frame {data.get('frame_id', '?')}")
tags += [f"-EXIF:UserComment={compact}",
"-Software=FireMapper embed_metadata.py",
f"-EXIF:ImageDescription={' '.join(label_bits) or 'FireMapper frame'}"]
return tags
# --------------------------------------------------------------------------- #
# core (importable - used by both the CLI and the GUI)
# --------------------------------------------------------------------------- #
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 collect_pairs(session: Path) -> list[tuple[Path, Path]]:
"""Every image under the session that has a same-basename .json sidecar."""
pairs = []
for img in sorted(session.rglob("*")):
if img.suffix.lower() in IMAGE_EXTS and img.with_suffix(".json").exists():
pairs.append((img, img.with_suffix(".json")))
return pairs
def load_manifest_cams(session: Path) -> dict:
"""camera id -> manifest entry (model / serial), if a manifest exists."""
cams: dict = {}
mpath = session / "manifest.json"
if mpath.exists():
try:
for cam in json.loads(mpath.read_text(encoding="utf-8")).get("cameras", []):
cams[cam["id"]] = cam
except Exception: # noqa: BLE001
pass
return cams
def embed_session(session, out_root=None, gps_source="position", exiftool=None,
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. `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:
raise RuntimeError(f"No image+json pairs found under {session}")
cams = load_manifest_cams(session)
et = exiftool or ensure_exiftool(None, Path(__file__).resolve().parent / "tools")
total = len(pairs)
log(f"Session : {session}")
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 ...")
# 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})")
return None
dst = out_root / rel
shutil.copy2(img, dst)
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"]
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 ...")
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 {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")
log(f"\nDone. {updated} files tagged.\nOutput: {out_root}")
return {"out_root": out_root, "pairs": total, "updated": updated}
# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def main() -> int:
ap = argparse.ArgumentParser(description="Embed FireMapper JSON sidecars into image EXIF/XMP.")
ap.add_argument("session", nargs="?", help="session folder (default: newest session_* here)")
ap.add_argument("--out", help="output folder (default: <session>_exif)")
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()
session = find_session(args.session)
if not session.is_dir():
sys.exit(f"Session folder not found: {session}")
if args.dry_run:
pairs = collect_pairs(session)
cams = load_manifest_cams(session)
if not pairs:
sys.exit(f"No image+json pairs found under {session}")
img, sidecar = pairs[0]
cam_id = img.relative_to(session).parts[0]
cam_id = cam_id if cam_id.startswith("cam") else None
data = json.loads(sidecar.read_text(encoding="utf-8"))
print(f"DRY RUN: {len(pairs)} pairs under {session}")
print("\n--- sample tags (first frame) ---")
print("\n".join(build_tags(data, cams, cam_id, args.gps_source)))
return 0
embed_session(session, args.out, args.gps_source, args.exiftool,
progress=lambda f, m: None, log=print, workers=args.workers)
return 0
if __name__ == "__main__":
raise SystemExit(main())