| """Flow IDM: predict W/A/S/D/Shift and yaw/pitch from a video. |
| |
| Usage: python inference.py VIDEO [--weights model.safetensors] [--config config.json] |
| |
| Pipeline: decode -> 128x128 RGB -> RAFT-small optical flow on consecutive sampled |
| frames -> FlowTransformer. Each window of (num_frames - 1) * frame_gap + 1 source |
| frames gives one prediction. Output is JSON on stdout. |
| Needs: torch, torchvision, av, opencv-python, safetensors, numpy. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import av |
| import cv2 |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from safetensors.torch import load_file |
| from torchvision.models.optical_flow import Raft_Small_Weights, raft_small |
|
|
| MAX_GAP_EMBED = 64 |
|
|
|
|
| class SignMagnitudeRotationHead(nn.Module): |
| """Per axis: a sign logit and a log(1 + |deg|) magnitude. Output (B, 4).""" |
|
|
| def __init__(self, feat_dim: int, head_hidden: int, dropout: float) -> None: |
| super().__init__() |
| self.yaw_sign = self._branch(feat_dim, head_hidden, dropout) |
| self.yaw_mag = self._branch(feat_dim, head_hidden, dropout) |
| self.pitch_sign = self._branch(feat_dim, head_hidden, dropout) |
| self.pitch_mag = self._branch(feat_dim, head_hidden, dropout) |
|
|
| @staticmethod |
| def _branch(feat_dim: int, head_hidden: int, dropout: float) -> nn.Sequential: |
| return nn.Sequential( |
| nn.ReLU(), |
| nn.Linear(feat_dim, head_hidden), |
| nn.ReLU(inplace=True), |
| nn.Dropout(dropout), |
| nn.Linear(head_hidden, 1), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return torch.cat([self.yaw_sign(x), self.yaw_mag(x), self.pitch_sign(x), self.pitch_mag(x)], dim=1) |
|
|
|
|
| class FlowTransformer(nn.Module): |
| """Conv stem per flow pair -> 64 tokens per pair -> joint space-time encoder -> CLS heads.""" |
|
|
| def __init__(self, cfg: dict) -> None: |
| super().__init__() |
| d, grid = cfg["token_dim"], cfg["patch_grid"] |
| self.grid = grid |
| self.max_pairs = cfg["num_flow_frames"] |
| self.stem = nn.Sequential( |
| nn.Conv2d(2, 64, kernel_size=4, stride=4), |
| nn.GroupNorm(8, 64), |
| nn.GELU(), |
| nn.Conv2d(64, 128, kernel_size=2, stride=2), |
| nn.GroupNorm(8, 128), |
| nn.GELU(), |
| nn.Conv2d(128, d, kernel_size=2, stride=2), |
| ) |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, d)) |
| self.pos_embed = nn.Parameter(torch.zeros(1, grid * grid, d)) |
| self.frame_embed = nn.Parameter(torch.zeros(1, self.max_pairs, d)) |
| self.gap_embed = nn.Embedding(MAX_GAP_EMBED, d) |
| layer = nn.TransformerEncoderLayer( |
| d_model=d, |
| nhead=cfg["num_heads"], |
| dim_feedforward=cfg["ffn_dim"], |
| dropout=cfg["dropout"], |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.encoder = nn.TransformerEncoder(layer, num_layers=cfg["num_layers"], enable_nested_tensor=False) |
| self.norm = nn.LayerNorm(d) |
| self.key_head = nn.Sequential( |
| nn.Linear(d, cfg["head_hidden"]), |
| nn.GELU(), |
| nn.Dropout(cfg["dropout"]), |
| nn.Linear(cfg["head_hidden"], cfg["num_keys"]), |
| ) |
| assert cfg["rotation_head_type"] == "sign_magnitude" |
| self.rotation_head = SignMagnitudeRotationHead(d, cfg["head_hidden"], cfg["dropout"]) |
|
|
| def forward(self, flow: torch.Tensor, frame_gap: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """flow: (B, 2*N, 128, 128), channels (fx, fy) interleaved per pair. frame_gap: (B,) long.""" |
| b, c, h, w = flow.shape |
| n = c // 2 |
| assert c % 2 == 0 and n <= self.max_pairs, f"bad channel count {c}" |
| assert h == w == self.grid * 16, f"expected {self.grid * 16}x{self.grid * 16}, got {h}x{w}" |
| feat = self.stem(flow.reshape(b * n, 2, h, w)) |
| tok = feat.flatten(2).transpose(1, 2) + self.pos_embed |
| tok = tok.view(b, n, self.grid * self.grid, -1) + self.frame_embed[:, :n].unsqueeze(2) |
| cls = self.cls_token.repeat(b, 1, 1) + self.gap_embed(frame_gap.clamp(0, MAX_GAP_EMBED - 1)).unsqueeze(1) |
| x = self.encoder(torch.cat([cls, tok.flatten(1, 2)], dim=1)) |
| x = self.norm(x[:, 0]) |
| return self.key_head(x), self.rotation_head(x) |
|
|
|
|
| def decode_sign_magnitude(pred: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """(B, 4) -> (yaw_deg, pitch_deg). Sign is sigmoid >= 0.5; magnitude is expm1 of the log-magnitude.""" |
| yaw_mag = torch.expm1(pred[:, 1]).clamp_min(0.0) |
| pitch_mag = torch.expm1(pred[:, 3]).clamp_min(0.0) |
| yaw_sign = torch.where(torch.sigmoid(pred[:, 0]) >= 0.5, 1.0, -1.0) |
| pitch_sign = torch.where(torch.sigmoid(pred[:, 2]) >= 0.5, 1.0, -1.0) |
| return yaw_sign * yaw_mag, pitch_sign * pitch_mag |
|
|
|
|
| def load_model(weights: Path, cfg: dict) -> FlowTransformer: |
| model = FlowTransformer(cfg) |
| model.load_state_dict(load_file(str(weights)), strict=True) |
| return model.eval() |
|
|
|
|
| def decode_frames(video: Path, resolution: int) -> tuple[np.ndarray, float]: |
| """All frames as (N, res, res, 3) uint8 RGB, squashed to a square with INTER_AREA.""" |
| frames: list[np.ndarray] = [] |
| with av.open(str(video)) as container: |
| rate = container.streams.video[0].average_rate |
| fps = float(rate) if rate else 24.0 |
| for frame in container.decode(video=0): |
| img = frame.to_ndarray(format="rgb24") |
| frames.append(cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_AREA)) |
| assert frames, f"no frames decoded from {video}" |
| return np.stack(frames), fps |
|
|
|
|
| @torch.no_grad() |
| def window_flow(raft: nn.Module, transform, frames: np.ndarray, device: torch.device) -> torch.Tensor: |
| """(F, H, W, 3) uint8 -> (2*(F-1), H, W): RAFT flow of each consecutive pair, (fx, fy) interleaved.""" |
| t = torch.from_numpy(frames).to(device).permute(0, 3, 1, 2).float() / 255.0 |
| a, b = transform(t[:-1], t[1:]) |
| flow = raft(a, b)[-1] |
| return flow.reshape(-1, flow.shape[-2], flow.shape[-1]) |
|
|
|
|
| def predict(video: Path, weights: Path, cfg: dict, device: torch.device, stride: int | None) -> dict: |
| model = load_model(weights, cfg).to(device) |
| raft = raft_small(weights=Raft_Small_Weights.DEFAULT).to(device).eval() |
| transform = Raft_Small_Weights.DEFAULT.transforms() |
| frames, fps = decode_frames(video, cfg["resolution"]) |
| gap, n_frames = cfg["frame_gap"], cfg["num_frames"] |
| span = (n_frames - 1) * gap + 1 |
| stride = stride or span |
| keys, thr = cfg["key_order"], cfg["key_thresholds"] |
| out = [] |
| for start in range(0, len(frames) - span + 1, stride): |
| clip = frames[start : start + span : gap] |
| flow = window_flow(raft, transform, clip, device).unsqueeze(0) |
| key_logits, rot = model(flow, torch.tensor([gap], device=device)) |
| prob = torch.sigmoid(key_logits)[0].cpu().tolist() |
| yaw, pitch = decode_sign_magnitude(rot) |
| out.append( |
| { |
| "start_frame": start, |
| "end_frame": start + span - 1, |
| "start_s": round(start / fps, 3), |
| "key_probabilities": dict(zip(keys, (round(p, 4) for p in prob))), |
| "keys_pressed": [k for k, p in zip(keys, prob) if p > thr[k]], |
| "yaw_deg": round(yaw[0].item(), 4), |
| "pitch_deg": round(pitch[0].item(), 4), |
| } |
| ) |
| return {"video": str(video), "fps": fps, "num_frames": len(frames), "window_frames": span, "windows": out} |
|
|
|
|
| def main() -> None: |
| here = Path(__file__).parent |
| ap = argparse.ArgumentParser(description=__doc__) |
| ap.add_argument("video", type=Path) |
| ap.add_argument("--weights", type=Path, default=here / "model.safetensors") |
| ap.add_argument("--config", type=Path, default=here / "config.json") |
| ap.add_argument("--stride", type=int, default=None, help="source frames between windows (default: no overlap)") |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| args = ap.parse_args() |
| cfg = json.loads(args.config.read_text()) |
| result = predict(args.video, args.weights, cfg, torch.device(args.device), args.stride) |
| print(json.dumps(result, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
