| """ |
| 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] |
| 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)) |
|
|
| |
| 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) |
| 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}') |
|
|
| |
| 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_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], |
| ] |
|
|
