MoFlow/data/dataloader_eth_ucy.py

import os
import pickle
import glob
import numpy as np
import math
from einops import rearrange
import torch
import matplotlib.pyplot as plt
from utils.normalization import normalize_min_max


def seq_collate_eth(batch):
    (index, past_traj, fut_traj, past_traj_orig, fut_traj_orig, traj_vel) = zip(*batch)
    indexes = torch.stack(index, dim=0)
    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)

    batch_size = torch.tensor(pre_motion_3D.shape[0]) ### bt 
    data = {
        'indexes': indexes,
        '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,
        'fut_traj_vel': fut_traj_vel,  
    }

    return data 


def seq_collate_imle_train(batch):
    (past_traj, fut_traj, past_traj_orig, fut_traj_orig, traj_vel, y_t, y_pred_data) = 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)
    y_t = torch.stack(y_t, dim=0)
    y_pred_data = torch.stack(y_pred_data,dim=0)

    batch_size = torch.tensor(pre_motion_3D.shape[0]) ### bt 
    data = {
        '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,
        'fut_traj_vel': fut_traj_vel,
        'y_t': y_t,
        'y_pred_data': y_pred_data
    }

    return data


def rotate_traj(past_rel, future_rel, past_abs, agents=2, rotate_time_frame=0, subset='eth'):
    past_rel = rearrange(past_rel, 'b a p d -> (b a) p d')
    past_abs = rearrange(past_abs, 'b a p d -> (b a) p d')
    future_rel = rearrange(future_rel, 'b a f d -> (b a) f d')
    past_diff = past_rel[:, rotate_time_frame]
    # past_diff = past[:, rotate_time_frame] - past[:, rotate_time_frame-1]
    past_theta = torch.atan(torch.div(past_diff[:, 1], past_diff[:, 0]+1e-5))
    past_theta = torch.where((past_diff[:, 0]<0), past_theta+math.pi, past_theta)
    
    rotate_matrix = torch.zeros((past_theta.size(0), 2, 2)).to(past_theta.device)
    rotate_matrix[:, 0, 0] = torch.cos(past_theta)
    rotate_matrix[:, 0, 1] = torch.sin(past_theta)
    rotate_matrix[:, 1, 0] = - torch.sin(past_theta)
    rotate_matrix[:, 1, 1] = torch.cos(past_theta)
    ### store the inverse of this rotate_matrix
    # inverse_rotate_matrix = rotate_matrix.transpose(1, 2)
    # np.save(f'inverse_rotate_matrix_{subset}.npy', inverse_rotate_matrix.detach().cpu().numpy())
    # exit()

    past_after = torch.matmul(rotate_matrix, past_rel.transpose(1, 2)).transpose(1, 2)
    future_after = torch.matmul(rotate_matrix, future_rel.transpose(1, 2)).transpose(1, 2)
    past_abs_after = torch.matmul(rotate_matrix, past_abs.transpose(1, 2)).transpose(1, 2)
    past_after = rearrange(past_after, '(b a) p d -> b a p d', a=agents)
    future_after = rearrange(future_after, '(b a) f d -> b a f d', a=agents)
    past_abs_after = rearrange(past_abs_after, '(b a) p d -> b a p d', a=agents)

    return past_after, future_after, past_abs_after


class ETHDataset(object):
    def __init__(self, cfg, training=True, data_dir = None, subset = None, rotate_time_frame=0, imle=False, type='original'):
        ### LED version of preprocessed data
        if type == 'LED':
            data_file_path = os.path.join(data_dir, type, '{:s}_data_{:s}.npy'.format(subset, 'train' if training else 'test'))
            num_file_path = os.path.join(data_dir, type, '{:s}_num_{:s}.npy'.format(subset, 'train' if training else 'test'))

            all_data = np.load(data_file_path)
            all_num = np.load(num_file_path)
            self.all_data = torch.Tensor(all_data)
            self.all_num = torch.Tensor(all_num) 
        elif type == 'original':
            ### Original version of preprocessed data
            data_file_path = os.path.join(data_dir, type, subset,'{:s}_{:s}.pkl'.format(subset, 'train' if training else 'test'))
            all_data = pickle.load(open(data_file_path, 'rb'))
            all_data = all_data['traj']
            self.all_data = torch.Tensor(all_data)   #[A, T, 2]
            self.all_data = self.all_data[:,None,:,:]  #[A, 1, T, 2]
        else:
            raise ValueError('Invalid type')


        self.cfg = cfg
        self.rotate_time_frame = rotate_time_frame
        self.imle = imle

        
        ### set the agent_num in the cfg
        cfg.agents = self.all_data.shape[1]
        cfg.MODEL.CONTEXT_ENCODER.AGENTS = cfg.agents
        
        ### compute past and future trajectories
        past_traj_abs = self.all_data[:,:,:cfg.past_frames]
        initial_pos = past_traj_abs[:, :, -1:]
        past_traj_rel = (past_traj_abs - initial_pos).contiguous()
        fut_traj = (self.all_data[:,:,cfg.past_frames:] - initial_pos).contiguous()
        if cfg.rotate:
            past_traj_rel, fut_traj, past_traj_abs = rotate_traj(past_traj_rel, fut_traj, past_traj_abs, cfg.agents, rotate_time_frame, subset)
        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)

        self.rotate_aug = cfg.rotate_aug and training

        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()
        
        ### record the original to avoid numerical errors
        self.past_traj_original_scale = past_traj
        self.fut_traj_original_scale = fut_traj
   
        ### min-max linear normalization
        if cfg.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 cfg.data_norm == 'original':
            self.past_traj = past_traj
            self.fut_traj = fut_traj


        """load distillation target"""
        if imle:
            os.makedirs(os.path.join(data_dir, f'imle/{subset}'), exist_ok=True)
            pkl_ls = sorted(glob.glob(os.path.join(data_dir, f'imle/{subset}/*train*.pkl')))

            keys_ls = ['past_traj', 'fut_traj', 'past_traj_original_scale', 'fut_traj_original_scale', 'fut_traj_vel', 'y_t', 'y_pred_data']
            imle_data_dict = {}
            total_scenes_loaded_ = 0
            for i_pkl, cur_pkl in enumerate(pkl_ls):
                data = pickle.load(open(cur_pkl, 'rb'))

                if i_pkl == 0:
                    self.imle_meta_data = data['meta_data']
                
                for key in keys_ls:
                    if key not in imle_data_dict:
                        imle_data_dict[key] = []
                    if key == 'y_t':
                        imle_data_dict[key].append(data[key][:, -1])
                    else:
                        imle_data_dict[key].append(data[key])

                total_scenes_loaded_ += data['past_traj'].shape[0]

                if total_scenes_loaded_ >= len(self.past_traj):
                    break

                if i_pkl == 0:
                    # y_t_original_scale_ = unnormalize_min_max(torch.from_numpy(data['y_t'][:, -1]), cfg.fut_traj_min, cfg.fut_traj_max, -1, 1)
                    # y_pred_data_original_scale_ = torch.from_numpy(data['y_pred_data'])
                    # assert torch.sum(torch.abs(y_t_original_scale_ - y_pred_data_original_scale_)) < 1e-5, 'IMLE data is not consistent'
                            
                    # past_tarj_original_scale_ = torch.from_numpy(data['past_traj_original_scale'])
                    # assert torch.sum(torch.abs(past_tarj_original_scale_[:10] - self.past_traj_original_scale[:10])) < 1e-5, 'IMLE data is not consistent'

                    pass

            # concat the data
            for key in keys_ls:
                imle_data_dict[key] = torch.from_numpy(np.concatenate(imle_data_dict[key], axis=0))[:len(self.past_traj)]

            self.imle_data_dict = imle_data_dict

    def __len__(self):
        return self.all_data.shape[0]

    def __getitem__(self, item):
        if self.imle:
            out = [
                    self.imle_data_dict['past_traj'][item], 
                    self.imle_data_dict['fut_traj'][item],
                    self.imle_data_dict['past_traj_original_scale'][item],
                    self.imle_data_dict['fut_traj_original_scale'][item],
                    self.imle_data_dict['fut_traj_vel'][item],
                    self.imle_data_dict['y_t'][item],
                    self.imle_data_dict['y_pred_data'][item]
                ]
        else:
            ### past traj, future traj, number of pedestrians (presumbly?), index
            past_traj_norm_scale = self.past_traj[item]                             # [A, P, 6]
            fut_traj_norm_scale = self.fut_traj[item]                               # [A, F, 2] 
            past_traj_original_scale = self.past_traj_original_scale[item]          # [A, P, 6]
            fut_traj_original_scale = self.fut_traj_original_scale[item]            # [A, F, 2]
            fut_traj_vel = self.fut_traj_vel[item]                                  # [A, F, 2]   

            if self.rotate_aug:
                A = past_traj_norm_scale.size(0)
                rot_angle = torch.rand(A) * 2 * math.pi

                past_traj_abs_o, past_traj_rel_o, past_traj_vel_o = past_traj_original_scale.chunk(3, dim=-1)   # [A, P, 2] * 3

                rotate_matrix = torch.zeros((rot_angle.size(0), 2, 2)).to(past_traj_norm_scale.device)          # [A, 2, 2]
                rotate_matrix[:, 0, 0] = torch.cos(rot_angle)
                rotate_matrix[:, 0, 1] = torch.sin(rot_angle)
                rotate_matrix[:, 1, 0] = - torch.sin(rot_angle)
                rotate_matrix[:, 1, 1] = torch.cos(rot_angle)

                past_traj_abs_o_rot = torch.matmul(rotate_matrix, past_traj_abs_o.transpose(1, 2)).transpose(1, 2)
                past_traj_rel_o_rot = torch.matmul(rotate_matrix, past_traj_rel_o.transpose(1, 2)).transpose(1, 2)
                fut_traj_o_rot = torch.matmul(rotate_matrix, fut_traj_original_scale.transpose(1, 2)).transpose(1, 2)

                past_traj_vel_o_rot = torch.cat((past_traj_rel_o_rot[:, :, 1:] - past_traj_rel_o_rot[:, :, :-1], torch.zeros_like(past_traj_rel_o_rot[:,:, -1:])), dim=2)
                past_traj_rot = torch.cat((past_traj_abs_o_rot, past_traj_rel_o_rot, past_traj_vel_o_rot), dim=-1)
                
                fut_traj_vel_o = torch.cat((fut_traj_o_rot[:, :, 1:] - fut_traj_o_rot[:,:, :-1], torch.zeros_like(fut_traj_o_rot[:, :, -1:])), dim=2)

                
                past_traj_norm_scale = normalize_min_max(past_traj_rot, self.cfg.past_traj_min, self.cfg.past_traj_max, -1, 1).contiguous()
                fut_traj_norm_scale = normalize_min_max(fut_traj_o_rot, self.cfg.fut_traj_min, self.cfg.fut_traj_max, -1, 1).contiguous()
                past_traj_original_scale = past_traj_rot
                fut_traj_original_scale = fut_traj_o_rot
                fut_traj_vel = fut_traj_vel_o

            out = [
                torch.Tensor([item]).to(torch.int32),
                past_traj_norm_scale,
                fut_traj_norm_scale,
                past_traj_original_scale,
                fut_traj_original_scale, 
                fut_traj_vel
            ]
        return out

class ETHDatasetSocialGAN:
    '''
    For the purpose of saving pickle files for train/val/test data for ETH-UCY dataset
    '''
    def __init__(self, data_dir, obs_len=8, pred_len=12, skip=1, min_ped=1, delim='\t', subset='eth'):
        super(ETHDatasetSocialGAN, self).__init__()

        self.max_peds_in_frame = 0
        self.data_dir = data_dir
        self.obs_len = obs_len
        self.pred_len = pred_len
        self.skip = skip
        self.seq_len = self.obs_len + self.pred_len
        self.delim = delim
        self.cur_frame_no = 0

        all_files_train = sorted(os.listdir(self.data_dir + '/train/'))
        all_files_test = sorted(os.listdir(self.data_dir + '/test/'))
        all_files_val = sorted(os.listdir(self.data_dir + '/val/'))
        all_files_train = [os.path.join(self.data_dir + '/train/', path) for path in all_files_train]
        all_files_test = [os.path.join(self.data_dir + '/test/', path) for path in all_files_test]
        all_files_val = [os.path.join(self.data_dir + '/val/', path) for path in all_files_val]
        data_dict = {'train': all_files_train, 'test': all_files_test, 'val': all_files_val}
        
        
        for type, all_files in data_dict.items():

            ### init the frame, seq_list and num_peds_in_seq while storing the train/val/test data
            num_peds_in_seq = []
            seq_list = []
            frame_list = []
            for path in all_files:    
                data = self.read_file(path, delim)
                frames = np.unique(data[:, 0]).tolist()
                frame_data = []
                for frame in frames:
                    frame_data.append(data[frame == data[:, 0], :])
                num_sequences = int(math.ceil((len(frames) - self.seq_len + 1) / skip))

                for idx in range(0, num_sequences * self.skip + 1, skip):
                    curr_seq_data = np.concatenate(frame_data[idx:idx + self.seq_len], axis=0)
                    peds_in_curr_seq = np.unique(curr_seq_data[:, 1])
                    self.max_peds_in_frame = max(self.max_peds_in_frame, len(peds_in_curr_seq))
                    curr_seq = np.zeros((len(peds_in_curr_seq), 2, self.seq_len))

                    num_peds_considered = 0
                    for _, ped_id in enumerate(peds_in_curr_seq):
                        curr_ped_seq = curr_seq_data[curr_seq_data[:, 1] == ped_id, :]
                        curr_ped_seq = np.around(curr_ped_seq, decimals=4)
                        pad_front = frames.index(curr_ped_seq[0, 0]) - idx
                        pad_end = frames.index(curr_ped_seq[-1, 0]) - idx + 1
                        if pad_end - pad_front != self.seq_len:
                            continue
                        curr_ped_seq = np.transpose(curr_ped_seq[:, 2:])
                        _idx = num_peds_considered
                        curr_seq[_idx, :, pad_front:pad_end] = curr_ped_seq
                        num_peds_considered += 1
                    if num_peds_considered > min_ped:
                        num_peds_in_seq.append(num_peds_considered)
                        seq_list.append(curr_seq[:num_peds_considered])
                        frame_list.append(frames[idx])

            self.num_seq = len(seq_list)
            self.trajs = np.concatenate(seq_list, axis=0).transpose(0, 2, 1)
            cum_start_idx = [0] + np.cumsum(num_peds_in_seq).tolist()
            self.seq_start_end = [(start, end) for start, end in zip(cum_start_idx, cum_start_idx[1:])]
            self.frame_list = np.array(frame_list, dtype=np.int32)
            traj_data = {
                'traj': self.trajs,
                'frame_list': self.frame_list,
                'num_peds_in_seq': num_peds_in_seq,
                'seq_start_end': self.seq_start_end
            }
            pickle.dump(traj_data, open(self.data_dir + '/{:s}_{:s}.pkl'.format(subset, type), 'wb'))


    def read_file(self, _path, delim):
        delim = delim if delim else self.delim
        data = []
        with open(_path, 'r') as f:
            for line in f:
                line = line.strip().split(delim)
                line = [float(i) for i in line]
                data.append(line)
        return np.asarray(data)