# 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