auto-iter 20260513-2231: GX019817 RoPE skip, 4 PLY done ready for stage06
This commit is contained in:
Poulpe
166
scripts/dvl_optical_imu.py
Normal file
166
scripts/dvl_optical_imu.py
Normal file
@@ -0,0 +1,166 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Optical DVL + IMU heading correction.
|
||||
Same as dvl_optical.py but rotates cam-frame flow to world frame using compass_hdg.
|
||||
|
||||
Usage:
|
||||
python3 dvl_optical_imu.py --frames-dir <dir> --bag-dir <auv_bags_dir> \
|
||||
--altitude 1.5 --fps 1.0 --start-iso ... --label ... \
|
||||
--out csv --plot png [--ref-csv ...]
|
||||
"""
|
||||
import argparse, csv, math, sys
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
import numpy as np
|
||||
import cv2
|
||||
from rosbags.highlevel import AnyReader
|
||||
|
||||
def load_heading(bag_dir, t_start, t_end):
|
||||
bags = sorted(Path(bag_dir).glob('*.mcap'))
|
||||
# filter empty
|
||||
bags = [b for b in bags if b.stat().st_size > 1000]
|
||||
headings = [] # list of (ts_s, heading_deg)
|
||||
for b in bags:
|
||||
try:
|
||||
with AnyReader([b]) as r:
|
||||
for conn, ts_ns, raw in r.messages(connections=[c for c in r.connections if c.topic == '/mavros/global_position/compass_hdg']):
|
||||
t = ts_ns / 1e9
|
||||
if t_start - 60 <= t <= t_end + 60:
|
||||
m = r.deserialize(raw, conn.msgtype)
|
||||
headings.append((t, m.data))
|
||||
except Exception as e:
|
||||
print(f'[warn] {b.name}: {e}', flush=True)
|
||||
headings.sort()
|
||||
return headings
|
||||
|
||||
def nearest_hdg(headings, t_target):
|
||||
if not headings: return None
|
||||
ts = [h[0] for h in headings]
|
||||
idx = np.searchsorted(ts, t_target)
|
||||
if idx == 0: return headings[0][1]
|
||||
if idx >= len(headings): return headings[-1][1]
|
||||
if abs(headings[idx][0] - t_target) < abs(headings[idx-1][0] - t_target):
|
||||
return headings[idx][1]
|
||||
return headings[idx-1][1]
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument('--frames-dir', required=True)
|
||||
ap.add_argument('--bag-dir', required=True)
|
||||
ap.add_argument('--altitude', type=float, default=1.5)
|
||||
ap.add_argument('--fov-deg', type=float, default=122.0)
|
||||
ap.add_argument('--fps', type=float, default=1.0)
|
||||
ap.add_argument('--start-iso', default='2026-05-05T00:00:00')
|
||||
ap.add_argument('--label', default='segment')
|
||||
ap.add_argument('--out', required=True)
|
||||
ap.add_argument('--plot', default=None)
|
||||
ap.add_argument('--ref-csv', default=None)
|
||||
args = ap.parse_args()
|
||||
|
||||
frames = sorted(Path(args.frames_dir).glob('frame_*.jpg'))
|
||||
print(f'[dvl] {len(frames)} frames', flush=True)
|
||||
|
||||
t0 = datetime.fromisoformat(args.start_iso).timestamp()
|
||||
t_end = t0 + len(frames) / args.fps
|
||||
|
||||
# Load heading
|
||||
print(f'[hdg] loading from {args.bag_dir}', flush=True)
|
||||
headings = load_heading(args.bag_dir, t0, t_end)
|
||||
print(f'[hdg] {len(headings)} samples loaded, t range: {headings[0][0]:.0f}-{headings[-1][0]:.0f}', flush=True)
|
||||
|
||||
W, H = 518, 294
|
||||
f = (W/2) / math.tan(math.radians(args.fov_deg/2))
|
||||
px_to_m = args.altitude / f
|
||||
print(f'[dvl] px_to_m={px_to_m:.5f}', flush=True)
|
||||
|
||||
rows = []
|
||||
rows.append({'frame_idx': 0, 'ts_s': t0, 'heading_deg': nearest_hdg(headings, t0) or 0, 'flow_x_px': 0, 'flow_y_px': 0,
|
||||
'speed_mps': 0, 'east_m': 0, 'north_m': 0})
|
||||
|
||||
prev = cv2.imread(str(frames[0]), cv2.IMREAD_GRAYSCALE)
|
||||
east_cum, north_cum = 0.0, 0.0
|
||||
|
||||
for i in range(1, len(frames)):
|
||||
curr = cv2.imread(str(frames[i]), cv2.IMREAD_GRAYSCALE)
|
||||
if curr is None: continue
|
||||
t_frame = t0 + i / args.fps
|
||||
hdg = nearest_hdg(headings, t_frame) or 0
|
||||
|
||||
flow = cv2.calcOpticalFlowFarneback(prev, curr, None, 0.5, 3, 21, 3, 5, 1.2, 0)
|
||||
fx_cam = np.median(flow[..., 0])
|
||||
fy_cam = np.median(flow[..., 1])
|
||||
|
||||
# convert px → m in CAM frame (cam right = +X_cam, cam down = +Y_cam image coord)
|
||||
dx_cam = -fx_cam * px_to_m # AUV moves opposite to flow
|
||||
dy_cam = -fy_cam * px_to_m
|
||||
|
||||
# Apply heading rotation: cam +X_cam = body forward? assume cam frame Y axis = AUV forward
|
||||
# The downward camera: cam +Y_image = body forward typically (or -Y if mounted otherwise)
|
||||
# heading = degrees clockwise from North in body frame
|
||||
# World rotation: rotate body (dy_cam = forward, dx_cam = right) by heading angle from north
|
||||
hdg_rad = math.radians(hdg)
|
||||
# body forward (north when hdg=0) component:
|
||||
# body_forward_m = dy_cam (assuming cam Y_image = forward)
|
||||
# body_right_m = dx_cam
|
||||
body_forward = dy_cam # may need sign flip depending on mounting; we'll see
|
||||
body_right = dx_cam
|
||||
# world East = forward*sin(hdg) + right*cos(hdg)
|
||||
# world North = forward*cos(hdg) - right*sin(hdg)
|
||||
de = body_forward * math.sin(hdg_rad) + body_right * math.cos(hdg_rad)
|
||||
dn = body_forward * math.cos(hdg_rad) - body_right * math.sin(hdg_rad)
|
||||
|
||||
east_cum += de
|
||||
north_cum += dn
|
||||
speed_mps = math.sqrt(de**2 + dn**2) * args.fps
|
||||
|
||||
rows.append({'frame_idx': i, 'ts_s': t_frame, 'heading_deg': hdg, 'flow_x_px': float(fx_cam), 'flow_y_px': float(fy_cam),
|
||||
'speed_mps': speed_mps, 'east_m': east_cum, 'north_m': north_cum})
|
||||
prev = curr
|
||||
if i % 100 == 0:
|
||||
print(f'[dvl] {i}/{len(frames)} hdg={hdg:.1f}° flow=({fx_cam:.1f},{fy_cam:.1f}) pos=({east_cum:.2f},{north_cum:.2f})', flush=True)
|
||||
|
||||
print(f'[dvl] done. Final ENU: ({east_cum:.2f}, {north_cum:.2f}) m', flush=True)
|
||||
|
||||
with open(args.out, 'w', newline='') as ff:
|
||||
w = csv.DictWriter(ff, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader(); w.writerows(rows)
|
||||
print(f'[out] {args.out}', flush=True)
|
||||
|
||||
if args.plot:
|
||||
import matplotlib
|
||||
matplotlib.use('Agg')
|
||||
import matplotlib.pyplot as plt
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 12))
|
||||
ax_xy, ax_hdg, ax_speed, ax_cmp = axes[0,0], axes[0,1], axes[1,0], axes[1,1]
|
||||
e = [r['east_m'] for r in rows]
|
||||
n = [r['north_m'] for r in rows]
|
||||
ax_xy.plot(e, n, '-b', linewidth=1.2)
|
||||
ax_xy.plot(e[0], n[0], 'go', markersize=10, label='start')
|
||||
ax_xy.plot(e[-1], n[-1], 'r^', markersize=10, label='end')
|
||||
ax_xy.set_xlabel('East (m)'); ax_xy.set_ylabel('North (m)'); ax_xy.set_title('DVL + IMU heading trajectory')
|
||||
ax_xy.set_aspect('equal'); ax_xy.legend(); ax_xy.grid(True, alpha=0.3)
|
||||
|
||||
hdgs = [r['heading_deg'] for r in rows]
|
||||
ax_hdg.plot(range(len(rows)), hdgs, '-c'); ax_hdg.set_xlabel('Frame'); ax_hdg.set_ylabel('Heading (deg)'); ax_hdg.set_title('Compass heading from MCAP'); ax_hdg.grid(True, alpha=0.3)
|
||||
|
||||
speeds = [r['speed_mps'] for r in rows]
|
||||
ax_speed.plot(range(len(rows)), speeds, '-r'); ax_speed.set_xlabel('Frame'); ax_speed.set_ylabel('Speed m/s'); ax_speed.set_title('Speed over time'); ax_speed.grid(True, alpha=0.3)
|
||||
|
||||
if args.ref_csv:
|
||||
try:
|
||||
with open(args.ref_csv) as fff:
|
||||
refrows = [r for r in csv.DictReader(fff) if r.get('segment','')==args.label or r.get('label','')==args.label]
|
||||
rx = [float(r['x']) for r in refrows]
|
||||
ry = [float(r['y']) for r in refrows]
|
||||
ax_cmp.plot(e, n, '-b', linewidth=1.2, label='DVL+IMU', alpha=0.7)
|
||||
ax_cmp.plot(rx, ry, '-r', linewidth=1.2, label='lingbot', alpha=0.7)
|
||||
ax_cmp.set_xlabel('X/East'); ax_cmp.set_ylabel('Y/North'); ax_cmp.set_title('Comparison'); ax_cmp.set_aspect('equal'); ax_cmp.legend(); ax_cmp.grid(True, alpha=0.3)
|
||||
except Exception as e: print(f'[ref] {e}')
|
||||
else:
|
||||
ax_cmp.set_title('(no reference)')
|
||||
|
||||
plt.suptitle(f'DVL+IMU heading — {args.label}')
|
||||
plt.tight_layout()
|
||||
plt.savefig(args.plot, dpi=130, bbox_inches='tight')
|
||||
print(f'[plot] {args.plot}', flush=True)
|
||||
|
||||
if __name__ == '__main__': main()
|
||||
Reference in New Issue
Block a user