multimodalart's picture
multimodalart HF Staff
Upload folder using huggingface_hub
f7dc16a verified
Raw History Blame Contribute Delete
16.9 kB
"""Build model inputs from raw per-player data at inference time.
Raw input CSV, one row per sample and one numbered column per player:
image_path{i} first frame of the clip (PNG/JPEG)
ticks_path{i} 64 Hz tick log (``ticks.parquet``) ending with the clip
caption{i} optional text prompt
case_id optional output number
Optional columns for longer tick logs, such as whole rounds:
start_frame first clip frame on the tick log's 32 fps timeline;
defaults to the last ``num_frames`` frames of the log
num_frames clip length; defaults to ``target_rgb_frames``
source{i} ``match_id=M/map_name=X/round=R/player=S`` identity
The tick log provides the keyboard/mouse actions, the absolute camera poses,
and the player trajectories used for grouping, which also read the 48 frames
before the clip when the log contains them. The first frame is encoded with
the Wan2.2 VAE as the clip's first latent frame. The resulting files use the
same formats as the precomputed example inputs.
"""
from __future__ import annotations
import os
from pathlib import Path
import tempfile
import numpy as np
import pandas as pd
import torch
from PIL import Image
from multics.grouping.state import row_identity, state_from_tick_files
VIDEO_FPS = 32.0
TICK_RATE = 64.0
TICKS_PER_FRAME = 2
HEIGHT = 352
WIDTH = 640
KEYBOARD_DIM = 23
LATENT_STRIDE = 4
POS_SCALE_M = 32.0
WORLD_UNIT_TO_METRE = 0.0254
STANDING_EYE_HEIGHT = 64.0
CROUCH_EYE_HEIGHT = 46.0
MOUSE_DEGREES_PER_UNIT = 15.0
# Tick action ids -> keyboard dimensions.
ACTION_TO_KEYBOARD = {
5: 11, # forward
3: 12, # back
8: 13, # move left
9: 14, # move right
6: 15, # jump
4: 16, # duck
12: 17, # walk
1: 19, # attack
13: 20, # use
2: 21, # secondary attack (scope)
10: 22, # reload
}
WEAPON_SWITCH_DIM = 18
WEAPON_CLASS_TO_KEYBOARD = {
"taser": 0,
"rifle": 2,
"sniper": 3,
"smg": 4,
"shotgun": 5,
"machine_gun": 6,
"pistol": 7,
"knife": 8,
"grenade": 9,
"objective": 10,
}
WEAPON_NAME_TO_CLASS = {
name: cls
for cls, names in {
"rifle": "ak47 aug famas galilar m4a1 m4a1_silencer m4a1_silencer_off sg556",
"sniper": "awp g3sg1 scar20 ssg08",
"smg": "bizon mac10 mp5sd mp7 mp9 p90 ump45",
"shotgun": "mag7 nova sawedoff xm1014",
"machine_gun": "m249 negev",
"pistol": (
"cz75a deagle elite fiveseven glock hkp2000 p250 revolver tec9 "
"usp_silencer usp_silencer_off uspsilenceroff"
),
"knife": (
"bayonet knife knife_butterfly knife_karambit knife_m9_bayonet "
"knife_stiletto m9bayonet"
),
"grenade": (
"decoy flashbang hegrenade highexplosivegrenade incgrenade inferno "
"molotov smokegrenade"
),
"objective": "c4 c4explosive planted_c4",
"taser": "taser",
}.items()
for name in names.split()
}
# Weapon value of a player without a weapon; "nan" is stored as a string.
NO_WEAPON = {None, "", "nan"}
# A single missing interior tick may be interpolated; larger gaps are errors.
MAX_INTERPOLATION_GAP_TICKS = 2
def tick_actions(path):
"""Return per-frame ``keyboard[N,23]`` and ``mouse[N,2]`` in degrees.
Each 32 fps frame combines two consecutive 64 Hz tick rows. ``mouse`` is
``[delta_yaw, delta_pitch]`` summed over both ticks. A weapon switch sets
dimension 18 and the dimension of the new weapon class; ticks without a
weapon (after death) never count as a switch.
"""
import pyarrow.parquet as pq
table = pq.read_table(
path, columns=["active_ids", "delta_pitch", "delta_yaw", "input_weapon"]
)
n_frames = table.num_rows // TICKS_PER_FRAME
keyboard = np.zeros((n_frames, KEYBOARD_DIM), dtype=np.float32)
mouse = np.zeros((n_frames, 2), dtype=np.float32)
active_ids = table.column("active_ids").to_pylist()
delta_pitch = (
table.column("delta_pitch").to_numpy(zero_copy_only=False).astype(np.float32)
)
delta_yaw = (
table.column("delta_yaw").to_numpy(zero_copy_only=False).astype(np.float32)
)
weapons = table.column("input_weapon").to_pylist()
previous = None
for frame in range(n_frames):
first = frame * TICKS_PER_FRAME
for tick in (first, first + 1):
for action in active_ids[tick] or ():
dim = ACTION_TO_KEYBOARD.get(int(action))
if dim is not None:
keyboard[frame, dim] = 1.0
mouse[frame, 0] += float(delta_yaw[tick])
mouse[frame, 1] += float(delta_pitch[tick])
weapon = weapons[tick]
if weapon in NO_WEAPON:
continue
if previous is not None and weapon != previous:
keyboard[frame, WEAPON_SWITCH_DIM] = 1.0
weapon_class = WEAPON_NAME_TO_CLASS.get(str(weapon))
if weapon_class is not None:
keyboard[frame, WEAPON_CLASS_TO_KEYBOARD[weapon_class]] = 1.0
previous = weapon
return keyboard, mouse
def clip_actions(path, start_frame, num_frames):
"""Keyboard/mouse tensors for one clip, mouse scaled to model units."""
keyboard, mouse = tick_actions(path)
if start_frame + num_frames > keyboard.shape[0]:
raise ValueError(
f"{path}: {keyboard.shape[0]} action frames do not cover "
f"[{start_frame}, {start_frame + num_frames})"
)
window = slice(start_frame, start_frame + num_frames)
scaled = mouse[window].copy() * (1.0 / MOUSE_DEGREES_PER_UNIT)
return {
"keyboard": torch.from_numpy(keyboard[window].copy()).float(),
"mouse": torch.from_numpy(scaled.astype(np.float32)).float(),
}
def _is_duck(active):
if active is None:
return False
if isinstance(active, str):
return active == "duck" or "duck" in active.split(",")
try:
return "duck" in active
except TypeError:
return False
def _eye_poses(path, start_frame, num_frames):
"""Eye pose ``[T,5]`` (metres, degrees), ticks, and interpolation mask."""
import pyarrow.parquet as pq
columns = ["tick", "x", "y", "z", "yaw", "pitch", "active"]
table = pq.read_table(path, columns=columns)
data = {}
for name in columns:
column = table[name].combine_chunks()
data[name] = (
column.to_pylist()
if name == "active"
else column.to_numpy(zero_copy_only=False)
)
ticks = np.asarray(data["tick"], dtype=np.int64)
if ticks.size == 0:
raise ValueError(f"Empty tick log: {path}")
if np.any(np.diff(ticks) <= 0):
raise ValueError(f"Ticks are not strictly increasing: {path}")
frames = start_frame + np.arange(num_frames, dtype=np.int64)
target = ticks[0] + TICKS_PER_FRAME * frames
loc = np.searchsorted(ticks, target)
valid = np.zeros(len(target), dtype=bool)
inside = loc < len(ticks)
valid[inside] = ticks[loc[inside]] == target[inside]
interpolated = ~valid
left = loc.copy()
right = loc.copy()
alpha = np.zeros(len(target), dtype=np.float64)
for index in np.flatnonzero(interpolated):
j = loc[index]
if not (
0 < j < len(ticks)
and ticks[j] - ticks[j - 1] == MAX_INTERPOLATION_GAP_TICKS
and target[index] - ticks[j - 1] == 1
and ticks[j] - target[index] == 1
):
raise ValueError(
f"{path}: frame {int(frames[index])} needs tick {int(target[index])}; "
"only one missing interior tick can be interpolated"
)
left[index], right[index], alpha[index] = j - 1, j, 0.5
def eye(indices):
xyz = np.stack(
[np.asarray(data[k], dtype=np.float64)[indices] for k in ("x", "y", "z")],
axis=-1,
)
xyz[:, 2] += np.asarray(
[
CROUCH_EYE_HEIGHT if _is_duck(data["active"][int(i)]) else
STANDING_EYE_HEIGHT
for i in indices
],
dtype=np.float64,
)
return xyz
xyz = (1.0 - alpha[:, None]) * eye(left) + alpha[:, None] * eye(right)
yaw_values = np.asarray(data["yaw"], dtype=np.float64)
yaw_delta = (yaw_values[right] - yaw_values[left] + 180.0) % 360.0 - 180.0
yaw = yaw_values[left] + alpha * yaw_delta
pitch_values = np.asarray(data["pitch"], dtype=np.float64)
pitch = (1.0 - alpha) * pitch_values[left] + alpha * pitch_values[right]
pose = np.concatenate([xyz, yaw[:, None], pitch[:, None]], axis=-1)
pose[:, :3] *= WORLD_UNIT_TO_METRE
return pose, target, interpolated
def _camera_to_world(pose):
"""``[T,5]`` eye poses -> ``[T,4,4]`` c2w with right/up/back local axes.
Yaw 0 points along +X and increases toward +Y; positive pitch looks down.
"""
yaw = np.deg2rad(pose[..., 3])
pitch = np.deg2rad(pose[..., 4])
forward = np.stack(
[np.cos(pitch) * np.cos(yaw), np.cos(pitch) * np.sin(yaw), -np.sin(pitch)],
axis=-1,
)
forward /= np.linalg.norm(forward, axis=-1, keepdims=True).clip(min=1e-12)
world_up = np.zeros_like(forward)
world_up[..., 2] = 1.0
right = np.cross(forward, world_up)
right /= np.linalg.norm(right, axis=-1, keepdims=True).clip(min=1e-12)
up = np.cross(right, forward)
up /= np.linalg.norm(up, axis=-1, keepdims=True).clip(min=1e-12)
out = np.broadcast_to(np.eye(4, dtype=np.float64), (*pose.shape[:-1], 4, 4)).copy()
out[..., :3, 0] = right
out[..., :3, 1] = up
out[..., :3, 2] = -forward
out[..., :3, 3] = pose[..., :3]
return out
def clip_camera(path, start_frame, num_frames, player_key):
"""Absolute camera poses at the latent frames of one clip."""
pose, ticks, interpolated = _eye_poses(path, start_frame, num_frames)
latent = np.arange(0, num_frames, LATENT_STRIDE, dtype=np.int64)
payload = {
"c2w_abs": _camera_to_world(pose[latent]).astype(np.float64),
"rgb_ticks": ticks,
"latent_ticks": ticks[latent],
"latent_rgb_indices": latent,
"interpolated_rgb_mask": interpolated,
"interpolated_latent_mask": interpolated[latent],
"rgb_source_frame_indices": start_frame + np.arange(num_frames, dtype=np.int64),
"player_key": np.asarray(player_key, dtype=np.int64),
"start_frame": np.asarray(start_frame, dtype=np.int64),
"num_rgb_frames": np.asarray(num_frames, dtype=np.int64),
"latent_stride": np.asarray(LATENT_STRIDE, dtype=np.int64),
"video_fps": np.asarray(VIDEO_FPS, dtype=np.float32),
"tick_rate": np.asarray(TICK_RATE, dtype=np.float32),
"world_unit_to_metre": np.asarray(WORLD_UNIT_TO_METRE, dtype=np.float32),
"pos_scale": np.asarray(POS_SCALE_M, dtype=np.float32),
"format_version": np.asarray(3, dtype=np.int32),
}
return {
key: torch.from_numpy(np.array(value, copy=True, order="C"))
for key, value in payload.items()
}
def load_first_frame(path, height=HEIGHT, width=WIDTH):
"""RGB uint8 ``[H,W,3]``; other sizes are scaled to cover and center-cropped."""
image = Image.open(path).convert("RGB")
w, h = image.size
if (w, h) != (width, height):
scale = max(width / w, height / h)
new_w, new_h = round(w * scale), round(h * scale)
image = image.resize((new_w, new_h), Image.BILINEAR)
left = int(round((new_w - width) / 2.0))
top = int(round((new_h - height) / 2.0))
image = image.crop((left, top, left + width, top + height))
return np.asarray(image)
def encode_first_frame(pipe, frame):
"""VAE latent ``[48,1,h,w]`` of one RGB uint8 frame, on the CPU."""
video = torch.from_numpy(frame[None]).float().div_(127.5).sub_(1.0)
video = video.permute(3, 0, 1, 2).contiguous().unsqueeze(0)
video = video.to(dtype=pipe.torch_dtype, device=pipe.device)
with torch.no_grad():
latent = pipe.encode_video(
video, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)
)
return latent[0].to("cpu")
def is_raw_metadata(table):
return "image_path1" in table.columns
def _value(row, column, default=None):
value = row.get(column)
if value is None or pd.isna(value) or not str(value).strip():
return default
return value
def _resolve(base, value):
path = Path(str(value))
return path if path.is_absolute() else Path(base) / path
def _clip_start_from_log_end(paths, num_frames):
import pyarrow.parquet as pq
lengths = {pq.ParquetFile(path).metadata.num_rows // TICKS_PER_FRAME for path in paths}
if len(lengths) != 1:
raise ValueError(
f"Tick logs of one row have different lengths {sorted(lengths)}; "
"set start_frame to choose the clip"
)
start = lengths.pop() - num_frames
if start < 0:
raise ValueError(f"Tick logs are shorter than {num_frames} frames: {paths}")
return start
def prepare_raw_inputs(table, dataset_path, rows, num_players, num_frames, output_dir):
"""Write actions, cameras, and grouping states for ``rows``.
Returns the path of a CSV in the precomputed-input format. Its first-frame
latents are written later by :func:`encode_raw_first_frames`, once the VAE
is loaded.
"""
output_dir = Path(output_dir).resolve()
output_dir.mkdir(parents=True, exist_ok=True)
# Isolate even concurrent runs of the same setting and case numbers.
output_dir = Path(tempfile.mkdtemp(prefix="run-", dir=output_dir))
rows = set(rows)
out_rows = []
for index, series in table.iterrows():
if index not in rows:
# Preserve source row positions without opening unselected inputs.
out_rows.append({})
continue
row = series.to_dict()
case = _value(row, "case_id")
folder = output_dir / (f"{int(case):02d}" if case is not None else f"row{index}")
count = int(_value(row, "num_frames", num_frames))
start = _value(row, "start_frame")
if start is None:
start = _clip_start_from_log_end(
[
_resolve(dataset_path, row[f"ticks_path{i}"])
for i in range(1, num_players + 1)
],
count,
)
start = int(start)
out = {"start_frame": start}
if case is not None:
out["case_id"] = int(case)
for i in range(1, num_players + 1):
source = _value(
row,
f"source{i}",
f"match_id=0/map_name=custom/round=00/player={i - 1:02d}",
)
out[f"video_path{i}"] = str(source)
out[f"video{i}_frame"] = count
out[f"caption{i}"] = _value(row, f"caption{i}", "")
out[f"image_path{i}"] = str(_resolve(dataset_path, row[f"image_path{i}"]))
out[f"ticks_path{i}"] = str(_resolve(dataset_path, row[f"ticks_path{i}"]))
identity = row_identity(out, num_players)
player_dirs = [folder / f"player={slot:02d}" for slot in identity[3]]
for i, player_dir in enumerate(player_dirs, start=1):
out[f"vae_path{i}"] = str(player_dir / "first_frame.latent.pt")
out[f"action_path{i}"] = str(player_dir / "action.pt")
out[f"camera_pose_path{i}"] = str(player_dir / "camera.pt")
out["grouping_state_path"] = str(folder / "grouping_state.npz")
out_rows.append(out)
ticks = [out[f"ticks_path{i}"] for i in range(1, num_players + 1)]
for i, (player_dir, path) in enumerate(zip(player_dirs, ticks), start=1):
player_dir.mkdir(parents=True, exist_ok=True)
torch.save(clip_actions(path, start, count), out[f"action_path{i}"])
torch.save(
clip_camera(path, start, count, identity[3][i - 1]),
out[f"camera_pose_path{i}"],
)
state_from_tick_files(ticks, identity).save(out["grouping_state_path"])
csv_path = output_dir / "metadata.csv"
csv_path.parent.mkdir(parents=True, exist_ok=True)
temporary = csv_path.with_name(f".{csv_path.name}.{os.getpid()}.tmp")
pd.DataFrame(out_rows).to_csv(temporary, index=False)
os.replace(temporary, csv_path)
return str(csv_path)
def encode_raw_first_frames(pipe, metadata_path, rows, num_players):
"""Encode the first frame of every player in ``rows`` of a prepared CSV."""
table = pd.read_csv(metadata_path)
pipe.load_models_to_device(["vae"])
for index in sorted(set(rows)):
row = table.iloc[index]
for i in range(1, num_players + 1):
frame = load_first_frame(row[f"image_path{i}"])
torch.save(encode_first_frame(pipe, frame), row[f"vae_path{i}"])