Add flow and pixel IDM weights (Apache 2.0)

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,8 @@
+IDM inverse dynamics models
+Copyright 2026 Reka AI
+
+The weights and code in this repository use the Apache License 2.0 (see LICENSE).
+
+RAFT-small optical flow weights: BSD-3-Clause, from torchvision
+(torchvision.models.optical_flow.raft_small). flow/inference.py downloads them at
+run time. They are not part of this repository and are not redistributed here.

README.md

@@ -0,0 +1,42 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: video-classification
+tags:
+- inverse-dynamics
+- camera-motion
+- optical-flow
+- counter-strike-2
+---
+
+# IDM: inverse dynamics models for camera motion
+
+Two models predict camera motion (W/A/S/D/Shift, yaw, pitch) from video.
+Both trained on Counter-Strike 2 renders.
+
+| | [flow](flow/) | [pixel](pixel/) |
+|---|---|---|
+| Input | Optical flow (RAFT-small) | Raw frames |
+| Parameters | 1,795,337 trained, plus 990,162 in frozen RAFT-small (2,785,499 in total) | 9,836,063 |
+| Output | One prediction per 17-frame window | One prediction per frame |
+| Real-video turn accuracy | 91.5 | 51.1 † |
+| Counter-Strike 2 turn accuracy | 84.9 | 70.9 † |
+| Dependencies | torch, torchvision, av, opencv, safetensors | torch, av, opencv, safetensors |
+
+† The number comes from a write-up and has no run log yet.
+
+```text
+idm-hf/
+  flow/   model.safetensors  config.json  inference.py  README.md
+  pixel/  model.safetensors  config.json  inference.py  README.md
+```
+
+Each folder runs alone: `python inference.py clip.mp4`.
+
+Status: preliminary release. The scores marked † are unconfirmed.
+
+## License
+
+- The weights and the code in this repository use the Apache License 2.0. See `LICENSE`.
+- RAFT-small weights use the BSD-3 license. They come from torchvision and are not in this repository. See `NOTICE`.
+- The models trained on Counter-Strike 2 gameplay captures. This repository holds no game assets, no clips and no training data.

flow/README.md

@@ -0,0 +1,77 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: video-classification
+tags: [inverse-dynamics, camera-motion, optical-flow, counter-strike-2]
+---
+
+# IDM flow model (flow_transformer_phase5_v4)
+
+This model predicts camera motion from video. It reads optical flow, not pixels.
+
+- Keys: W, A, S, D, Shift (one probability each).
+- Rotation: yaw and pitch in degrees.
+
+Status: preliminary release. The scores marked † are unconfirmed.
+
+## How it works
+
+```text
+video -> 5 frames (every 4th) -> RAFT-small flow x4 pairs -> FlowTransformer -> keys + yaw/pitch
+          128x128                 frozen, downloaded           1.80 M params
+```
+
+- RAFT-small is a frozen torchvision model. Its weights are not in this folder.
+  The script downloads them on first run (BSD-3 license).
+- The FlowTransformer has 1,795,337 parameters. It has 4 layers, width 192 and 4 heads.
+- One window is 17 source frames (5 sampled frames, gap 4). Each window gives one prediction.
+
+## Use
+
+```bash
+pip install torch torchvision av opencv-python safetensors numpy
+python inference.py clip.mp4 --stride 17
+```
+
+The script prints JSON. Each window has `key_probabilities`, `keys_pressed`, `yaw_deg` and `pitch_deg`.
+
+`keys_pressed` uses one threshold per key, tuned at 5 frames: W 0.54, A 0.40, S 0.54, D 0.70, Shift 0.20.
+Yaw and pitch are signed. The sign is a classifier. The size comes from a log-magnitude regressor.
+
+## Files
+
+| File | Content |
+|---|---|
+| `model.safetensors` | FlowTransformer weights only. No optimizer state. Float32. |
+| `config.json` | Architecture, thresholds, frame count, frame gap. |
+| `inference.py` | Standalone script. No internal imports. |
+
+## Results
+
+| Test set | Turn | Three-way |
+|---|---|---|
+| Real videos | 91.5 | 85.9 |
+| Action-camera footage | 72.0 † | 75.0 † |
+| Counter-Strike 2 | 84.9 | 84.5 |
+
+- † means the number comes from a write-up. It has no run log yet. Treat it as unconfirmed.
+- Key exact match on the held-out CS2 split (4,590 clips, 35,091 rows):
+  0.7960 at threshold 0.5 and 0.8390 with the tuned thresholds.
+- Best validation loss in the checkpoint: 0.8741 (epoch 10).
+
+## Limits
+
+- The model trained on Counter-Strike 2 renders only.
+- The default frame gap is 4 source frames, tuned on game footage. Other gaps change the result.
+- Real walking is slower than game movement. On walking-tour video a gap of 12 gave a strong forward signal.
+  The gap comes from `frame_gap` in `config.json`. Copy the file, set the value, and pass it with `--config`.
+  Use `--stride` to match the new window length: (num_frames - 1) * frame_gap + 1 source frames.
+- The gap embedding accepts values from 0 to 63.
+- Yaw and pitch are per window, not per frame.
+
+## License
+
+- The weights and `inference.py` use the Apache License 2.0. See `LICENSE` in the repository root.
+- RAFT-small weights: BSD-3, from torchvision.
+- The models trained on Counter-Strike 2 gameplay captures. This repository holds no game assets, no clips and no training data.
+- The terms of the evaluation datasets are not checked. Do not redistribute their clips.

flow/config.json

@@ -0,0 +1,23 @@
+{
+  "model_type": "idm-flow-transformer",
+  "name": "flow_transformer_phase5_v4",
+  "num_keys": 5,
+  "key_order": ["W", "A", "S", "D", "Shift"],
+  "key_thresholds": {"W": 0.54, "A": 0.40, "S": 0.54, "D": 0.70, "Shift": 0.20},
+  "thresholds_tuned_at_num_frames": 5,
+  "token_dim": 192,
+  "num_layers": 4,
+  "num_heads": 4,
+  "ffn_dim": 512,
+  "head_hidden": 256,
+  "dropout": 0.1,
+  "num_flow_frames": 9,
+  "patch_grid": 8,
+  "rotation_head_type": "sign_magnitude",
+  "resolution": 128,
+  "num_frames": 5,
+  "frame_gap": 4,
+  "flow_backbone": "torchvision.models.optical_flow.raft_small (Raft_Small_Weights.DEFAULT, frozen, not in this repo)",
+  "num_parameters": 1795337,
+  "training_checkpoint": {"epoch": 10, "best_val_loss": 0.8741227431995112, "run_ledger_status": "CANDIDATE"}
+}

flow/inference.py

@@ -0,0 +1,190 @@
+"""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]  # (F-1, 2, H, W)
+    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()

flow/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24a5dbfcf83306794ae5f1d4afb032a41f06faa5620224fe5015e36f30353ce9
+size 7189300

pixel/README.md

@@ -0,0 +1,70 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: video-classification
+tags: [inverse-dynamics, camera-motion, counter-strike-2]
+---
+
+# IDM pixel model (Z1-106k-full)
+
+This model predicts camera motion from raw video frames. It has no optical-flow stage.
+
+- Keys: W, A, S, D, Shift (one probability per frame).
+- Rotation: yaw and pitch in degrees, per frame.
+
+## How it works
+
+```text
+video -> 128x128 RGB -> IMPALA CNN per frame -> temporal attention -> keys + yaw/pitch per frame
+         ImageNet norm    3 blocks, 64/128/128    2 layers, width 512
+```
+
+- The model has 9,836,063 parameters.
+- The rotation head is a hybrid: a bin classifier plus a residual.
+  Yaw has 13 bins. Pitch has 11 bins. The result is the bin centre plus the residual.
+- The script cuts the video into 128-frame windows and averages any overlap.
+  A video with fewer than 16 frames gives no prediction.
+
+## Use
+
+```bash
+pip install torch av opencv-python safetensors numpy
+python inference.py clip.mp4
+```
+
+The script prints JSON with one row per frame:
+`key_probabilities`, `keys_pressed` (threshold 0.5), `yaw_deg` and `pitch_deg`.
+
+## Files
+
+| File | Content |
+|---|---|
+| `model.safetensors` | Weights only. Float32. The CNN keys are stored as `cnn.blocks.*`. |
+| `config.json` | Architecture and windowing. |
+| `inference.py` | Standalone script. No internal imports. |
+
+## Results
+
+| Test set | Turn | Three-way |
+|---|---|---|
+| Real videos | 51.1 † | 30.8 † |
+| Action-camera footage | 76.0 † | 58.3 † |
+| Counter-Strike 2 | 70.9 † | 77.5 † |
+
+- † means the number comes from a write-up. It has no run log yet. Treat it as unconfirmed.
+- The internal eval report shows about 0% on out-of-distribution real-world and marketplace footage.
+
+## Limits
+
+- The model trained on Counter-Strike 2 renders only.
+- It does not work on real-world footage. Use the flow model for that.
+- The 128x128 resize squashes 16:9 frames to a square. The training data did the same.
+- The training window was 150 frames. The reference inference code used 128. This folder uses 128.
+- `num_heads` (8) is not stored in the checkpoint. It is the default of the training code.
+  The weights load with `strict=True`, and the output matches the original code exactly.
+
+## License
+
+- The weights and `inference.py` use the Apache License 2.0. See `LICENSE` in the repository root.
+- The models trained on Counter-Strike 2 gameplay captures. This repository holds no game assets, no clips and no training data.
+- The terms of the evaluation datasets are not checked.

pixel/config.json

@@ -0,0 +1,23 @@
+{
+  "model_type": "idm-impala-vpt",
+  "name": "Z1-106k-full",
+  "num_keys": 5,
+  "key_order": ["W", "A", "S", "D", "Shift"],
+  "key_threshold": 0.5,
+  "channels": [64, 128, 128],
+  "gn_groups": 8,
+  "spatial_pool_size": 4,
+  "temporal_hidden": 512,
+  "temporal_layers": 2,
+  "num_heads": 8,
+  "head_hidden": 256,
+  "dropout": 0.2,
+  "rotation_head_type": "hybrid",
+  "resolution": 128,
+  "window_frames": 128,
+  "window_stride": 128,
+  "min_frames": 16,
+  "training_window_frames": 150,
+  "num_parameters": 9836063,
+  "notes": "num_heads is not stored in the checkpoint; 8 is the default in the training code. The training window was 150 frames; the reference inference library used 128."
+}

pixel/inference.py

@@ -0,0 +1,229 @@
+"""Pixel IDM: predict W/A/S/D/Shift and yaw/pitch per frame from a video.
+
+Usage: python inference.py VIDEO [--weights model.safetensors] [--config config.json]
+
+Pipeline: decode -> 128x128 RGB (squashed) -> ImageNet normalisation -> windows of
+window_frames -> per-frame IMPALA CNN + temporal attention -> per-frame heads.
+Overlapping windows are averaged. Output is JSON on stdout.
+Needs: torch, 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
+import torch.nn.functional as F
+from safetensors.torch import load_file
+
+YAW_BIN_CENTERS = [-70, -2.2, -0.7, -0.28, -0.12, -0.04, 0, 0.04, 0.12, 0.28, 0.7, 2.2, 70]
+PITCH_BIN_CENTERS = [-16, -0.35, -0.13, -0.06, -0.02, 0, 0.02, 0.06, 0.13, 0.35, 16]
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, ch: int, gn_groups: int) -> None:
+        super().__init__()
+        self.gn1 = nn.GroupNorm(gn_groups, ch)
+        self.conv1 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
+        self.gn2 = nn.GroupNorm(gn_groups, ch)
+        self.conv2 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.conv1(F.relu(self.gn1(x)))
+        out = self.conv2(F.relu(self.gn2(out)))
+        return out + x
+
+
+class ImpalaBlock(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, gn_groups: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_ch, out_ch, 3, padding=1)
+        self.pool = nn.MaxPool2d(3, stride=2, padding=1)
+        self.res1 = ResBlock(out_ch, gn_groups)
+        self.res2 = ResBlock(out_ch, gn_groups)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.res2(self.res1(self.pool(self.conv(x))))
+
+
+class ImpalaCNN(nn.Module):
+    def __init__(self, channels: list[int], gn_groups: int, pool: int) -> None:
+        super().__init__()
+        chs = [3, *channels]
+        self.blocks = nn.Sequential(*[ImpalaBlock(chs[i], chs[i + 1], gn_groups) for i in range(len(channels))])
+        self.pool = nn.AdaptiveAvgPool2d(pool)
+        self.feat_dim = channels[-1] * pool * pool
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.pool(self.blocks(x)).flatten(1)
+
+
+class PointwiseMLP(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fc1 = nn.Linear(dim, dim * 4)
+        self.fc2 = nn.Linear(dim * 4, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.fc2(F.gelu(self.fc1(self.norm(x))))
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
+        self.mlp = PointwiseMLP(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        n = self.norm(x)
+        x = x + self.attn(n, n, n, need_weights=False)[0]
+        return self.mlp(x)
+
+
+def sub_head(feat_dim: int, hidden: int, dropout: float, out_dim: int) -> nn.Sequential:
+    return nn.Sequential(
+        nn.ReLU(), nn.Linear(feat_dim, hidden), nn.ReLU(inplace=True), nn.Dropout(dropout), nn.Linear(hidden, out_dim)
+    )
+
+
+class HybridRotationHead(nn.Module):
+    """Output (B, 26): yaw bin logits (13), yaw residual, pitch bin logits (11), pitch residual."""
+
+    def __init__(self, feat_dim: int, hidden: int, dropout: float) -> None:
+        super().__init__()
+        self.yaw_cls = sub_head(feat_dim, hidden, dropout, len(YAW_BIN_CENTERS))
+        self.yaw_res = sub_head(feat_dim, hidden, dropout, 1)
+        self.pitch_cls = sub_head(feat_dim, hidden, dropout, len(PITCH_BIN_CENTERS))
+        self.pitch_res = sub_head(feat_dim, hidden, dropout, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.cat([self.yaw_cls(x), self.yaw_res(x), self.pitch_cls(x), self.pitch_res(x)], dim=1)
+
+
+class IDMImpalaVPT(nn.Module):
+    """Per-frame IMPALA CNN -> temporal attention over the window -> per-frame heads."""
+
+    def __init__(self, cfg: dict) -> None:
+        super().__init__()
+        hid = cfg["temporal_hidden"]
+        self.cnn = ImpalaCNN(cfg["channels"], cfg["gn_groups"], cfg["spatial_pool_size"])
+        self.cnn_proj = nn.Linear(self.cnn.feat_dim, hid)
+        self.pre_temporal_norm = nn.LayerNorm(hid)
+        self.temporal_blocks = nn.Sequential(
+            *[ResidualAttentionBlock(hid, cfg["num_heads"]) for _ in range(cfg["temporal_layers"])]
+        )
+        self.post_temporal = nn.Sequential(nn.ReLU(), nn.Linear(hid, hid), nn.LayerNorm(hid))
+        self.key_head = sub_head(hid, cfg["head_hidden"], cfg["dropout"], cfg["num_keys"])
+        assert cfg["rotation_head_type"] == "hybrid"
+        self.rotation_head = HybridRotationHead(hid, cfg["head_hidden"], cfg["dropout"])
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """x: (B, N, 3, H, W) normalised frames -> key logits (B, N, 5), rotation raw (B, N, 26)."""
+        b, n = x.shape[:2]
+        f = self.cnn(x.flatten(0, 1)).view(b, n, -1)
+        f = self.post_temporal(self.temporal_blocks(self.pre_temporal_norm(self.cnn_proj(f))))
+        flat = f.reshape(b * n, -1)
+        return self.key_head(flat).view(b, n, -1), self.rotation_head(flat).view(b, n, -1)
+
+
+def decode_hybrid_rotation(raw: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """(N, 26) -> (yaw_deg, pitch_deg): argmax bin centre plus residual."""
+    ny = len(YAW_BIN_CENTERS)
+    yaw = np.array(YAW_BIN_CENTERS, dtype=np.float32)[raw[:, :ny].argmax(1)] + raw[:, ny]
+    pitch_logits = raw[:, ny + 1 : ny + 1 + len(PITCH_BIN_CENTERS)]
+    pitch = np.array(PITCH_BIN_CENTERS, dtype=np.float32)[pitch_logits.argmax(1)] + raw[:, -1]
+    return yaw, pitch
+
+
+def load_model(weights: Path, cfg: dict) -> IDMImpalaVPT:
+    model = IDMImpalaVPT(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
+
+
+def generate_windows(total: int, size: int, stride: int, min_frames: int) -> list[tuple[int, int]]:
+    if total < min_frames:
+        return []
+    if total < size:
+        return [(0, total)]
+    windows, start = [], 0
+    while start < total:
+        end = min(start + size, total)
+        windows.append((start, end))
+        if end == total:
+            break
+        start += stride
+    return windows
+
+
+@torch.no_grad()
+def predict(video: Path, weights: Path, cfg: dict, device: torch.device) -> dict:
+    model = load_model(weights, cfg).to(device)
+    frames, fps = decode_frames(video, cfg["resolution"])
+    total = len(frames)
+    windows = generate_windows(total, cfg["window_frames"], cfg["window_stride"], cfg["min_frames"])
+    assert windows, f"video has {total} frames; need at least {cfg['min_frames']}"
+    norm = torch.from_numpy(((frames.astype(np.float32) / 255.0 - IMAGENET_MEAN) / IMAGENET_STD)).permute(0, 3, 1, 2)
+    dim = cfg["num_keys"] + 26
+    acc, count = np.zeros((total, dim), np.float64), np.zeros(total, np.float64)
+    for start, end in windows:
+        x = norm[start:end].unsqueeze(0).to(device)
+        key_logits, rot = model(x)
+        raw = torch.cat([torch.sigmoid(key_logits), rot], dim=-1)[0].cpu().numpy()
+        acc[start:end] += raw
+        count[start:end] += 1
+    acc /= count[:, None]
+    nk = cfg["num_keys"]
+    yaw, pitch = decode_hybrid_rotation(acc[:, nk:].astype(np.float32))
+    prob = acc[:, :nk]
+    rows = [
+        {
+            "frame": i,
+            "t_s": round(i / fps, 3),
+            "key_probabilities": dict(zip(cfg["key_order"], (round(float(p), 4) for p in prob[i]))),
+            "keys_pressed": [k for k, p in zip(cfg["key_order"], prob[i]) if p > cfg["key_threshold"]],
+            "yaw_deg": round(float(yaw[i]), 4),
+            "pitch_deg": round(float(pitch[i]), 4),
+        }
+        for i in range(total)
+    ]
+    return {"video": str(video), "fps": fps, "num_frames": total, "frames": rows}
+
+
+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("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = ap.parse_args()
+    cfg = json.loads(args.config.read_text())
+    print(json.dumps(predict(args.video, args.weights, cfg, torch.device(args.device)), indent=2))
+
+
+if __name__ == "__main__":
+    main()

pixel/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4498b70cb763b3595599a717d6a80a22b581356fee3a62dacf402aba38636953
+size 39353324