| |
| """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") |
|
|
