"""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 # px per 0.1 s SPEED_SCALE = 30.0 # m/s YAW_SCALE = 8.0 # px per 0.1 s HZ = 10.0 # feature grid SUB = 3 # decode at 30 fps, keep every 3rd frame IMU_DIM = 5 # g_lon, g_lat, yaw/60, |a|-1, vibration RMS 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 # lap times from smoothed phase crossings 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)