| """SpeedNet v4 — ego-speed estimation from onboard/POV video. |
| |
| Self-contained inference module. No project dependencies: just |
| torch, opencv-python, numpy (and optionally pandas for smoothing). |
| |
| Pipeline |
| -------- |
| 1. Decode the video, resize frames to 320x180. |
| 2. Dense optical flow (Farneback) between consecutive frames, resized to |
| 128x72. Flow is normalized to "pixels per 1/15 s" so any source frame |
| rate works (multiply raw flow by fps/15). |
| 3. A low-resolution RGB context frame (128x72, sampled ~1 Hz) feeds a small |
| context branch that lets the model infer scene scale (indoor kart track |
| vs. open highway look identical in flow magnitude but differ in scale). |
| 4. Flow features + context embedding run through a GRU. Two inference |
| modes exist (sequence length acts as a scale cue for this model): |
| short windows (default; roads/trails/open environments) and |
| long_context=True (closed-course/track footage such as indoor karting). |
| |
| Typical usage |
| ------------- |
| from modeling_speednet import SpeedNet, predict_video |
| import torch |
| |
| model = SpeedNet() |
| model.load_state_dict(torch.load("speednet_v5.pt", map_location="cuda")) |
| result = predict_video(model, "my_onboard_clip.mp4", device="cuda") |
| # result["t"] -> timestamps in seconds |
| # result["speed_mps"]-> speed in m/s (multiply by 3.6 for km/h) |
| """ |
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| FLOW_SCALE = 8.0 |
| SPEED_SCALE = 40.0 |
| BASE_FPS = 15.0 |
|
|
|
|
| class SpeedNet(nn.Module): |
| """Flow CNN + RGB context branch -> GRU -> speed (v4 architecture).""" |
|
|
| def __init__(self, hidden=192, ctx_dim=64): |
| super().__init__() |
|
|
| def blk(i, o): |
| return [nn.Conv2d(i, o, 3, 2, 1), nn.BatchNorm2d(o), nn.SiLU()] |
|
|
| self.cnn = nn.Sequential(*blk(2, 32), *blk(32, 64), *blk(64, 128), |
| *blk(128, 256), nn.AdaptiveAvgPool2d(1), |
| nn.Flatten()) |
| self.ctx_cnn = nn.Sequential(*blk(3, 16), *blk(16, 32), *blk(32, 64), |
| nn.AdaptiveAvgPool2d(1), nn.Flatten(), |
| nn.Linear(64, ctx_dim), nn.SiLU()) |
| self.gru = nn.GRU(256 + ctx_dim, hidden, num_layers=2, |
| batch_first=True) |
| self.head = nn.Linear(hidden, 1) |
|
|
| def forward(self, flow, ctx, h0=None): |
| """flow: [B, T, 2, 72, 128] (normalized), ctx: [B, 3, 72, 128] |
| (RGB/255 - 0.5). Returns (speed [B, T] normalized, hidden).""" |
| B, T = flow.shape[:2] |
| f = self.cnn(flow.flatten(0, 1)).view(B, T, -1) |
| c = self.ctx_cnn(ctx)[:, None, :].expand(B, T, -1) |
| y, h = self.gru(torch.cat([f, c], -1), h0) |
| return self.head(y).squeeze(-1), h |
|
|
|
|
| def extract_flow_and_ctx(video_path, max_seconds=None): |
| """Decode video -> (flow [N,72,128,2] float32 normalized to 15 fps |
| convention, ctx [M,72,128,3] uint8 at ~1 Hz, fps).""" |
| import cv2 |
| cap = cv2.VideoCapture(video_path) |
| fps = cap.get(cv2.CAP_PROP_FPS) or BASE_FPS |
| fps_scale = fps / BASE_FPS |
| ctx_step = max(1, int(round(fps))) |
| prev, flows, ctxs = None, [], [] |
| i = 0 |
| while True: |
| ok, frame = cap.read() |
| if not ok: |
| break |
| if max_seconds and i / fps > max_seconds: |
| break |
| if frame.shape[:2] != (180, 320): |
| frame = cv2.resize(frame, (320, 180)) |
| if i % ctx_step == 0: |
| ctxs.append(cv2.cvtColor(cv2.resize(frame, (128, 72)), |
| cv2.COLOR_BGR2RGB)) |
| g = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) |
| if prev is not None: |
| fl = cv2.calcOpticalFlowFarneback(prev, g, None, |
| 0.5, 3, 15, 3, 5, 1.2, 0) |
| fl *= fps_scale |
| flows.append(cv2.resize(fl, (128, 72))) |
| prev = g |
| i += 1 |
| cap.release() |
| return (np.stack(flows).astype(np.float32), |
| np.stack(ctxs).astype(np.uint8), fps) |
|
|
|
|
| @torch.no_grad() |
| def predict_speed(model, flow, ctx, fps, device="cuda", long_context=False): |
| """Inference in one of two modes (see README "Inference modes"): |
| |
| - short (default): hidden state reset every 16 frames. Matches the |
| published GPS-domain metrics (roads, trails, open environments). |
| - long_context=True: hidden state carried across the clip. Use for |
| closed-course/track footage (e.g. indoor karting) where absolute |
| scale must be inferred from long temporal context. |
| |
| The model was trained with two window lengths, so sequence length acts |
| as a scale cue — pick the mode matching your footage. |
| """ |
| model.eval().to(device) |
| ctx_hz = max(1, int(round(fps))) |
| reset = 10**9 if long_context else 16 |
| out = np.empty(len(flow), np.float32) |
| for r0 in range(0, len(flow), reset): |
| seg = flow[r0:r0 + reset] |
| h = None |
| for s in range(0, len(seg), 240): |
| x = torch.from_numpy(seg[s:s + 240] |
| .transpose(0, 3, 1, 2))[None].to(device) \ |
| / FLOW_SCALE |
| ci = ctx[min((r0 + s) // ctx_hz, len(ctx) - 1)] \ |
| .astype(np.float32) / 255. - .5 |
| c = torch.from_numpy(ci.transpose(2, 0, 1))[None].to(device) |
| p, h = model(x, c, h) |
| out[r0 + s:r0 + s + x.shape[1]] = p[0].cpu().numpy() * SPEED_SCALE |
| return np.clip(out, 0, None) |
|
|
|
|
| predict_streaming = predict_speed |
|
|
|
|
| def smooth(speed_mps, fps, seconds=1.0): |
| """Rolling-median smoothing (recommended post-processing).""" |
| try: |
| import pandas as pd |
| k = max(3, int(round(seconds * fps)) | 1) |
| return pd.Series(speed_mps).rolling(k, center=True, min_periods=1) \ |
| .median().to_numpy() |
| except ImportError: |
| return speed_mps |
|
|
|
|
| def predict_video(model, video_path, device="cuda", max_seconds=None, |
| long_context=False): |
| """End-to-end: video file -> dict(t, speed_mps, speed_smooth_mps, fps). |
| Set long_context=True for closed-course/track footage (indoor karting).""" |
| flow, ctx, fps = extract_flow_and_ctx(video_path, max_seconds) |
| speed = predict_speed(model, flow, ctx, fps, device, |
| long_context=long_context) |
| t = (np.arange(len(speed)) + 1) / fps |
| return dict(t=t, speed_mps=speed, |
| speed_smooth_mps=smooth(speed, fps), fps=fps) |
|
|
