| """KartNet v3 — track-specialized telemetry from onboard karting video. |
| |
| Estimates, per frame at 10 Hz, from video alone (no GPS — it does not |
| exist indoors): |
| |
| speed_mps vehicle speed (m/s) |
| u position around the lap in [0, 1) as (sin, cos) — lap |
| times fall out of the wrap points |
| yaw camera yaw rate (px of horizontal image shift per 0.1 s) |
| |
| If the recording carries an IMU (DJI/GoPro action cameras do), pass it in: |
| a small optional branch measurably improves speed (per-lap mean MAE |
| 1.9 -> 1.1 km/h on the held-out validation race). Without IMU the model |
| runs on video alone — one checkpoint serves both cases. |
| |
| Trained on 29 hours of onboard footage from Racehall Aarhus, auto-labelled |
| by corpus self-localization (no sensor ground truth; validated against the |
| venue's official SMS-Timing lap times to 0.2 s median). The model is |
| SPECIALIZED to this venue and layout by design — see the model card. |
| |
| Self-contained: torch, opencv-python, numpy (pandas optional, for IMU). |
| |
| Typical usage |
| ------------- |
| from modeling_kartnet import KartNet, extract_features, predict |
| import torch |
| |
| model = KartNet() |
| model.load_state_dict(torch.load("kartnet_v3.pt", map_location="cuda")) |
| feats = extract_features("onboard.mp4") # ~2x realtime on GPU |
| out = predict(model, feats, device="cuda") |
| # out["t"], out["speed_mps"], out["u"], out["lap_times"] |
| """ |
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| FLOW_SCALE = 6.0 |
| SPEED_SCALE = 30.0 |
| YAW_SCALE = 8.0 |
| HZ = 10.0 |
| SUB = 3 |
| IMU_DIM = 5 |
|
|
|
|
| def _blk(i, o, s=2): |
| return [nn.Conv2d(i, o, 3, s, 1), nn.BatchNorm2d(o), nn.SiLU()] |
|
|
|
|
| class KartNet(nn.Module): |
| """Flow CNN + appearance CNN + optional IMU MLP -> GRU -> 3 heads. |
| |
| A presence flag accompanies the IMU channels so the GRU can tell |
| "no sensors" apart from "sensors reading zero g"; the branch was |
| dropped randomly during training, so video-only inference is a |
| first-class mode, not a degradation. |
| """ |
|
|
| def __init__(self, hidden=256, app_dim=128, imu_dim=32): |
| super().__init__() |
| self.flow_cnn = nn.Sequential(*_blk(2, 32), *_blk(32, 64), |
| *_blk(64, 128), *_blk(128, 256), |
| nn.AdaptiveAvgPool2d(1), nn.Flatten()) |
| self.app_cnn = nn.Sequential(*_blk(1, 24), *_blk(24, 48), |
| *_blk(48, 96), *_blk(96, 192), |
| nn.AdaptiveAvgPool2d((2, 3)), |
| nn.Flatten(), |
| nn.Linear(192 * 6, app_dim), nn.SiLU()) |
| self.imu_mlp = nn.Sequential(nn.Linear(IMU_DIM + 1, 32), nn.SiLU(), |
| nn.Linear(32, imu_dim), nn.SiLU()) |
| self.gru = nn.GRU(256 + app_dim + imu_dim, hidden, num_layers=2, |
| batch_first=True) |
| self.head_v = nn.Linear(hidden, 1) |
| self.head_u = nn.Linear(hidden, 2) |
| self.head_w = nn.Linear(hidden, 1) |
|
|
| def forward(self, flow, gray, h0=None, app_mask=None, imu=None, |
| imu_mask=None): |
| B, T = flow.shape[:2] |
| f = self.flow_cnn(flow.flatten(0, 1).contiguous()).view(B, T, -1) |
| a = self.app_cnn(gray.flatten(0, 1).contiguous()).view(B, T, -1) |
| if app_mask is not None: |
| a = a * app_mask |
| if imu is None: |
| imu = flow.new_zeros(B, T, IMU_DIM) |
| flag = flow.new_zeros(B, T, 1) |
| else: |
| flag = imu.new_ones(B, T, 1) |
| if imu_mask is not None: |
| imu = imu * imu_mask |
| flag = flag * imu_mask |
| i = self.imu_mlp(torch.cat([imu, flag], -1)) |
| y, h = self.gru(torch.cat([f, a, i], -1), h0) |
| u = self.head_u(y) |
| return dict(v=self.head_v(y).squeeze(-1), |
| u=u / (u.norm(dim=-1, keepdim=True) + 1e-6), |
| w=self.head_w(y).squeeze(-1)), h |
|
|
|
|
| def extract_features(video_path, max_seconds=None): |
| """Video -> dict(flow [N,2,72,128] f32, gray [N,72,128] u8, fps=10). |
| |
| Decodes at 30 fps internally (Farneback flow between kept 10 Hz frames, |
| matching training).""" |
| import cv2, subprocess |
| DW, DH, FW, FH = 320, 180, 128, 72 |
| cmd = ['ffmpeg', '-v', 'error', '-i', video_path, |
| '-vf', f'fps={HZ*SUB},scale={DW}:{DH}:flags=area,format=gray', |
| '-f', 'rawvideo', '-'] |
| p = subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=DW * DH * 64) |
| prev, flows, grays = None, [], [] |
| i = 0 |
| while True: |
| b = p.stdout.read(DW * DH) |
| if len(b) < DW * DH: |
| break |
| if max_seconds and i / (HZ * SUB) > max_seconds: |
| break |
| if i % SUB == 0: |
| g = np.frombuffer(b, np.uint8).reshape(DH, DW) |
| grays.append(cv2.resize(g, (FW, FH), |
| interpolation=cv2.INTER_AREA)) |
| if prev is None: |
| flows.append(np.zeros((FH, FW, 2), np.float32)) |
| else: |
| fl = cv2.calcOpticalFlowFarneback(prev, g, None, |
| 0.5, 3, 21, 3, 5, 1.2, 0) |
| flows.append(cv2.resize(fl, (FW, FH), |
| interpolation=cv2.INTER_AREA)) |
| prev = g |
| i += 1 |
| p.wait() |
| return dict(flow=np.stack(flows).transpose(0, 3, 1, 2) |
| .astype(np.float32), |
| gray=np.stack(grays), hz=HZ) |
|
|
|
|
| def imu_features(n_frames, t_imu, accel_xyz_g, yaw_rate_dps, |
| gforce_lon=None, gforce_lat=None): |
| """Optional IMU channels on the 10 Hz frame grid. |
| |
| t_imu must be seconds on the VIDEO's clock (embedded telemetry from the |
| same file already is). accel_xyz_g: [M,3]. The vibration channel needs |
| the full-rate stream (>=50 Hz) — tyre/engine vibration rises with speed |
| and cannot be recovered from downsampled data.""" |
| tg = np.arange(n_frames) / HZ |
| ax, ay, az = (np.asarray(accel_xyz_g)[:, k] for k in range(3)) |
| mag = np.sqrt(ax ** 2 + ay ** 2 + az ** 2) |
| trend = np.convolve(mag, np.ones(25) / 25, 'same') |
| hf = mag - trend |
| bins = np.clip((np.asarray(t_imu) * HZ).astype(int), 0, n_frames - 1) |
| acc = np.zeros(n_frames); cnt = np.zeros(n_frames) |
| np.add.at(acc, bins, hf ** 2) |
| np.add.at(cnt, bins, 1.0) |
| vib = np.sqrt(acc / np.maximum(cnt, 1)) |
| m = cnt > 0 |
| vib = np.interp(tg, tg[m], vib[m]) |
|
|
| def rs(v): |
| v = np.asarray(v, np.float64) |
| ok = np.isfinite(v) |
| return np.interp(tg, np.asarray(t_imu)[ok], v[ok]) |
|
|
| gl = rs(gforce_lon) if gforce_lon is not None else np.zeros(n_frames) |
| gt = rs(gforce_lat) if gforce_lat is not None else np.zeros(n_frames) |
| out = np.stack([gl, gt, rs(yaw_rate_dps) / 60.0, |
| rs(mag) - 1.0, vib], 1).astype(np.float32) |
| return np.nan_to_num(out, nan=0.0) |
|
|
|
|
| @torch.no_grad() |
| def predict(model, feats, imu=None, device="cuda", win=64): |
| """Windowed inference (the model is length-invariant by training, so |
| window size is a throughput knob, not an accuracy mode).""" |
| model.eval().to(device) |
| flow, gray = feats['flow'], feats['gray'] |
| n = (min(len(flow), len(gray)) // win) * win |
| vs, us, ws = [], [], [] |
| B = 16 |
| for i in range(0, n // win, B): |
| j = min(i + B, n // win) |
| sl = slice(i * win, j * win) |
| f = torch.from_numpy(flow[sl]).to(device) \ |
| .view(j - i, win, 2, 72, 128) / FLOW_SCALE |
| g = (torch.from_numpy(gray[sl].astype(np.float32)).to(device) |
| .view(j - i, win, 1, 72, 128) / 255. - .5) |
| kw = {} |
| if imu is not None: |
| kw = dict(imu=torch.from_numpy(imu[sl]).to(device) |
| .view(j - i, win, -1), |
| imu_mask=torch.ones(j - i, 1, 1, device=device)) |
| out, _ = model(f, g, **kw) |
| vs.append(out['v'].reshape(-1).cpu().numpy() * SPEED_SCALE) |
| us.append(out['u'].reshape(-1, 2).cpu().numpy()) |
| ws.append(out['w'].reshape(-1).cpu().numpy() * YAW_SCALE) |
| v = np.concatenate(vs); u = np.concatenate(us) |
| t = np.arange(len(v)) / HZ |
| |
| ang = np.arctan2(u[:, 0], u[:, 1]) |
| k = 11 |
| c = np.convolve(np.cos(ang), np.ones(k) / k, 'same') |
| s = np.convolve(np.sin(ang), np.ones(k) / k, 'same') |
| ph = np.remainder(np.arctan2(s, c) / (2 * np.pi), 1.0) |
| armed, cross, last = False, [], -1e9 |
| for i, p in enumerate(ph): |
| if p > 0.7: |
| armed = True |
| elif armed and p < 0.3 and (i - last) > 45 * HZ: |
| cross.append(i); last = i; armed = False |
| laps = np.diff(np.array(cross) / HZ) if len(cross) > 1 else np.array([]) |
| return dict(t=t, speed_mps=v, u=ph, yaw=np.concatenate(ws), |
| lap_boundaries_s=np.array(cross) / HZ, lap_times_s=laps) |
|
|
