# https://claude.ai/public/artifacts/ab215f0c-c8d7-4974-a90b-288fe4bb43be

#!/usr/bin/env python3

"""

Select the four corner frames of a drone survey area from aerial video.

Two modes:

  telemetry -- if a DJI-style .SRT sidecar with [latitude]/[longitude] is present,

               the flight track comes straight from GPS (true north-up).

  visual -- otherwise, estimate the track by accumulating frame-to-frame

               similarity transforms (rough visual odometry on the ground plane).

In both modes the track is reduced to its minimum-area rotated rectangle; the

frame nearest each rectangle corner is emitted, ordered bottom-left then

clockwise (BL -> TL -> TR -> BR).

Usage:

    python survey_corners.py VIDEO [-o OUTDIR] [--srt FILE] [--fps 2] [--width 640]

"""

import argparse

import math

import os

import re

import sys

import cv2

import numpy as np

SRT_LAT = re.compile(r"\[latitude\s*:\s*([-\d.]+)\]")

SRT_LON = re.compile(r"\[long(?:i)?tude\s*:\s*([-\d.]+)\]")

# ---------------------------------------------------------------- telemetry

def track_from_srt(path):

    """Return (Nx2 array of local metres, Nx1 array of seconds) or None."""

    text = open(path, "r", errors="ignore").read()

    lats = [float(m) for m in SRT_LAT.findall(text)]

    lons = [float(m) for m in SRT_LON.findall(text)]

    if len(lats) < 4 or len(lats) != len(lons):

        return None

    lat0 = math.radians(np.mean(lats))

    # equirectangular projection, fine over a survey-sized area

    east = (np.array(lons) - np.mean(lons)) * 111320.0 * math.cos(lat0)

    north = (np.array(lats) - np.mean(lats)) * 110540.0

    return np.column_stack([east, north])

# ------------------------------------------------------------------ visual

def track_from_video(cap, sample_fps, width):

    """Accumulate similarity transforms into a rough ground track.

    World frame is aligned to the first frame's heading, so 'bottom-left' is

    relative to the drone's initial orientation, not to true north.

    """

    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0

    step = max(1, int(round(src_fps / sample_fps)))

    orb = cv2.ORB_create(1500)

    matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

    positions, frame_idx = [], []

    pos = np.zeros(2)

    heading = 0.0 # cumulative yaw, radians

    prev_kp = prev_des = None

    scale = None

    i = 0

    while True:

        ok = cap.grab()

        if not ok:

            break

        if i % step:

            i += 1

            continue

        ok, frame = cap.retrieve()

        if not ok:

            break

        if scale is None:

            scale = width / float(frame.shape[1])

        small = cv2.resize(frame, None, fx=scale, fy=scale)

        gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)

        gray = cv2.createCLAHE(2.0, (8, 8)).apply(gray)

        kp, des = orb.detectAndCompute(gray, None)

        if prev_des is not None and des is not None and len(des) > 10:

            matches = matcher.match(prev_des, des)

            if len(matches) >= 12:

                matches = sorted(matches, key=lambda m: m.distance)[:400]

                src = np.float32([prev_kp[m.queryIdx].pt for m in matches])

                dst = np.float32([kp[m.trainIdx].pt for m in matches])

                M, _ = cv2.estimateAffinePartial2D(

                    src, dst, method=cv2.RANSAC, ransacReprojThreshold=3.0

                )

                if M is not None:

                    # scene shift in image space; camera moves the other way

                    dx, dy = -M[0, 2], -M[1, 2]

                    dyaw = math.atan2(M[1, 0], M[0, 0])

                    c, s = math.cos(heading), math.sin(heading)

                    # rotate into world frame, flip y so +y is "up" on the map

                    pos = pos + np.array([c * dx - s * dy, -(s * dx + c * dy)])

                    heading += dyaw

        positions.append(pos.copy())

        frame_idx.append(i)

        prev_kp, prev_des = kp, des

        i += 1

    return np.array(positions), np.array(frame_idx)

# ------------------------------------------------------------------ corners

def order_clockwise_from_bottom_left(box, centre):

    """Order 4 points BL -> TL -> TR -> BR in a y-up coordinate frame."""

    ang = np.array([math.atan2(p[1] - centre[1], p[0] - centre[0]) % (2 * math.pi)

                    for p in box])

    target = 5 * math.pi / 4 # 225 deg = bottom-left

    diff = np.abs((ang - target + math.pi) % (2 * math.pi) - math.pi)

    start = int(np.argmin(diff))

    order = list(np.argsort(-ang)) # clockwise = decreasing angle

    k = order.index(start)

    return [box[j] for j in order[k:] + order[:k]]

def corner_samples(track):

    """Indices into `track` of the four survey-area corners, BL-first clockwise."""

    pts = track.astype(np.float32)

    rect = cv2.minAreaRect(pts)

    box = cv2.boxPoints(rect)

    ordered = order_clockwise_from_bottom_left(box, np.array(rect[0]))

    return [int(np.argmin(np.linalg.norm(pts - c, axis=1))) for c in ordered]

def grab_frame(cap, index):

    cap.set(cv2.CAP_PROP_POS_FRAMES, index)

    ok, frame = cap.read()

    return frame if ok else None

# --------------------------------------------------------------------- main

def main():

    ap = argparse.ArgumentParser()

    ap.add_argument("video")

    ap.add_argument("-o", "--outdir", default="corners")

    ap.add_argument("--srt", help="telemetry sidecar; defaults to VIDEO.srt if present")

    ap.add_argument("--fps", type=float, default=2.0, help="sampling rate")

    ap.add_argument("--width", type=int, default=640, help="analysis width in px")

    args = ap.parse_args()

    cap = cv2.VideoCapture(args.video)

    if not cap.isOpened():

        sys.exit(f"cannot open {args.video}")

    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    srt = args.srt or os.path.splitext(args.video)[0] + ".srt"

    track = track_from_srt(srt) if os.path.exists(srt) else None

    if track is not None:

        mode = "telemetry"

        frame_idx = np.linspace(0, max(total - 1, 0), len(track)).astype(int)

    else:

        mode = "visual"

        track, frame_idx = track_from_video(cap, args.fps, args.width)

        if len(track) < 8:

            sys.exit("not enough usable frames to estimate a track")

    picks = corner_samples(track)

    labels = ["1-bottom-left", "2-top-left", "3-top-right", "4-bottom-right"]

    os.makedirs(args.outdir, exist_ok=True)

    print(f"mode: {mode} samples: {len(track)}")

    for label, s in zip(labels, picks):

        fi = int(frame_idx[s])

        frame = grab_frame(cap, fi)

        if frame is None:

            print(f"{label}: frame {fi} unreadable")

            continue

        out = os.path.join(args.outdir, f"{label}.jpg")

        cv2.imwrite(out, frame)

        t = fi / (cap.get(cv2.CAP_PROP_FPS) or 30.0)

        print(f"{label}: frame {fi} t={t:7.2f}s xy=({track[s][0]:.1f}, {track[s][1]:.1f}) -> {out}")

    cap.release()

if __name__ == "__main__":

    main()

 

# Results:

mode: visual samples: 1100

1-bottom-left: frame 4065 t= 135.50s xy=(-252.0, -1591.2) -> corners\1-bottom-left.jpg

2-top-left: frame 2475 t= 82.50s xy=(-1809.0, 317.5) -> corners\2-top-left.jpg

3-top-right: frame 1275 t= 42.50s xy=(-371.8, 1783.0) -> corners\3-top-right.jpg

4-bottom-right: frame 4560 t= 152.00s xy=(261.0, -695.0) -> corners\4-bottom-right.jpg