Spaces:
Running on Zero
Running on Zero
| """Lossless tick states; model camera embeddings cannot replace this input.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass, fields | |
| from pathlib import Path | |
| import re | |
| import numpy as np | |
| class GroupingState: | |
| ticks: np.ndarray | |
| pos: np.ndarray | |
| yaw: np.ndarray | |
| pitch: np.ndarray | |
| alive: np.ndarray | |
| present: np.ndarray | |
| player_slots: np.ndarray | |
| round_tick_start: int | |
| round_tick_end: int # inclusive, before cropping | |
| match_id: int | |
| map_name: str | |
| round_no: int | |
| clip_start_frame: int | |
| clip_num_frames: int | |
| def num_players(self): | |
| return len(self.player_slots) | |
| def frame_to_tick(self, frame): | |
| return int(self.round_tick_start + round(frame * 64.0 / 32.0)) | |
| def frame_window(self, start_frame, num_frames): | |
| return slice( | |
| int(np.searchsorted(self.ticks, self.frame_to_tick(start_frame))), | |
| int( | |
| np.searchsorted( | |
| self.ticks, self.frame_to_tick(start_frame + num_frames) | |
| ) | |
| ), | |
| ) | |
| def eye_pos(self): | |
| # Addition in float32, THEN float64 conversion in relevance is deliberate. | |
| eye = self.pos.copy() | |
| eye[..., 2] += 64.0 | |
| return eye | |
| def validate(self): | |
| t, p = len(self.ticks), self.num_players | |
| if not t or p < 2 or p > 10: | |
| raise ValueError("Grouping state requires ticks and 2–10 players") | |
| shapes = { | |
| "ticks": ((t,), np.int64), | |
| "pos": ((t, p, 3), np.float32), | |
| "yaw": ((t, p), np.float32), | |
| "pitch": ((t, p), np.float32), | |
| "alive": ((t, p), np.bool_), | |
| "present": ((t, p), np.bool_), | |
| "player_slots": ((p,), np.int64), | |
| } | |
| for name, (shape, dtype) in shapes.items(): | |
| value = getattr(self, name) | |
| if value.shape != shape or value.dtype != dtype: | |
| raise ValueError(f"Invalid {name}: expected {shape}/{np.dtype(dtype)}") | |
| if not np.all(np.diff(self.ticks) == 1): | |
| raise ValueError( | |
| "Grouping tick grid must be continuous (retain missing ticks)" | |
| ) | |
| if len(set(self.player_slots.tolist())) != p or np.any( | |
| (self.player_slots < 0) | (self.player_slots >= 10) | |
| ): | |
| raise ValueError( | |
| "Grouping player_slots must be distinct physical slots in [0,9]" | |
| ) | |
| if self.clip_start_frame < 0 or self.clip_num_frames <= 0: | |
| raise ValueError("Invalid grouping clip range") | |
| lo = self.frame_to_tick(max(0, self.clip_start_frame - 48)) | |
| hi = self.frame_to_tick(self.clip_start_frame + self.clip_num_frames) - 1 | |
| if ( | |
| self.ticks[0] > lo | |
| or self.ticks[-1] < hi | |
| or self.round_tick_start > self.ticks[0] | |
| or self.round_tick_end < self.ticks[-1] | |
| ): | |
| raise ValueError( | |
| "Grouping state does not cover the clip and its 48-frame history" | |
| ) | |
| active = self.alive & self.present | |
| for name in ("pos", "yaw", "pitch"): | |
| if not np.isfinite(getattr(self, name)[active]).all(): | |
| raise ValueError(f"Non-finite {name} for an active player") | |
| def save(self, path): | |
| self.validate() | |
| payload = {f.name: getattr(self, f.name) for f in fields(self)} | |
| np.savez_compressed(path, format_version=np.int32(1), **payload) | |
| def load(cls, path): | |
| with np.load(path, allow_pickle=False) as data: | |
| expected = {f.name for f in fields(cls)} | |
| if set(data.files) != expected | {"format_version"}: | |
| raise ValueError(f"Invalid grouping state fields: {path}") | |
| if data["format_version"].shape != () or data["format_version"].item() != 1: | |
| raise ValueError(f"Unsupported grouping state version: {path}") | |
| values = {name: data[name].copy() for name in expected} | |
| for name in ( | |
| "round_tick_start", | |
| "round_tick_end", | |
| "match_id", | |
| "round_no", | |
| "clip_start_frame", | |
| "clip_num_frames", | |
| ): | |
| value = values[name] | |
| if value.shape != () or value.dtype.kind not in "iu": | |
| raise ValueError(f"{name} must be an integer scalar: {path}") | |
| values[name] = int(value) | |
| if values["map_name"].shape != () or values["map_name"].dtype.kind != "U": | |
| raise ValueError(f"map_name must be a string scalar: {path}") | |
| values["map_name"] = str(values["map_name"].item()) | |
| result = cls(**values) | |
| result.validate() | |
| return result | |
| def row_identity(row, num_players): | |
| """CSV order is distinct from physical player slots, especially for P6.""" | |
| pattern = r"match_id=(\d+)/map_name=([^/]+)/round=(\d+)/player=(\d+)(?:/|$)" | |
| matches = [ | |
| re.search(pattern, str(row[f"video_path{i}"]).replace("\\", "/")) | |
| for i in range(1, num_players + 1) | |
| ] | |
| if any(m is None for m in matches): | |
| raise ValueError( | |
| "Grouping requires video_path1..N containing round/player identity" | |
| ) | |
| keys = {(int(m[1]), m[2], int(m[3])) for m in matches} | |
| if len(keys) != 1: | |
| raise ValueError("All grouping players must belong to the same round") | |
| slots = [int(m[4]) for m in matches] | |
| if len(set(slots)) != num_players: | |
| raise ValueError("Duplicate physical players in grouping row") | |
| def integer(value, name): | |
| if isinstance(value, (bool, np.bool_)) or not np.isfinite(float(value)): | |
| raise ValueError(f"{name} must be an integer") | |
| result = int(value) | |
| if float(value) != result: | |
| raise ValueError(f"{name} must be an integer") | |
| return result | |
| start = integer(row["start_frame"], "start_frame") | |
| lengths = [ | |
| integer(row[f"video{i}_frame"], f"video{i}_frame") | |
| for i in range(1, num_players + 1) | |
| ] | |
| if start < 0 or min(lengths) <= 0: | |
| raise ValueError("Invalid grouping frame range") | |
| return (*next(iter(keys)), slots, start, min(lengths)) | |
| def state_from_parquet(rounds_root, row, num_players): | |
| """Read only required columns, align full-round states, then crop losslessly.""" | |
| match, map_name, round_no, slots, start, count = row_identity(row, num_players) | |
| folder = ( | |
| Path(rounds_root) | |
| / f"match_id={match}" | |
| / f"map_name={map_name}" | |
| / f"round={round_no:02d}" | |
| ) | |
| return state_from_tick_files( | |
| [folder / f"player={s:02d}" / "ticks.parquet" for s in slots], | |
| (match, map_name, round_no, slots, start, count), | |
| ) | |
| def state_from_tick_files(paths, identity): | |
| """Build a state from one tick file per player, in CSV player order. | |
| ``identity`` is ``(match_id, map_name, round_no, player_slots, start_frame, | |
| num_frames)``, as returned by :func:`row_identity`. | |
| """ | |
| import pandas as pd | |
| match, map_name, round_no, slots, start, count = identity | |
| num_players = len(paths) | |
| columns = ["tick", "x", "y", "z", "yaw", "pitch", "is_alive"] | |
| frames = [ | |
| pd.read_parquet(path, columns=columns).sort_values("tick", kind="stable") | |
| for path in paths | |
| ] | |
| if any(df.empty for df in frames): | |
| raise ValueError("Empty player tick file") | |
| tick0 = min(int(df.tick.iloc[0]) for df in frames) | |
| tick1 = max(int(df.tick.iloc[-1]) for df in frames) | |
| ticks = np.arange(tick0, tick1 + 1, dtype=np.int64) | |
| shape = (len(ticks), num_players) | |
| pos = np.full((*shape, 3), np.nan, np.float32) | |
| yaw, pitch = (np.full(shape, np.nan, np.float32) for _ in range(2)) | |
| alive, present = (np.zeros(shape, bool) for _ in range(2)) | |
| for k, df in enumerate(frames): | |
| idx = df.tick.to_numpy(np.int64) - tick0 | |
| for j, name in enumerate(("x", "y", "z")): | |
| pos[idx, k, j] = df[name].to_numpy(np.float32) | |
| yaw[idx, k] = df.yaw.to_numpy(np.float32) | |
| pitch[idx, k] = df.pitch.to_numpy(np.float32) | |
| alive[idx, k] = df.is_alive.to_numpy(bool) | |
| present[idx, k] = True | |
| last = np.flatnonzero(present[:, k])[-1] | |
| pos[last + 1 :, k] = pos[last, k] | |
| yaw[last + 1 :, k] = yaw[last, k] | |
| pitch[last + 1 :, k] = pitch[last, k] | |
| crop = slice(max(0, start - 48) * 2, (start + count) * 2) | |
| result = GroupingState( | |
| ticks[crop], | |
| pos[crop], | |
| yaw[crop], | |
| pitch[crop], | |
| alive[crop], | |
| present[crop], | |
| np.asarray(slots, np.int64), | |
| tick0, | |
| tick1, | |
| match, | |
| map_name, | |
| round_no, | |
| start, | |
| count, | |
| ) | |
| result.validate() | |
| return result | |
| def eligible_players(world, sl): | |
| present = world.present[sl] | |
| grid = present.any(1) if len(present) else np.zeros(0, bool) | |
| if not grid.any(): | |
| return np.zeros(world.num_players, bool) | |
| return present[grid].mean(0) >= 1.0 | |
| def update_slices(world, sl, f0, every): | |
| n = int(round((sl.stop - sl.start) / 2.0)) | |
| if every <= 0 or every >= n: | |
| return [(f0, sl)] | |
| out = [] | |
| for k in range(0, n, every): | |
| ssl = world.frame_window(f0 + k, min(every, n - k)) | |
| if ssl.stop > ssl.start: | |
| out.append((f0 + k, ssl)) | |
| return out | |