File size: 3,024 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 | """
Sport-dataset dataloader for LED (soccer / football).
Schema mirrors NBADataset (data/dataloader_nba.py) so the trainer can reuse
its own collate and loop logic with minimal changes, but parameterized by:
- data_dir : absolute directory with train.npy / val.npy [/ test.npy]
- num_agents : agents per scene (23 for soccer/football)
- split : 'train' or 'val' (val is used as the test set too)
The raw .npy files are shape (N, 30, A, 2). We split 10 past / 20 future
and return per-sample tensors in the same tuple layout as NBADataset so
the existing seq_collate logic in the sport trainer can produce the same
'pre_motion_3D'/'fut_motion_3D'/... dict.
"""
import os
import numpy as np
import torch
from torch.utils.data import Dataset
def sport_seq_collate(data):
(pre_motion_3D, fut_motion_3D, pre_motion_mask, fut_motion_mask) = zip(*data)
pre_motion_3D = torch.stack(pre_motion_3D, dim=0)
fut_motion_3D = torch.stack(fut_motion_3D, dim=0)
pre_motion_mask = torch.stack(pre_motion_mask, dim=0)
fut_motion_mask = torch.stack(fut_motion_mask, dim=0)
return {
'pre_motion_3D': pre_motion_3D,
'fut_motion_3D': fut_motion_3D,
'fut_motion_mask': fut_motion_mask,
'pre_motion_mask': pre_motion_mask,
'traj_scale': 1,
'pred_mask': None,
'seq': 'sport',
}
class SportDataset(Dataset):
"""LED-style dataset for soccer/football raw .npy files."""
def __init__(
self,
data_dir: str,
num_agents: int,
obs_len: int = 10,
pred_len: int = 20,
split: str = 'train',
):
super().__init__()
assert split in ('train', 'val'), f'unknown split {split}'
self.obs_len = obs_len
self.pred_len = pred_len
self.seq_len = obs_len + pred_len
self.num_agents = num_agents
fname = f'{split}.npy'
path = os.path.join(data_dir, fname)
if not os.path.isfile(path):
raise FileNotFoundError(f'missing {path}')
trajs = np.load(path).astype(np.float32) # (N, 30, A, 2)
assert trajs.ndim == 4 and trajs.shape[1] == self.seq_len, \
f'expected (N, {self.seq_len}, A, 2), got {trajs.shape}'
assert trajs.shape[2] == num_agents, \
f'expected {num_agents} agents, got {trajs.shape[2]}'
self.data_len = len(trajs)
print(f'[SportDataset] {split} {path} → {trajs.shape}')
# [N, A, T, 2] (LED convention: agent dim before time)
self.traj_abs = torch.from_numpy(trajs).permute(0, 2, 1, 3).contiguous()
def __len__(self):
return self.data_len
def __getitem__(self, index):
pre = self.traj_abs[index, :, :self.obs_len, :]
fut = self.traj_abs[index, :, self.obs_len:, :]
pre_mask = torch.ones(self.num_agents, self.obs_len)
fut_mask = torch.ones(self.num_agents, self.pred_len)
return [pre, fut, pre_mask, fut_mask]
|