| #!/usr/bin/env python3 | |
| """Minimal example: onboard karting video -> speed + lap times. | |
| python example.py my_onboard_clip.mp4 | |
| """ | |
| import sys | |
| import numpy as np | |
| import torch | |
| from modeling_kartnet import KartNet, extract_features, predict | |
| video = sys.argv[1] | |
| dev = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| model = KartNet() | |
| model.load_state_dict(torch.load('kartnet_v3.pt', map_location=dev)) | |
| print('extracting features ...') | |
| feats = extract_features(video) | |
| out = predict(model, feats, device=dev) | |
| print(f"frames: {len(out['t'])} ({out['t'][-1]:.0f}s)") | |
| print(f"speed: median {np.median(out['speed_mps'])*3.6:.1f} km/h, " | |
| f"p95 {np.percentile(out['speed_mps'],95)*3.6:.1f} km/h") | |
| if len(out['lap_times_s']): | |
| print('lap times:', ' '.join(f"{x:.2f}" for x in out['lap_times_s'])) | |
| np.savez(video.rsplit('.', 1)[0] + '_kartnet.npz', **out) | |
| print('saved', video.rsplit('.', 1)[0] + '_kartnet.npz') | |