File size: 5,562 Bytes
0d52562 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """
Sport-dataset dataloader for MoFlow (soccer / football).
Mirrors the structure of dataloader_nba.py's NBADatasetMinMax but:
- Loads from absolute sport data_dir (train.npy / val.npy)
- No 94/28 court rescale (soccer/football data is already in field units)
- Agent count taken from cfg.agents (23 for soccer/football)
- Supports per-scene normalization of the abs channel (matching the LED/MID sport setup)
- Soccer has no test split → we use val.npy as both val and test
"""
import os
import numpy as np
import torch
from torch.utils.data import Dataset
from utils.normalization import normalize_min_max, normalize_sqrt
from utils.utils import rotate_trajs_x_direction
def seq_collate_sport(batch):
(past_traj, fut_traj, past_traj_orig, fut_traj_orig, traj_vel) = zip(*batch)
pre_motion_3D = torch.stack(past_traj, dim=0)
fut_motion_3D = torch.stack(fut_traj, dim=0)
pre_motion_3D_orig = torch.stack(past_traj_orig, dim=0)
fut_motion_3D_orig = torch.stack(fut_traj_orig, dim=0)
fut_traj_vel = torch.stack(traj_vel, dim=0)
B = pre_motion_3D.shape[0]
A = pre_motion_3D.shape[1] # per-sample agent count
batch_size = torch.tensor(B)
traj_mask = torch.zeros(B * A, B * A)
for i in range(B):
traj_mask[i * A:(i + 1) * A, i * A:(i + 1) * A] = 1.
return {
'batch_size': batch_size,
'past_traj': pre_motion_3D,
'fut_traj': fut_motion_3D,
'past_traj_original_scale': pre_motion_3D_orig,
'fut_traj_original_scale': fut_motion_3D_orig,
'traj_mask': traj_mask,
'fut_traj_vel': fut_traj_vel,
}
class SportDatasetMinMax(Dataset):
def __init__(
self,
obs_len=10,
pred_len=20,
training=True,
num_scenes=None,
test_scenes=None,
overfit=False,
cfg=None,
data_dir='',
rotate=False,
data_norm='min_max',
):
super().__init__()
self.obs_len = obs_len
self.pred_len = pred_len
self.seq_len = obs_len + pred_len
self.traj_mean = torch.FloatTensor(cfg.traj_mean).unsqueeze(0).unsqueeze(0).unsqueeze(0)
self.per_scene_norm = bool(cfg.get('per_scene_norm', False))
# Soccer has no test split → use val.npy for both roles.
split = 'train' if training else 'val'
path = os.path.join(data_dir, f'{split}.npy')
if not os.path.isfile(path):
raise FileNotFoundError(path)
trajs = np.load(path).astype(np.float32) # (N, 30, A, 2)
if num_scenes is not None and training:
trajs = trajs[:num_scenes]
if test_scenes is not None and not training:
trajs = trajs[:test_scenes]
self.data_len = len(trajs)
print(f'[SportDatasetMinMax] {split}: {path} → {trajs.shape}')
# (N, A, T, 2)
self.traj_abs = torch.from_numpy(trajs).type(torch.float).permute(0, 2, 1, 3)
self.actor_num = self.traj_abs.shape[1]
pre_motion_3D = self.traj_abs[:, :, :self.obs_len, :]
fut_motion_3D = self.traj_abs[:, :, self.obs_len:, :]
initial_pos = pre_motion_3D[:, :, -1:]
fut_traj = (fut_motion_3D - initial_pos).contiguous()
if self.per_scene_norm:
# Scene centroid at the last observed frame: [N, 1, 1, 2]
scene_center = pre_motion_3D[:, :, -1, :].mean(dim=1, keepdim=True).unsqueeze(2)
past_traj_abs = (pre_motion_3D - scene_center).contiguous()
else:
past_traj_abs = (pre_motion_3D - self.traj_mean).contiguous()
past_traj_rel = (pre_motion_3D - initial_pos).contiguous()
if rotate:
past_traj_rel, fut_traj, past_traj_abs = rotate_trajs_x_direction(
past_traj_rel, fut_traj, past_traj_abs)
past_traj_vel = torch.cat(
(past_traj_rel[:, :, 1:] - past_traj_rel[:, :, :-1],
torch.zeros_like(past_traj_rel[:, :, -1:])), dim=2)
past_traj = torch.cat((past_traj_abs, past_traj_rel, past_traj_vel), dim=-1)
self.fut_traj_vel = torch.cat(
(fut_traj[:, :, 1:] - fut_traj[:, :, :-1],
torch.zeros_like(fut_traj[:, :, -1:])), dim=2)
if training:
cfg.fut_traj_max = fut_traj.max()
cfg.fut_traj_min = fut_traj.min()
cfg.past_traj_max = past_traj.max()
cfg.past_traj_min = past_traj.min()
self.past_traj_original_scale = past_traj
self.fut_traj_original_scale = fut_traj
self.data_norm = data_norm
if data_norm == 'min_max':
self.past_traj = normalize_min_max(
past_traj, cfg.past_traj_min, cfg.past_traj_max, -1, 1).contiguous()
self.fut_traj = normalize_min_max(
fut_traj, cfg.fut_traj_min, cfg.fut_traj_max, -1, 1).contiguous()
elif data_norm == 'sqrt':
sqrt_a_ = torch.tensor([cfg.sqrt_x_a, cfg.sqrt_y_a], device=past_traj.device)
sqrt_b_ = torch.tensor([cfg.sqrt_x_b, cfg.sqrt_y_b], device=past_traj.device)
self.past_traj = past_traj
self.fut_traj = normalize_sqrt(fut_traj, sqrt_a_, sqrt_b_).contiguous()
def __len__(self):
return self.data_len
def __getitem__(self, index):
return [
self.past_traj[index],
self.fut_traj[index],
self.past_traj_original_scale[index],
self.fut_traj_original_scale[index],
self.fut_traj_vel[index],
]
|