extract_frames.py

#!/usr/bin/env python3
"""
extract_frames.py

Extract geotagged frames from a DJI drone video for photogrammetry.

Key improvements over extract_and_geotag.py:
  - Time-interval-based sampling (--interval-sec): picks sharpest frame in each N-second window
  - EXIF GPS injection via piexif so OpenDroneMap reads coords directly from JPEG
  - Camera EXIF metadata (Make/Model/FocalLength) for ODM sensor DB lookup
  - Video-prefixed output naming to avoid collisions when processing multiple videos
  - Relaxed default similarity threshold (0.98) for photogrammetry spatial coverage

Usage:
    python extract_frames.py \
        --video /path/to/DJI_0925.MP4 \
        --telemetry /path/to/DJI_0925_telemetry.csv \
        --out photogrammetry_dataset \
        --interval-sec 1.0

Output:
    <out>/images/<prefix>_f000000.jpg  (JPEG with GPS EXIF)
    <out>/dataset_summary.csv   (appended per run)
    <out>/dataset_geotags.csv   (appended per run)
"""

import argparse
import bisect
import csv
import datetime
import fcntl
import json
import math
import os
import shutil
import struct
import subprocess
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any

import cv2
import numpy as np
import piexif
import piexif.helper

cv2.setNumThreads(1)  # Let ThreadPoolExecutor handle concurrency, not OpenCV


@dataclass
class Telemetry:
    t_ms_srt: int
    lat: float
    lon: float
    abs_alt: float
    rel_alt: float
    focal_len: float  # 35mm-equivalent focal length
    fnum: float


def haversine_m(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """Great-circle distance in metres between two WGS84 points."""
    R = 6_371_000  # Earth mean radius, metres
    rlat1, rlat2 = math.radians(lat1), math.radians(lat2)
    dlat = math.radians(lat2 - lat1)
    dlon = math.radians(lon2 - lon1)
    a = math.sin(dlat / 2) ** 2 + math.cos(rlat1) * math.cos(rlat2) * math.sin(dlon / 2) ** 2
    return R * 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))


def _center_crop(frame: np.ndarray, scale: float) -> np.ndarray:
    """Return a 1:1 center crop of the frame.

    The crop dimensions are (h*scale, w*scale) — same pixel count as
    downsampling by `scale`, but preserving native resolution for accurate
    Laplacian blur scoring.  Returns a numpy view (O(1), no copy).
    """
    h, w = frame.shape[:2]
    ch, cw = int(h * scale), int(w * scale)
    y0 = (h - ch) // 2
    x0 = (w - cw) // 2
    return frame[y0 : y0 + ch, x0 : x0 + cw]


def load_telemetry_csv(path: Path) -> List[Telemetry]:
    rows: List[Telemetry] = []
    with path.open("r", encoding="utf-8") as f:
        r = csv.DictReader(f)
        for row in r:
            rows.append(
                Telemetry(
                    t_ms_srt=int(float(row["t_ms_srt"])),
                    lat=float(row["latitude"]),
                    lon=float(row["longitude"]),
                    abs_alt=float(row["abs_alt"]),
                    rel_alt=float(row["rel_alt"]),
                    # Older CSVs from parse_drone_srt.py may not have these columns
                    focal_len=float(row.get("focal_len") or 0),
                    fnum=float(row.get("fnum") or 0),
                )
            )
    rows.sort(key=lambda x: x.t_ms_srt)
    if not rows:
        raise RuntimeError("No telemetry rows loaded.")
    return rows


def nearest_telemetry(samples: List[Telemetry], t_ms: int) -> Telemetry:
    lo, hi = 0, len(samples) - 1
    best = samples[0]
    best_diff = abs(samples[0].t_ms_srt - t_ms)
    while lo <= hi:
        mid = (lo + hi) // 2
        diff = samples[mid].t_ms_srt - t_ms
        ad = abs(diff)
        if ad < best_diff:
            best_diff = ad
            best = samples[mid]
        if diff < 0:
            lo = mid + 1
        elif diff > 0:
            hi = mid - 1
        else:
            return samples[mid]
    return best


def interpolate_telemetry(samples: List[Telemetry], t_ms: int, _keys: Optional[List[int]] = None) -> Telemetry:
    """Linearly interpolate between the two SRT entries bracketing t_ms.

    Falls back to the nearest endpoint when t_ms is outside the telemetry range.
    The _keys list is a pre-built sorted list of t_ms_srt values (built once and
    cached by the caller for performance).
    """
    if _keys is None:
        _keys = [s.t_ms_srt for s in samples]
    idx = bisect.bisect_left(_keys, t_ms)
    if idx == 0:
        return samples[0]
    if idx >= len(samples):
        return samples[-1]
    a, b = samples[idx - 1], samples[idx]
    span = b.t_ms_srt - a.t_ms_srt
    if span == 0:
        return a
    ratio = (t_ms - a.t_ms_srt) / span
    return Telemetry(
        t_ms_srt=t_ms,
        lat=a.lat + ratio * (b.lat - a.lat),
        lon=a.lon + ratio * (b.lon - a.lon),
        abs_alt=a.abs_alt + ratio * (b.abs_alt - a.abs_alt),
        rel_alt=a.rel_alt + ratio * (b.rel_alt - a.rel_alt),
        focal_len=a.focal_len,
        fnum=a.fnum,
    )


def blur_score(frame) -> float:
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    return cv2.Laplacian(gray, cv2.CV_64F).var()


def frame_hist(frame):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    hist = cv2.calcHist([hsv], [0, 1], None, [32, 32], [0, 180, 0, 256])
    return cv2.normalize(hist, hist).flatten()


@dataclass
class WindowResult:
    """Result of selecting the best frame from a time window."""
    frame: Optional[np.ndarray]
    blur: float
    frame_idx: int
    is_fallback: bool
    window_best_score: float  # best blur score seen in window (for diagnostics)


def select_window_frame(
    frames: List[Tuple[int, np.ndarray]],
    score_scale: float,
    blur_thresh: float,
    blur_floor: float,
) -> WindowResult:
    """Select the best frame from a list of (frame_idx, frame) candidates.

    Returns a WindowResult. result.frame is None if no candidate passes blur_floor.
    """
    best_frame = None
    best_blur = -1.0
    best_idx = -1
    fallback_frame = None
    fallback_blur = -1.0
    fallback_idx = -1
    window_best_score = -1.0

    for idx, frame in frames:
        crop = _center_crop(frame, score_scale)
        b = blur_score(crop)
        if b > window_best_score:
            window_best_score = b
        if b >= blur_thresh and b > best_blur:
            best_blur = b
            best_frame = frame
            best_idx = idx
        elif b >= blur_floor and b > fallback_blur:
            fallback_blur = b
            fallback_frame = frame
            fallback_idx = idx

    if best_frame is not None:
        return WindowResult(best_frame, best_blur, best_idx, False, window_best_score)
    if fallback_frame is not None:
        return WindowResult(fallback_frame, fallback_blur, fallback_idx, True, window_best_score)
    return WindowResult(None, -1.0, -1, False, window_best_score)


def _deg_to_rational(value: float) -> Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]:
    """
    Convert decimal degrees to EXIF GPS rational (deg, min, sec) tuple.
    Each component is (numerator, denominator).
    Uses high-precision seconds to preserve ~1m GPS accuracy.
    """
    deg = int(abs(value))
    minutes_f = (abs(value) - deg) * 60.0
    minutes = int(minutes_f)
    seconds_f = (minutes_f - minutes) * 60.0
    # Store seconds with denominator 100_000 for sub-mm precision;
    # eliminates EXIF quantization as a GPS accuracy bottleneck.
    seconds_num = round(seconds_f * 100_000)
    return (deg, 1), (minutes, 1), (seconds_num, 100_000)


def _float_to_rational(value: float, denom: int = 100) -> Tuple[int, int]:
    return (round(abs(value) * denom), denom)


def build_gps_exif(
    tel: Telemetry,
    alt_source: str = "abs",
    frame_dt: Optional[datetime.datetime] = None,
) -> Dict:
    """Build piexif GPS IFD dict from a Telemetry sample.

    alt_source: 'abs' uses abs_alt (WGS84 ellipsoidal), 'rel' uses rel_alt (relative to takeoff).
    frame_dt: if provided, adds GPSDateStamp and GPSTimeStamp.
    """
    lat_rational = _deg_to_rational(tel.lat)
    lon_rational = _deg_to_rational(tel.lon)
    alt = tel.abs_alt if alt_source == "abs" else tel.rel_alt
    alt_rational = _float_to_rational(alt, denom=100)

    gps_ifd: Dict[int, Any] = {
        piexif.GPSIFD.GPSVersionID: (2, 3, 0, 0),
        piexif.GPSIFD.GPSLatitudeRef: b"N" if tel.lat >= 0 else b"S",
        piexif.GPSIFD.GPSLatitude: lat_rational,
        piexif.GPSIFD.GPSLongitudeRef: b"E" if tel.lon >= 0 else b"W",
        piexif.GPSIFD.GPSLongitude: lon_rational,
        piexif.GPSIFD.GPSAltitudeRef: 0,  # 0 = above sea level
        piexif.GPSIFD.GPSAltitude: alt_rational,
        piexif.GPSIFD.GPSMeasureMode: b"3",  # 3D fix
    }
    if frame_dt is not None:
        gps_ifd[piexif.GPSIFD.GPSDateStamp] = frame_dt.strftime("%Y:%m:%d").encode()
        sec_frac = frame_dt.second + frame_dt.microsecond / 1_000_000
        gps_ifd[piexif.GPSIFD.GPSTimeStamp] = (
            (frame_dt.hour, 1),
            (frame_dt.minute, 1),
            (round(sec_frac * 1000), 1000),
        )
    return gps_ifd


def build_exif_bytes(
    tel: Telemetry,
    image_width: int,
    image_height: int,
    camera_make: str = "DJI",
    camera_model: str = "FC2204",
    crop_factor: float = 2.0,
    alt_source: str = "abs",
    frame_dt: Optional[datetime.datetime] = None,
) -> bytes:
    """Build a complete piexif EXIF blob to embed in a JPEG.

    focal_len from the SRT is the 35mm-equivalent focal length (e.g. 24mm for Mavic 3).
    The actual focal length is derived via: actual_fl = equiv_fl / crop_factor.
    """
    actual_fl_mm = (tel.focal_len / crop_factor) if tel.focal_len > 0 else (24.0 / crop_factor)
    fl_rational = _float_to_rational(actual_fl_mm, denom=100)
    equiv_fl = round(tel.focal_len) if tel.focal_len > 0 else 24

    zeroth_ifd = {
        piexif.ImageIFD.Make: camera_make.encode(),
        piexif.ImageIFD.Model: camera_model.encode(),
        piexif.ImageIFD.XResolution: (72, 1),
        piexif.ImageIFD.YResolution: (72, 1),
        piexif.ImageIFD.ResolutionUnit: 2,      # inch
        piexif.ImageIFD.ImageWidth: image_width,
        piexif.ImageIFD.ImageLength: image_height,
    }

    exif_ifd: Dict[int, Any] = {
        piexif.ExifIFD.FocalLength: fl_rational,
        piexif.ExifIFD.FocalLengthIn35mmFilm: equiv_fl,
        piexif.ExifIFD.FNumber: _float_to_rational(tel.fnum if tel.fnum > 0 else 2.8, denom=10),
        piexif.ExifIFD.PixelXDimension: image_width,
        piexif.ExifIFD.PixelYDimension: image_height,
    }
    if frame_dt is not None:
        exif_ifd[piexif.ExifIFD.DateTimeOriginal] = frame_dt.strftime("%Y:%m:%d %H:%M:%S").encode()
        subsec = f"{frame_dt.microsecond // 1000:03d}"
        exif_ifd[piexif.ExifIFD.SubSecTimeOriginal] = subsec.encode()

    gps_ifd = build_gps_exif(tel, alt_source=alt_source, frame_dt=frame_dt)

    exif_dict = {
        "0th": zeroth_ifd,
        "Exif": exif_ifd,
        "GPS": gps_ifd,
    }
    return piexif.dump(exif_dict)


def write_jpeg_with_exif(frame, out_path: Path, exif_bytes: bytes, quality: int = 95) -> None:
    """
    Write frame as JPEG then insert the EXIF blob.
    piexif.insert(exif_bytes, src_path, dst_path) works in-place so we write
    the bare JPEG first, then overwrite it with the EXIF-tagged version.
    """
    cv2.imwrite(str(out_path), frame, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
    piexif.insert(exif_bytes, str(out_path), str(out_path))


def _probe_codec(video_path: Path) -> str:
    """Return the video codec name (e.g. 'h264', 'hevc') via ffprobe."""
    if not shutil.which("ffprobe"):
        return "h264"
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "quiet", "-select_streams", "v:0",
             "-show_entries", "stream=codec_name", "-of", "csv=p=0",
             str(video_path)],
            capture_output=True, text=True, timeout=10,
        )
        codec = r.stdout.strip()
        return codec if codec else "h264"
    except Exception:
        return "h264"


def _cuvid_available(codec: str = "h264") -> bool:
    """Return True if system ffmpeg has the cuvid decoder for the given codec."""
    if not shutil.which("ffmpeg"):
        return False
    cuvid_name = {"h264": "h264_cuvid", "hevc": "hevc_cuvid"}.get(codec, f"{codec}_cuvid")
    try:
        r = subprocess.run(
            ["ffmpeg", "-decoders"], capture_output=True, text=True, timeout=5
        )
        return cuvid_name in r.stdout
    except Exception:
        return False


class FFmpegReader:
    """
    Drop-in replacement for cv2.VideoCapture that uses the system ffmpeg binary.

    Decodes via NVIDIA NVDEC (h264_cuvid / hevc_cuvid) when use_gpu=True,
    falling back to multi-threaded CPU decode.  Emits raw BGR24 frames via a
    subprocess pipe so the rest of the pipeline is unchanged.

    Interface mirrors cv2.VideoCapture: isOpened(), get(prop), read(), release().
    """

    def __init__(self, video_path: Path, use_gpu: bool = True) -> None:
        self.video_path = video_path
        self._width = 0
        self._height = 0
        self._fps = 0.0
        self._frame_count = 0
        self._frame_bytes = 0
        self._codec = "h264"
        self._creation_time: Optional[str] = None
        self._proc: Optional[subprocess.Popen] = None
        self._opened = False

        self._probe()
        self._start(use_gpu)

    def _probe(self) -> None:
        """Populate fps, dimensions, frame_count, codec, and creation_time via ffprobe."""
        result = subprocess.run(
            [
                "ffprobe", "-v", "quiet",
                "-print_format", "json",
                "-show_streams", "-show_format",
                str(self.video_path),
            ],
            capture_output=True, text=True, timeout=15,
        )
        data = json.loads(result.stdout)
        for stream in data.get("streams", []):
            if stream.get("codec_type") != "video":
                continue
            self._width = int(stream["width"])
            self._height = int(stream["height"])
            self._codec = stream.get("codec_name", "h264")
            fps_str = stream.get("avg_frame_rate") or stream.get("r_frame_rate", "60/1")
            num, den = map(int, fps_str.split("/"))
            self._fps = num / den if den else 0.0
            if "nb_frames" in stream:
                self._frame_count = int(stream["nb_frames"])
            else:
                dur = float(stream.get("duration", 0))
                self._frame_count = int(dur * self._fps)
            self._frame_bytes = self._width * self._height * 3  # BGR24
            break
        # Extract creation_time from format tags (ISO 8601, e.g. "2025-10-17T14:32:15.000000Z")
        fmt_tags = data.get("format", {}).get("tags", {})
        self._creation_time = fmt_tags.get("creation_time")

    def _start(self, use_gpu: bool) -> None:
        """Launch the ffmpeg subprocess."""
        if use_gpu:
            cuvid_map = {"h264": "h264_cuvid", "hevc": "hevc_cuvid"}
            decoder = cuvid_map.get(self._codec, f"{self._codec}_cuvid")
            cmd = [
                "ffmpeg",
                "-hwaccel", "cuda",
                "-c:v", decoder,
                "-i", str(self.video_path),
                "-f", "rawvideo", "-pix_fmt", "bgr24",
                "pipe:1",
            ]
        else:
            cmd = [
                "ffmpeg",
                "-threads", "4",
                "-i", str(self.video_path),
                "-f", "rawvideo", "-pix_fmt", "bgr24",
                "pipe:1",
            ]
        self._proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            bufsize=self._frame_bytes * 4,
        )
        self._opened = True

    # ── cv2.VideoCapture-compatible interface ──────────────────────────────

    def isOpened(self) -> bool:
        return self._opened and self._proc is not None and self._proc.poll() is None

    def get(self, prop: int) -> float:
        if prop == cv2.CAP_PROP_FPS:
            return self._fps
        if prop == cv2.CAP_PROP_FRAME_COUNT:
            return float(self._frame_count)
        if prop == cv2.CAP_PROP_FRAME_WIDTH:
            return float(self._width)
        if prop == cv2.CAP_PROP_FRAME_HEIGHT:
            return float(self._height)
        return 0.0

    def read(self) -> Tuple[bool, Optional[np.ndarray]]:
        if not self._proc:
            return False, None
        raw = self._proc.stdout.read(self._frame_bytes)
        if len(raw) < self._frame_bytes:
            return False, None
        frame = np.frombuffer(raw, dtype=np.uint8).reshape(
            self._height, self._width, 3
        ).copy()  # copy to own the memory; frombuffer gives a read-only view
        return True, frame

    def release(self) -> None:
        if self._proc:
            try:
                self._proc.stdout.close()
            except OSError:
                pass
            self._proc.wait()
            self._proc = None
        self._opened = False


def main():
    ap = argparse.ArgumentParser(
        description="Extract geotagged frames from DJI drone video for photogrammetry."
    )
    ap.add_argument("--video", type=Path, required=True, help="Input MP4 video file")
    ap.add_argument(
        "--telemetry", type=Path, required=True,
        help="Telemetry CSV from parse_drone_srt.py"
    )
    ap.add_argument("--out", type=Path, required=True, help="Output directory for frames and CSVs")
    ap.add_argument(
        "--prefix", type=str, default=None,
        help="Prefix for output frame filenames (default: video stem digits, e.g. '0925')"
    )
    ap.add_argument(
        "--interval-sec", type=float, default=1.0,
        help="Extract the sharpest frame from each N-second window (default: 1.0). "
             "Use 0 to disable interval mode and use --frame-step instead."
    )
    ap.add_argument(
        "--frame-step", type=int, default=5,
        help="Fallback: consider every Nth frame when --interval-sec is 0 (default: 5)"
    )
    ap.add_argument(
        "--blur-thresh", type=float, default=150.0,
        help="Reject frames with Laplacian blur score below this (default: 150.0)"
    )
    ap.add_argument(
        "--blur-floor", type=float, default=50.0,
        help="Fallback blur threshold: when no frame in a window passes --blur-thresh, "
             "accept the sharpest frame if it scores above this floor (default: 50.0). "
             "Set to 0 to accept any frame as fallback. Set equal to --blur-thresh to disable."
    )
    ap.add_argument(
        "--sim-thresh", type=float, default=0.98,
        help="Secondary duplicate filter: reject if HSV histogram correlation to previous "
             "accepted frame > this (default: 0.98). Runs after the distance check "
             "(--min-distance-m). Set to 1.0 to disable and rely solely on distance."
    )
    ap.add_argument(
        "--min-distance-m", type=float, default=3.0,
        help="Primary duplicate filter: reject frame if GPS distance to previous accepted "
             "frame is less than this (metres, default: 3.0). Set to 0 to disable distance "
             "filtering. More robust than histogram similarity over homogeneous terrain."
    )
    ap.add_argument(
        "--max-telemetry-delta-ms", type=int, default=250,
        help="Warn if nearest telemetry is farther than this many ms (default: 250)"
    )
    ap.add_argument(
        "--jpeg-quality", type=int, default=95,
        help="JPEG output quality 1-100 (default: 95)"
    )
    ap.add_argument(
        "--start-idx", type=int, default=0,
        help="Starting index for frame filenames (for global numbering across videos)"
    )
    ap.add_argument(
        "--score-scale", type=float, default=0.25,
        help="Scoring region scale (default: 0.25). For blur: 1:1 center crop of "
             "(h*scale, w*scale) pixels preserving native detail. For histogram: full-frame "
             "resize to (h*scale, w*scale). Does not affect JPEG output quality."
    )
    ap.add_argument(
        "--gpu-decode", action=argparse.BooleanOptionalAction, default=True,
        help="Use system ffmpeg NVDEC for video decode (default: True). "
             "Use --no-gpu-decode for CPU-only OpenCV decode."
    )
    ap.add_argument(
        "--alt-source", choices=["abs", "rel"], default="abs",
        help="Which altitude to write in GPS EXIF: 'abs' = absolute/ellipsoidal "
             "(DJI abs_alt, WGS84), 'rel' = relative to takeoff. Default: abs."
    )
    ap.add_argument(
        "--creation-time", type=str, default=None,
        help="Video creation time (ISO 8601, e.g. '2025-10-17T14:32:15.000000Z'). "
             "Auto-detected from video metadata when using --gpu-decode."
    )
    ap.add_argument(
        "--camera-make", type=str, default="DJI",
        help="EXIF Make tag (default: DJI)"
    )
    ap.add_argument(
        "--camera-model", type=str, default="FC2204",
        help="EXIF Model tag for ODM sensor DB lookup (default: FC2204 = DJI Mavic 3)"
    )
    ap.add_argument(
        "--crop-factor", type=float, default=2.0,
        help="Sensor crop factor for deriving actual FL from 35mm equiv (default: 2.0)"
    )

    args = ap.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    (args.out / "images").mkdir(exist_ok=True)

    # Derive prefix from video stem if not specified (e.g. "DJI_0925" -> "0925")
    prefix = args.prefix
    if prefix is None:
        stem = args.video.stem  # e.g. "DJI_0925"
        digits = "".join(c for c in stem if c.isdigit())
        prefix = digits if digits else stem

    telemetry = load_telemetry_csv(args.telemetry)
    telemetry_keys = [s.t_ms_srt for s in telemetry]

    # Probe video codec to pick the right NVDEC decoder
    video_codec = _probe_codec(args.video)
    use_gpu = args.gpu_decode and _cuvid_available(video_codec)
    if args.gpu_decode and not use_gpu:
        cuvid_name = {"h264": "h264_cuvid", "hevc": "hevc_cuvid"}.get(video_codec, f"{video_codec}_cuvid")
        print(f"  [decode] {cuvid_name} not available, falling back to OpenCV CPU decode")

    if use_gpu:
        cap = FFmpegReader(args.video, use_gpu=True)
        cuvid_name = {"h264": "h264_cuvid", "hevc": "hevc_cuvid"}.get(video_codec, f"{video_codec}_cuvid")
        decode_mode = f"ffmpeg/NVDEC ({cuvid_name})"
    else:
        cap = cv2.VideoCapture(str(args.video))
        decode_mode = "OpenCV/CPU"

    if not cap.isOpened():
        raise RuntimeError(f"Cannot open video: {args.video}")

    # Resolve video creation time for EXIF timestamps
    video_start_dt: Optional[datetime.datetime] = None
    if args.creation_time:
        video_start_dt = datetime.datetime.fromisoformat(args.creation_time.replace("Z", "+00:00"))
    elif use_gpu and isinstance(cap, FFmpegReader) and cap._creation_time:
        video_start_dt = datetime.datetime.fromisoformat(cap._creation_time.replace("Z", "+00:00"))

    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps <= 0:
        raise RuntimeError("Could not read FPS from video.")
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    use_interval = args.interval_sec > 0
    interval_frames = max(1, round(fps * args.interval_sec)) if use_interval else args.frame_step

    print(
        f"Video: {args.video.name}  FPS={fps:.3f}  Frames={total_frames}  "
        f"Duration={total_frames/fps:.1f}s  Decode={decode_mode}"
    )
    if use_interval:
        floor_info = f"  blur-floor={args.blur_floor}" if args.blur_floor < args.blur_thresh else ""
        print(
            f"Mode: interval={args.interval_sec}s "
            f"({interval_frames} frames/window)  "
            f"blur>={args.blur_thresh}{floor_info}  min-dist>={args.min_distance_m}m  "
            f"sim<={args.sim_thresh}  score-scale={args.score_scale}"
        )
    else:
        print(
            f"Mode: frame-step={args.frame_step}  "
            f"blur>={args.blur_thresh}  min-dist>={args.min_distance_m}m  "
            f"sim<={args.sim_thresh}  score-scale={args.score_scale}"
        )

    frames_meta: List[Dict[str, Any]] = []
    geotags: List[Dict[str, Any]] = []
    pending_writes: List[tuple] = []  # (future, meta_dict, geotag_dict)

    prev_hist = None
    prev_tel: Optional[Telemetry] = None
    saved_idx = args.start_idx
    frame_idx = 0
    fallback_count = 0
    gap_count = 0

    with ThreadPoolExecutor(max_workers=4) as write_pool:
        if use_interval:
            # Interval mode: scan all frames in each window, keep the sharpest one that
            # passes blur threshold (or blur floor as fallback), distance check, and
            # similarity check.
            window_start_frame = 0

            while True:
                window_end_frame = window_start_frame + interval_frames

                # Collect frames for this window
                window_frames: List[Tuple[int, np.ndarray]] = []
                while frame_idx < window_end_frame:
                    ret, frame = cap.read()
                    if not ret:
                        break
                    window_frames.append((frame_idx, frame.copy()))
                    frame_idx += 1

                if not window_frames and frame_idx < window_end_frame:
                    break  # End of video

                # Select best frame using primary threshold + fallback floor
                wr = select_window_frame(
                    window_frames, args.score_scale,
                    args.blur_thresh, args.blur_floor,
                )

                if wr.frame is None:
                    if window_frames:
                        # No frame even above blur_floor — log the gap
                        window_center_ms = int(
                            (((window_start_frame + window_end_frame) / 2) / fps) * 1000.0
                        )
                        gap_tel = interpolate_telemetry(telemetry, window_center_ms, telemetry_keys)
                        print(
                            f"  GAP: window {window_start_frame}-{window_end_frame} "
                            f"(t={window_center_ms}ms) — no frame above blur_floor="
                            f"{args.blur_floor:.0f}, best_score={wr.window_best_score:.1f}, "
                            f"GPS=({gap_tel.lat:.5f}, {gap_tel.lon:.5f}), "
                            f"alt={gap_tel.rel_alt:.1f}m"
                        )
                        gap_count += 1
                    window_start_frame = window_end_frame
                    continue

                best_frame = wr.frame
                best_blur = wr.blur
                best_frame_idx = wr.frame_idx
                is_fallback = wr.is_fallback
                if is_fallback:
                    fallback_count += 1

                # Compute telemetry for the candidate (needed for distance check)
                t_ms = int((best_frame_idx / fps) * 1000.0)
                tel = interpolate_telemetry(telemetry, t_ms, telemetry_keys)
                tel_delta = abs(tel.t_ms_srt - t_ms)

                # Primary filter: GPS distance to previous accepted frame
                dist: Optional[float] = None
                if args.min_distance_m > 0 and prev_tel is not None:
                    dist = haversine_m(prev_tel.lat, prev_tel.lon, tel.lat, tel.lon)
                    if dist < args.min_distance_m:
                        window_start_frame = window_end_frame
                        continue

                # Secondary filter: HSV histogram similarity
                best_sf = cv2.resize(best_frame, (0, 0), fx=args.score_scale, fy=args.score_scale,
                                     interpolation=cv2.INTER_AREA)
                h = frame_hist(best_sf)
                if prev_hist is not None:
                    sim = cv2.compareHist(prev_hist, h, cv2.HISTCMP_CORREL)
                    if sim > args.sim_thresh:
                        window_start_frame = window_end_frame
                        continue
                else:
                    sim = None

                prev_hist = h
                prev_tel = tel

                h_px, w_px = best_frame.shape[:2]
                frame_dt = (video_start_dt + datetime.timedelta(milliseconds=t_ms)) if video_start_dt else None
                exif_bytes = build_exif_bytes(
                    tel, w_px, h_px,
                    camera_make=args.camera_make, camera_model=args.camera_model,
                    crop_factor=args.crop_factor, alt_source=args.alt_source,
                    frame_dt=frame_dt,
                )

                filename = f"{prefix}_f{saved_idx:06d}.jpg"
                out_path = args.out / "images" / filename
                meta = {
                    "video_source": args.video.name,
                    "filename": filename,
                    "frame_idx": best_frame_idx,
                    "t_ms_video": t_ms,
                    "blur": f"{best_blur:.3f}",
                    "is_fallback": "1" if is_fallback else "",
                    "sim_to_prev": "" if sim is None else f"{sim:.5f}",
                    "dist_to_prev_m": "" if dist is None else f"{dist:.2f}",
                    "nearest_t_ms_srt": tel.t_ms_srt,
                    "telemetry_delta_ms": tel_delta,
                }
                geotag = {
                    "video_source": args.video.name,
                    "filename": filename,
                    "t_ms_video": t_ms,
                    "nearest_t_ms_srt": tel.t_ms_srt,
                    "telemetry_delta_ms": tel_delta,
                    "latitude": tel.lat,
                    "longitude": tel.lon,
                    "abs_alt": tel.abs_alt,
                    "rel_alt": tel.rel_alt,
                }
                pending_writes.append((
                    write_pool.submit(write_jpeg_with_exif, best_frame, out_path, exif_bytes, args.jpeg_quality),
                    meta, geotag,
                ))

                if tel_delta > args.max_telemetry_delta_ms:
                    print(
                        f"  WARNING: {filename} telemetry delta {tel_delta}ms "
                        f"(t_video={t_ms}, t_srt={tel.t_ms_srt})"
                    )

                saved_idx += 1
                window_start_frame = window_end_frame

                if not ret:
                    break

        else:
            # Frame-step mode (original behaviour, kept for backward compatibility)
            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                if frame_idx % args.frame_step != 0:
                    frame_idx += 1
                    continue

                # Center crop for blur scoring; full-frame resize for histogram
                crop = _center_crop(frame, args.score_scale)
                b = blur_score(crop)
                if b < args.blur_thresh:
                    frame_idx += 1
                    continue

                # Compute telemetry (needed for distance check)
                t_ms = int((frame_idx / fps) * 1000.0)
                tel = interpolate_telemetry(telemetry, t_ms, telemetry_keys)
                tel_delta = abs(tel.t_ms_srt - t_ms)

                # Primary filter: GPS distance
                dist: Optional[float] = None
                if args.min_distance_m > 0 and prev_tel is not None:
                    dist = haversine_m(prev_tel.lat, prev_tel.lon, tel.lat, tel.lon)
                    if dist < args.min_distance_m:
                        frame_idx += 1
                        continue

                # Secondary filter: HSV histogram similarity
                sf = cv2.resize(frame, (0, 0), fx=args.score_scale, fy=args.score_scale,
                                interpolation=cv2.INTER_AREA)
                h = frame_hist(sf)
                if prev_hist is not None:
                    sim = cv2.compareHist(prev_hist, h, cv2.HISTCMP_CORREL)
                    if sim > args.sim_thresh:
                        frame_idx += 1
                        continue
                else:
                    sim = None

                prev_hist = h
                prev_tel = tel

                h_px, w_px = frame.shape[:2]
                frame_dt = (video_start_dt + datetime.timedelta(milliseconds=t_ms)) if video_start_dt else None
                exif_bytes = build_exif_bytes(
                    tel, w_px, h_px,
                    camera_make=args.camera_make, camera_model=args.camera_model,
                    crop_factor=args.crop_factor, alt_source=args.alt_source,
                    frame_dt=frame_dt,
                )

                filename = f"{prefix}_f{saved_idx:06d}.jpg"
                out_path = args.out / "images" / filename
                meta = {
                    "video_source": args.video.name,
                    "filename": filename,
                    "frame_idx": frame_idx,
                    "t_ms_video": t_ms,
                    "blur": f"{b:.3f}",
                    "sim_to_prev": "" if sim is None else f"{sim:.5f}",
                    "dist_to_prev_m": "" if dist is None else f"{dist:.2f}",
                    "nearest_t_ms_srt": tel.t_ms_srt,
                    "telemetry_delta_ms": tel_delta,
                }
                geotag = {
                    "video_source": args.video.name,
                    "filename": filename,
                    "t_ms_video": t_ms,
                    "nearest_t_ms_srt": tel.t_ms_srt,
                    "telemetry_delta_ms": tel_delta,
                    "latitude": tel.lat,
                    "longitude": tel.lon,
                    "abs_alt": tel.abs_alt,
                    "rel_alt": tel.rel_alt,
                }
                pending_writes.append((
                    write_pool.submit(write_jpeg_with_exif, frame.copy(), out_path, exif_bytes, args.jpeg_quality),
                    meta, geotag,
                ))

                if tel_delta > args.max_telemetry_delta_ms:
                    print(
                        f"  WARNING: {filename} telemetry delta {tel_delta}ms "
                        f"(t_video={t_ms}, t_srt={tel.t_ms_srt})"
                    )

                saved_idx += 1
                frame_idx += 1

        cap.release()
    # write_pool.shutdown(wait=True) called here — all JPEG writes complete

    # Collect write results; skip any that failed
    for fut, meta, geotag in pending_writes:
        try:
            fut.result()
            frames_meta.append(meta)
            geotags.append(geotag)
        except Exception as e:
            print(f"  WARNING: failed to write {meta['filename']}: {e}")

    if not frames_meta:
        print(f"  No frames saved (all rejected by blur/distance/sim filters).")
        return saved_idx

    # Append to dataset-wide CSVs (create headers on first write)
    _append_csv(
        args.out / "dataset_summary.csv",
        frames_meta,
        ["video_source", "filename", "frame_idx", "t_ms_video",
         "blur", "is_fallback", "sim_to_prev", "dist_to_prev_m", "nearest_t_ms_srt", "telemetry_delta_ms"],
    )
    _append_csv(
        args.out / "dataset_geotags.csv",
        geotags,
        ["video_source", "filename", "t_ms_video", "nearest_t_ms_srt",
         "telemetry_delta_ms", "latitude", "longitude", "abs_alt", "rel_alt"],
    )

    fallback_msg = f"  ({fallback_count} fallback)" if fallback_count > 0 else ""
    gap_msg = f"  ({gap_count} uncoverable gaps)" if gap_count > 0 else ""
    print(
        f"  Saved {len(frames_meta)} frames{fallback_msg}{gap_msg}  "
        f"(global idx {args.start_idx}–{saved_idx - 1})  → {args.out}"
    )
    return saved_idx


def _append_csv(path: Path, rows: List[Dict], fieldnames: List[str]) -> None:
    """Append rows to a CSV file, writing header only if the file is new/empty.

    Uses an exclusive file lock so concurrent video subprocesses don't interleave rows
    or double-write the header.
    """
    with path.open("a", newline="", encoding="utf-8") as f:
        fcntl.flock(f, fcntl.LOCK_EX)
        try:
            # Use fstat for actual on-disk size — f.tell() returns the position
            # captured at open() time and is stale when other writers have
            # already appended since this fd was opened.
            write_header = os.fstat(f.fileno()).st_size == 0
            w = csv.DictWriter(f, fieldnames=fieldnames)
            if write_header:
                w.writeheader()
            w.writerows(rows)
            f.flush()  # ensure bytes reach disk before releasing the lock
        finally:
            fcntl.flock(f, fcntl.LOCK_UN)


if __name__ == "__main__":
    main()