#!/usr/bin/env python3
"""Estimate drone altitude gain from nadir RGB video with LingBot-Map.
The script fits a ground plane to LingBot-Map world points and measures the
camera-to-plane distance for every frame. Monocular reconstruction has an
unknown global scale, so metric output requires --initial-altitude-m (or a
separate metric registration step). This is an analytics example, not a flight-
control or safety-critical altitude source.
Run from any directory after cloning/installing LingBot-Map:
python altitude_gain_lingbot.py \
--lingbot-repo /opt/lingbot-map \
--model-path /models/lingbot-map.pt \
--video-path flight.mp4 \
--initial-altitude-m 18.4 \
--output-csv altitude.csv \
--summary-json altitude_summary.json
"""
from __future__ import annotations
import argparse
import csv
import json
import math
import sys
from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace
import cv2
import numpy as np
import torch
def effective_sample_fps(video_path: str, requested_fps: int, stride: int) -> float:
cap = cv2.VideoCapture(video_path)
source_fps = float(cap.get(cv2.CAP_PROP_FPS) or 30.0)
cap.release()
interval = max(1, round(source_fps / requested_fps))
return source_fps / interval / max(1, stride)
def fit_ground_plane_ransac(
points: np.ndarray,
optical_axis: np.ndarray,
camera_center: np.ndarray,
threshold: float,
iterations: int,
max_tilt_deg: float,
seed: int,
) -> tuple[np.ndarray, float, np.ndarray]:
"""Return oriented unit normal, offset, and inlier mask for n.x + d = 0."""
if len(points) < 100:
raise ValueError(f"Too few candidate ground points: {len(points)}")
rng = np.random.default_rng(seed)
cos_limit = math.cos(math.radians(max_tilt_deg))
optical_axis = optical_axis / np.linalg.norm(optical_axis)
best_mask = None
best_count = -1
best_median = np.inf
for _ in range(iterations):
p0, p1, p2 = points[rng.choice(len(points), 3, replace=False)]
n = np.cross(p1 - p0, p2 - p0)
norm = np.linalg.norm(n)
if norm < 1e-9:
continue
n /= norm
if abs(float(n @ optical_axis)) < cos_limit:
continue
d = -float(n @ p0)
residuals = np.abs(points @ n + d)
mask = residuals <= threshold
count = int(mask.sum())
med = float(np.median(residuals[mask])) if count else np.inf
if count > best_count or (count == best_count and med < best_median):
best_mask, best_count, best_median = mask, count, med
if best_mask is None or best_count < 100:
raise RuntimeError(
"Ground-plane RANSAC failed. Check that the video is nadir-looking, "
"lower --confidence-threshold, or increase --plane-threshold-ratio."
)
inliers = points[best_mask]
centroid = inliers.mean(axis=0)
_, _, vh = np.linalg.svd(inliers - centroid, full_matrices=False)
n = vh[-1]
n /= np.linalg.norm(n)
d = -float(n @ centroid)
# Orient the plane normal toward the camera side of the ground plane.
if float(n @ camera_center + d) < 0:
n, d = -n, -d
residuals = np.abs(points @ n + d)
refined_mask = residuals <= threshold
return n, d, refined_mask
def moving_median(values: np.ndarray, window: int) -> np.ndarray:
if window <= 1:
return values.copy()
window = window if window % 2 == 1 else window + 1
radius = window // 2
padded = np.pad(values, (radius, radius), mode="edge")
return np.array(
[np.median(padded[i : i + window]) for i in range(len(values))],
dtype=np.float64,
)
def collect_plane_points(
world_points: np.ndarray,
confidence: np.ndarray,
frame_count: int,
confidence_threshold: float,
center_fraction: float,
pixel_stride: int,
max_points: int,
seed: int,
) -> np.ndarray:
samples = []
for i in range(min(frame_count, world_points.shape[0])):
pts = np.asarray(world_points[i])
conf = np.asarray(confidence[i]).squeeze()
if pts.ndim != 3 or pts.shape[-1] != 3:
raise ValueError(f"Unexpected world_points[{i}] shape: {pts.shape}")
h, w = pts.shape[:2]
y_margin = int((1.0 - center_fraction) * h / 2.0)
x_margin = int((1.0 - center_fraction) * w / 2.0)
ys = slice(y_margin, max(y_margin + 1, h - y_margin), pixel_stride)
xs = slice(x_margin, max(x_margin + 1, w - x_margin), pixel_stride)
p = pts[ys, xs].reshape(-1, 3)
c = conf[ys, xs].reshape(-1)
valid = np.isfinite(p).all(axis=1) & np.isfinite(c) & (c >= confidence_threshold)
samples.append(p[valid])
points = np.concatenate(samples, axis=0) if samples else np.empty((0, 3))
if len(points) > max_points:
rng = np.random.default_rng(seed)
points = points[rng.choice(len(points), max_points, replace=False)]
return points.astype(np.float64, copy=False)
def build_model_args(args: argparse.Namespace) -> SimpleNamespace:
return SimpleNamespace(
mode=args.mode,
model_path=args.model_path,
image_size=args.image_size,
patch_size=14,
enable_3d_rope=True,
max_frame_num=1024,
kv_cache_sliding_window=args.kv_cache_sliding_window,
num_scale_frames=args.num_scale_frames,
use_sdpa=args.use_sdpa,
camera_num_iterations=args.camera_num_iterations,
)
def run_lingbot(args: argparse.Namespace):
repo = Path(args.lingbot_repo).resolve()
if not (repo / "demo.py").is_file():
raise FileNotFoundError(f"LingBot-Map demo.py not found under {repo}")
sys.path.insert(0, str(repo))
from demo import load_images, load_model, postprocess
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device.type != "cuda":
raise RuntimeError("This example expects a CUDA GPU for LingBot-Map inference.")
images, paths, _ = load_images(
video_path=args.video_path,
fps=args.fps,
stride=args.stride,
first_k=args.first_k,
image_size=args.image_size,
patch_size=14,
)
if images.shape[0] < max(3, args.num_scale_frames):
raise ValueError("The sampled video contains too few frames for reconstruction.")
model_args = build_model_args(args)
model = load_model(model_args, device)
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
if getattr(model, "aggregator", None) is not None:
model.aggregator = model.aggregator.to(dtype=dtype)
images = images.to(device)
frame_count = int(images.shape[0])
keyframe_interval = args.keyframe_interval
if keyframe_interval is None:
keyframe_interval = math.ceil(frame_count / 320) if args.mode == "streaming" and frame_count > 320 else 1
output_device = torch.device("cpu") if args.offload_to_cpu else None
with torch.no_grad(), torch.autocast(device_type="cuda", dtype=dtype):
if args.mode == "streaming":
predictions = model.inference_streaming(
images,
num_scale_frames=args.num_scale_frames,
keyframe_interval=keyframe_interval,
output_device=output_device,
)
else:
predictions = model.inference_windowed(
images,
window_size=args.window_size,
overlap_size=args.overlap_size,
overlap_keyframes=args.overlap_keyframes,
num_scale_frames=args.num_scale_frames,
keyframe_interval=keyframe_interval,
output_device=output_device,
)
if args.offload_to_cpu:
images_for_post = predictions.get("images")
if images_for_post is None:
raise KeyError("LingBot-Map did not return images while CPU offload was enabled.")
else:
images_for_post = images
predictions, _ = postprocess(predictions, images_for_post)
return predictions, paths, keyframe_interval
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(
description="Estimate relative or metric drone altitude gain using LingBot-Map."
)
p.add_argument("--lingbot-repo", required=True, help="Path to cloned Robbyant/lingbot-map")
p.add_argument("--model-path", required=True, help="Path to lingbot-map.pt")
p.add_argument("--video-path", required=True, help="Nadir/top-down drone video")
p.add_argument("--output-csv", default="altitude_gain.csv")
p.add_argument("--summary-json", default="altitude_summary.json")
p.add_argument(
"--initial-altitude-m",
type=float,
default=None,
help="Known camera height above the fitted ground plane at the baseline frames; enables metres",
)
p.add_argument("--fps", type=int, default=10, help="Requested video sampling rate")
p.add_argument("--stride", type=int, default=1, help="Additional stride after FPS sampling")
p.add_argument("--first-k", type=int, default=None)
p.add_argument("--image-size", type=int, default=518)
p.add_argument("--mode", choices=("streaming", "windowed"), default="streaming")
p.add_argument("--num-scale-frames", type=int, default=8)
p.add_argument("--keyframe-interval", type=int, default=None)
p.add_argument("--kv-cache-sliding-window", type=int, default=64)
p.add_argument("--window-size", type=int, default=128)
p.add_argument("--overlap-size", type=int, default=16)
p.add_argument("--overlap-keyframes", type=int, default=16)
p.add_argument("--camera-num-iterations", type=int, default=4)
p.add_argument("--use-sdpa", action="store_true", help="Use PyTorch SDPA instead of FlashInfer")
p.add_argument(
"--offload-to-cpu",
action=argparse.BooleanOptionalAction,
default=True,
help="Offload predictions to CPU during inference",
)
p.add_argument("--plane-frames", type=int, default=8)
p.add_argument("--baseline-frames", type=int, default=3)
p.add_argument("--confidence-threshold", type=float, default=1.5)
p.add_argument("--center-fraction", type=float, default=0.8)
p.add_argument("--pixel-stride", type=int, default=4)
p.add_argument("--max-plane-points", type=int, default=30000)
p.add_argument(
"--plane-threshold-ratio",
type=float,
default=0.01,
help="RANSAC distance tolerance as a fraction of initial camera-to-point range",
)
p.add_argument("--ransac-iterations", type=int, default=750)
p.add_argument("--max-plane-tilt-deg", type=float, default=45.0)
p.add_argument("--smooth-window", type=int, default=5)
p.add_argument("--seed", type=int, default=7)
return p.parse_args()
def main() -> None:
args = parse_args()
predictions, paths, keyframe_interval = run_lingbot(args)
required = {"extrinsic", "world_points", "world_points_conf"}
missing = required.difference(predictions)
if missing:
raise KeyError(f"LingBot-Map output is missing: {sorted(missing)}")
extrinsic = np.asarray(predictions["extrinsic"])
world_points = np.asarray(predictions["world_points"])
confidence = np.asarray(predictions["world_points_conf"])
if extrinsic.ndim == 4 and extrinsic.shape[0] == 1:
extrinsic = extrinsic[0]
if world_points.ndim == 5 and world_points.shape[0] == 1:
world_points = world_points[0]
if confidence.ndim == 4 and confidence.shape[0] == 1:
confidence = confidence[0]
camera_centers = extrinsic[:, :3, 3].astype(np.float64)
camera_forward = extrinsic[:, :3, 2].astype(np.float64)
mean_forward = camera_forward[: min(args.plane_frames, len(camera_forward))].mean(axis=0)
plane_points = collect_plane_points(
world_points,
confidence,
args.plane_frames,
args.confidence_threshold,
args.center_fraction,
args.pixel_stride,
args.max_plane_points,
args.seed,
)
initial_range = np.median(np.linalg.norm(plane_points - camera_centers[0], axis=1))
threshold = max(1e-6, args.plane_threshold_ratio * float(initial_range))
normal, offset, inliers = fit_ground_plane_ransac(
plane_points,
mean_forward,
camera_centers[0],
threshold,
args.ransac_iterations,
args.max_plane_tilt_deg,
args.seed,
)
altitude_model = camera_centers @ normal + offset
altitude_model = moving_median(altitude_model, args.smooth_window)
baseline_n = min(max(1, args.baseline_frames), len(altitude_model))
baseline_model = float(np.median(altitude_model[:baseline_n]))
if baseline_model <= 0:
raise RuntimeError("Estimated baseline camera-to-ground distance is non-positive.")
if args.initial_altitude_m is not None:
scale = args.initial_altitude_m / baseline_model
altitude = altitude_model * scale
baseline = float(np.median(altitude[:baseline_n]))
unit = "m"
else:
scale = 1.0
altitude = altitude_model
baseline = baseline_model
unit = "lingbot_scale_units"
gain = altitude - baseline
sample_fps = effective_sample_fps(args.video_path, args.fps, args.stride)
timestamps = np.arange(len(altitude), dtype=np.float64) / sample_fps
output_csv = Path(args.output_csv)
output_csv.parent.mkdir(parents=True, exist_ok=True)
with output_csv.open("w", newline="", encoding="utf-8") as f:
writer = csv.writer(f)
writer.writerow(
["frame_index", "timestamp_s", "source_frame", f"altitude_{unit}", f"gain_{unit}"]
)
for i, (ts, alt, delta) in enumerate(zip(timestamps, altitude, gain)):
source = paths[i] if i < len(paths) else ""
writer.writerow([i, f"{ts:.6f}", source, f"{alt:.9f}", f"{delta:.9f}"])
plane_residuals = np.abs(plane_points @ normal + offset)
summary = {
"video": str(Path(args.video_path).resolve()),
"frames_processed": int(len(altitude)),
"effective_sample_fps": sample_fps,
"mode": args.mode,
"keyframe_interval": keyframe_interval,
"unit": unit,
"metric_scale": scale if args.initial_altitude_m is not None else None,
"baseline_altitude": baseline,
"final_altitude": float(altitude[-1]),
"final_gain": float(gain[-1]),
"maximum_gain": float(np.max(gain)),
"minimum_gain": float(np.min(gain)),
"ground_plane_normal_lingbot_world": normal.tolist(),
"ground_plane_offset_lingbot_world": offset,
"plane_candidate_points": int(len(plane_points)),
"plane_inlier_ratio": float(np.mean(inliers)),
"plane_median_absolute_residual_model_units": float(np.median(plane_residuals[inliers])),
"notes": [
"Metric altitude is scaled from the supplied initial camera height above the fitted plane.",
"The estimate assumes the visible reference surface is approximately static and planar.",
"Do not use this output for navigation, obstacle avoidance, or regulatory altitude reporting.",
],
}
summary_path = Path(args.summary_json)
summary_path.parent.mkdir(parents=True, exist_ok=True)
summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
print(json.dumps(summary, indent=2))
print(f"Wrote {output_csv}")
print(f"Wrote {summary_path}")
if __name__ == "__main__":
main()
No comments:
Post a Comment