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