| """SpeedNet: optical-flow -> speed regression for onboard/POV footage. |
| |
| Input per timestep: Farneback optical flow between consecutive frames of |
| 320x180 @ 15 fps video, resized to 128x72, 2 channels (px/frame at 320x180 |
| scale), divided by FLOW_SCALE. Output: speed in m/s after * SPEED_SCALE. |
| |
| IMPORTANT — inference protocol: |
| Run the GRU in STREAMING mode (carry hidden state across chunks, see |
| `predict_streaming`). Resetting the hidden state every few frames collapses |
| indoor-domain predictions to ~half scale; the domain-scale calibration lives |
| in long GRU context. |
| """ |
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| FLOW_SCALE = 8.0 |
| SPEED_SCALE = 40.0 |
|
|
|
|
| class SpeedNet(nn.Module): |
| def __init__(self, hidden=192): |
| 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.gru = nn.GRU(256, hidden, num_layers=2, batch_first=True) |
| self.head = nn.Linear(hidden, 1) |
|
|
| def forward(self, x): |
| B, T = x.shape[:2] |
| f = self.cnn(x.flatten(0, 1)).view(B, T, -1) |
| h, _ = self.gru(f) |
| return self.head(h).squeeze(-1) |
|
|
|
|
| @torch.no_grad() |
| def predict_streaming(model, flow, device='cuda', chunk=240): |
| """flow: [N, 72, 128, 2] float (px/frame). Returns speed m/s per frame.""" |
| model.eval() |
| h = None |
| out = [] |
| for s in range(0, len(flow), chunk): |
| x = torch.from_numpy(flow[s:s + chunk].astype(np.float32) |
| .transpose(0, 3, 1, 2)).to(device) / FLOW_SCALE |
| f = model.cnn(x)[None] |
| y, h = model.gru(f, h) |
| out.append(model.head(y)[0, :, 0].cpu().numpy() * SPEED_SCALE) |
| return np.concatenate(out) |
|
|
|
|
| def compute_flow(video_path): |
| """Farneback flow matching the training distribution. |
| Video should be (or be resized to) 320x180 @ 15 fps.""" |
| import cv2 |
| cap = cv2.VideoCapture(video_path) |
| prev, flows = None, [] |
| while True: |
| ok, frame = cap.read() |
| if not ok: |
| break |
| if frame.shape[:2] != (180, 320): |
| frame = cv2.resize(frame, (320, 180)) |
| 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) |
| flows.append(cv2.resize(fl, (128, 72))) |
| prev = g |
| cap.release() |
| return np.stack(flows) |
|
|
