File size: 1,408 Bytes
2162853 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | #!/usr/bin/env python3
"""SpeedNet v4 — minimal example.
Estimate speed from any onboard/POV video and save a CSV + plot:
pip install torch opencv-python-headless numpy pandas matplotlib
python example.py my_clip.mp4
Outputs: my_clip_speed.csv (t, speed_mps, speed_kmh) and my_clip_speed.png.
"""
import sys
import numpy as np
import torch
from modeling_speednet import SpeedNet, predict_video
video = sys.argv[1]
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SpeedNet()
model.load_state_dict(torch.load("speednet_v4.pt", map_location=device))
r = predict_video(model, video, device=device)
import pandas as pd
stem = video.rsplit(".", 1)[0]
pd.DataFrame({"t": r["t"], "speed_mps": r["speed_smooth_mps"],
"speed_kmh": r["speed_smooth_mps"] * 3.6}) \
.to_csv(f"{stem}_speed.csv", index=False)
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(12, 4))
ax.plot(r["t"], r["speed_mps"] * 3.6, lw=0.6, alpha=0.4, label="raw")
ax.plot(r["t"], r["speed_smooth_mps"] * 3.6, lw=1.5, label="smoothed (1 s)")
ax.set_xlabel("time (s)")
ax.set_ylabel("speed (km/h)")
ax.legend()
ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig(f"{stem}_speed.png", dpi=110)
print(f"mean {r['speed_smooth_mps'].mean()*3.6:.1f} km/h, "
f"max {r['speed_smooth_mps'].max()*3.6:.1f} km/h -> "
f"{stem}_speed.csv / .png")
|