├── README.md ├── __pycache__ ├── kalman_model.cpython-35.pyc ├── model.cpython-35.pyc ├── utils.cpython-35.pyc └── vanilla_gru.cpython-35.pyc ├── kalman_model.py ├── main.py ├── model.py ├── temp_.py ├── temp_rnn.py ├── utils.py └── vanilla_gru.py /README.md: -------------------------------------------------------------------------------- 1 | # trajectory-prediction-for-KalmanPrediction-and-DeepLearning 2 | There is no further plan for developing some projects from this repository, however, for studying the trajectory prediction research, this working code would be a good start. 3 | 4 | ## Requirements: 5 | * Python 3.x (I confirm that the scripts run in Python 3.5.3) 6 | * pytorch 7 | * matplotlib (for visualizing results) 8 | 9 | ## Description: 10 | This repository is for studying a trajectory prediction using Kalman filter and deep learning models. 11 | There are some confused parts on the script, but this code will be a good start for the trajectory prediction study. 12 | Curretly, below **models** are implemented for trajectory-prediction. 13 | 14 | * **Kalman Filter model (KF)**: Using prediction and correction step in KF, update covariance matrix. At the trajectory-prediction step, it uses **Constant Velocity** or **Constant Acceleration** model. 15 | * **Sequence-to-Sequence model based on GRU cell** 16 | 17 | ## Dataset: 18 | * [**Apolloscape**](http://apolloscape.auto/trajectory.html) 19 | : Apolloscape dataset is collected under various lighting conditions and traffic densities in Beijing, China. In this repo., we have used 'prediction_train.zip' file. 20 | Each line in the file contains (frame_id, object_id, object_type, position_x, position_y, position_z, object_length, object_width, object_height, heading). 21 | 22 | * [**Lyft**](https://self-driving.lyft.com/level5/data/) 23 | : Lyft data is collected under various urban conditions for 1000 hours driving. The dataset has 'zarr' format. For preprocessing data on your shoes, please check [this repository](https://github.com/zarr-developers/zarr-python). 24 | 25 | ## Further Details of This Repository. 26 | * **main.py**: main fuction for running whole scripts. You can change hyper-parameters, dataset and a model. 27 | * **model.py**: Training or testing a model 28 | * **kalman_model.py**: a basic kalman filter model for trajectory prediction. It uses the prediction only. However, as the prediction step is on-going, the covariance is getting larger as the kalman filter does. As a result, you can sample the trajectory prediction results from Kalman model. 29 | * **vanilla_gru.py**: a deep-learning model for learning the trajectory sequence. Currently, it consists of a Encoder-Decoder architecture. 30 | * **utils.py**: some utitlity functions, mainly preprocessing data. 31 | 32 | ### References. 33 | * http://apolloscape.auto/trajectory.html 34 | * https://self-driving.lyft.com/level5/data/ 35 | * https://arxiv.org/pdf/1908.11472.pdf 36 | 37 | 38 | **Please contact me (msk930512@snu.ac.kr) if you have any questions.** 39 | 40 | -------------------------------------------------------------------------------- /__pycache__/kalman_model.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BenMSK/trajectory-prediction-for-KalmanPrediction-and-DeepLearning/1505f433ddeb1cb9c7885443e128f96f79a9b2fd/__pycache__/kalman_model.cpython-35.pyc -------------------------------------------------------------------------------- /__pycache__/model.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BenMSK/trajectory-prediction-for-KalmanPrediction-and-DeepLearning/1505f433ddeb1cb9c7885443e128f96f79a9b2fd/__pycache__/model.cpython-35.pyc -------------------------------------------------------------------------------- /__pycache__/utils.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BenMSK/trajectory-prediction-for-KalmanPrediction-and-DeepLearning/1505f433ddeb1cb9c7885443e128f96f79a9b2fd/__pycache__/utils.cpython-35.pyc -------------------------------------------------------------------------------- /__pycache__/vanilla_gru.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BenMSK/trajectory-prediction-for-KalmanPrediction-and-DeepLearning/1505f433ddeb1cb9c7885443e128f96f79a9b2fd/__pycache__/vanilla_gru.cpython-35.pyc -------------------------------------------------------------------------------- /kalman_model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy.linalg import inv 3 | 4 | class KalmanModel: 5 | 6 | def __init__(self, args, acceleration=False): 7 | self.obs_len = args.obs_length 8 | self.pred_len = args.pred_length 9 | 10 | self.p0 = 15#15 11 | self.q = .03#0.03 12 | self.r = .03#0.03 13 | self.data_fps = 2 14 | self.dt = 1/self.data_fps 15 | self.A, self.H = self.init_matrices(acceleration, dt = self.dt) 16 | self.m_dim, self.n_dim = self.H.shape 17 | 18 | self.Q = np.diag(np.full(self.n_dim, self.q))# Process Noise 19 | self.R = np.diag(np.full(self.m_dim, self.r))# Observation Noise 20 | self.P0 = np.diag(np.full(self.n_dim, self.p0))# Covariance matrix 21 | self.K = np.zeros((self.m_dim, self.n_dim)) 22 | self.P = None 23 | 24 | def init_matrices(self, acceleration, dt=0.1): 25 | if acceleration: # use acceleration 26 | # transition matrix x x' y y' x'' y'' 27 | A = np.array([[1, 1 * dt, 0, 0, 0.5 * dt * dt, 0], # x 28 | [0, 1, 0, 0, 1 * dt, 0], # x' 29 | [0, 0, 1, 1 * dt, 0, 0.5 * dt * dt], # y 30 | [0, 0, 0, 1, 0, 1 * dt], # y' 31 | [0, 0, 0, 0, 1, 0], # x'' 32 | [0, 0, 0, 0, 0, 1]]) # y'' 33 | # 6x6 34 | 35 | H = np.array([[1, 0, 0, 0, 0, 0], 36 | [0, 0, 1, 0, 0, 0]])# 2x6 37 | else: 38 | # transition matrix x x' y y' 39 | A = np.array([[1, 1 * dt, 0, 0], 40 | [0, 1, 0, 0], 41 | [0, 0, 1, 1 * dt], 42 | [0, 0, 0, 1]]) 43 | 44 | H = np.array([[1, 0, 0, 0], # m x n m = 2, n = 4 or 6 45 | [0, 0, 1, 0]]) 46 | return A, H 47 | 48 | def PredictTraj(self, trajectory): 49 | ''' 50 | Kalman 51 | : The whole system consists of calculating the filter parameter through n observation data 52 | Then, using that parameter, it is going to estimate the k frames trajectory! 53 | ''' 54 | 55 | x_hat_previous = None 56 | 57 | # Get agent's 58 | # agents_obs_data: Contains [x,y,type] of each agent in each history' frame. 59 | # the last frame is 'frame_idx' which means the current frame 60 | pred_trajectory = np.empty((self.pred_len, self.n_dim))#n_dim could be (x, x', y, y') or (x, x', y, y', x'', y'') 61 | self.K = np.zeros((self.m_dim, self.n_dim)) 62 | self.P = self.P0 63 | 64 | if len(trajectory) <= 1: 65 | return None 66 | 67 | # State initialization. 68 | x0, y0 = trajectory[0] 69 | 70 | v0x = (trajectory[1, 0] - x0) / self.dt#velocity 71 | v0y = (trajectory[1, 1] - y0) / self.dt 72 | if self.n_dim == 6: 73 | x_hat_previous = np.array([x0, v0x, y0, v0y, 0, 0]) 74 | elif self.n_dim == 4: 75 | x_hat_previous = np.array([x0, v0x, y0, v0y]) 76 | covariance_list = [None]*self.pred_len 77 | 78 | # Take the first n frames and fit the kalman filter on them. 79 | for j, z_k in enumerate(trajectory):#range(len(trajectory)): 80 | x_hat_new = self.fit(x_hat_previous, z_k) # predict and correct 81 | x_hat_previous = x_hat_new 82 | 83 | for u in range(self.pred_len): 84 | x_hat_new = self.predict(x_hat_previous) 85 | covariance_list[u] = self.P 86 | pred_trajectory[u] = x_hat_new 87 | x_hat_previous = x_hat_new 88 | 89 | return pred_trajectory, covariance_list 90 | 91 | def fit(self, xhat_previous, z_k): 92 | """ 93 | main iteration: we need the state estimate x_hat_minus k-1 at previous step and the current measurement z_k 94 | 95 | :param xhat_previous: previous a posteriori prediction 96 | :param z_k: current measurement: (tx,ty) tuple 97 | :return: new a posteriori prediction 98 | """ 99 | 100 | # prediction 101 | xhat_k_minus = self.predict(xhat_previous) # predict updates self.P (P_minus) 102 | P_minus = self.P 103 | 104 | # innovation 105 | residual = z_k - np.dot(self.H, xhat_k_minus) 106 | inv_HPHTR = inv(np.dot(np.dot(self.H, P_minus), self.H.T) + self.R) 107 | 108 | # correction (update) 109 | self.K = np.dot(np.dot(P_minus, self.H.T), inv_HPHTR)# Calculate Kalman gain 110 | xhat_k_new = xhat_k_minus + np.dot(self.K, residual) 111 | self.P = np.dot((np.eye(self.n_dim) - np.dot(self.K, self.H)), P_minus) 112 | 113 | return xhat_k_new 114 | 115 | def predict(self, xhat_k_previous): 116 | xhat_k_minus = np.dot(self.A, xhat_k_previous) # update previous state estimate with state transition matrix 117 | self.P = np.dot(np.dot(self.A, self.P), self.A.T) + self.Q # this is P minus; Covariacne 118 | 119 | return xhat_k_minus 120 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import argparse 3 | import os 4 | # import pickle 5 | # import time 6 | 7 | import torch 8 | 9 | from utils import DataSet, PreprocessData 10 | from torch.utils.data import DataLoader 11 | from vanilla_gru import NNPred 12 | from model import VanillaGRU 13 | from kalman_model import KalmanModel 14 | import matplotlib.pyplot as plt 15 | 16 | CUDA = True 17 | DEVICE = 'cuda:0' 18 | TRAIN = False 19 | TEST = True 20 | DATASET = "Apol"# "Apol", "Lyft" 21 | 22 | def main(): 23 | parser = argparse.ArgumentParser() 24 | 25 | # Set a torch manual seed 26 | parser.add_argument( 27 | "--seed", 28 | type=int, 29 | default=42, 30 | help="Random seed of torch" 31 | ) 32 | 33 | parser.add_argument( 34 | "--wandb", 35 | type=bool, 36 | default=False, 37 | help="Wandb trigger" 38 | ) 39 | 40 | parser.add_argument( 41 | '--cuda', '-g', 42 | action='store_true', 43 | help='GPU option', 44 | default=CUDA 45 | ) 46 | 47 | parser.add_argument( 48 | '--device', '-d', 49 | help='cuda device option', 50 | default=DEVICE, 51 | type=str 52 | ) 53 | 54 | parser.add_argument( 55 | "--raw_data_dir", 56 | type=str, 57 | default='../data/{}/'.format(DATASET), 58 | help="Data Directory" 59 | ) 60 | 61 | parser.add_argument( 62 | "--obs_length", 63 | type=int, 64 | default=6, 65 | help="History length of the trajectory" 66 | ) 67 | 68 | parser.add_argument( 69 | "--pred_length", 70 | type=int, 71 | default=10, 72 | help="predicted length of the trajectory" 73 | ) 74 | 75 | parser.add_argument( 76 | "--batch_size", 77 | type=int, 78 | default=32, 79 | help="mini-batch size of Data" 80 | ) 81 | 82 | parser.add_argument( 83 | "--train", 84 | type=bool, 85 | default=TRAIN, 86 | help="Activate the train mode. If not, the test mode is activated" 87 | ) 88 | 89 | parser.add_argument( 90 | "--test", 91 | type=bool, 92 | default=TEST, 93 | help="Activate the train mode. If not, the test mode is activated" 94 | ) 95 | 96 | parser.add_argument( 97 | "--epoch", 98 | type=int, 99 | default=100, 100 | help="Epoch!" 101 | ) 102 | 103 | print("Using {} dataset....".format(DATASET)) 104 | 105 | args = parser.parse_args() 106 | args.dataset_name = DATASET 107 | 108 | ## if you are using 10 Hz data and want to divide the frame to 2Hz, you have to 10/2 = 5// only int is considered, here. 109 | args.frame_interval = 1# skip rate 110 | args.train_val_test_ratio = (0.7, 0.2, 0.1) 111 | args.data_file = os.path.join("../data/", "{}-dataset-seed{}.cpkl".format(DATASET, args.seed))#pickle contains 'train', 'val', 'test' file name and 'scale param' 112 | # args.data_file = os.path.join("./data/", "dataset-seed{}.not.gaussian.cpkl".format(args.seed)) 113 | args.load_name = "none" 114 | if not os.path.exists(args.data_file): 115 | PreprocessData(args) 116 | 117 | args.save_name = "{}-{}-{}.S2S_GRUmodel.seed{}.tar".format(DATASET, args.obs_length, args.pred_length, args.seed) 118 | args.load_name = "21.Apol-6-10.S2S_GRUmodel.seed42.tar".format(args.seed) 119 | 120 | model = VanillaGRU(args) 121 | # for i in range(12300, 12500): 122 | # hist, hist_mask, fut, fut_mask, ref_pose, AgentInfo = model.test_dataset[i] 123 | # kalman_model = KalmanModel(args) 124 | # pred, _ = kalman_model.PredictTraj(hist) 125 | # print("pred: ", pred) 126 | # if AgentInfo[-1] == 3: 127 | # plt.plot(hist[:,0], hist[:,1], "ko-", alpha=0.8) 128 | # plt.plot(fut[:,0], fut[:,1], "g^-") 129 | # plt.plot(pred[:len(fut),0], pred[:len(fut),2], "r^-") 130 | # plt.axis('equal') 131 | 132 | # plt.show() 133 | 134 | if args.train: 135 | model.train() 136 | 137 | if args.test: 138 | model.load() 139 | model.test() 140 | 141 | 142 | if __name__ == "__main__": 143 | main() 144 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch import optim 4 | from vanilla_gru import NNPred 5 | from utils import DataSet, maskedMSE, maskedLastPositionLoss, class_objtype 6 | from torch.utils.data import DataLoader 7 | from torch.autograd import gradcheck 8 | 9 | import os 10 | import matplotlib.pyplot as plt 11 | import numpy as np 12 | import time 13 | import wandb 14 | import colorama 15 | from colorama import Fore, Style 16 | 17 | class VanillaGRU: 18 | def __init__(self, args): 19 | 20 | self.args = {} 21 | self.args['dataset'] = args.dataset_name 22 | self.args["batch_size"] = args.batch_size 23 | self.args['cuda'] = args.cuda 24 | self.args['device'] = args.device 25 | self.args['epoch'] = args.epoch 26 | 27 | self.args['in_length'] = args.obs_length 28 | self.args['out_length'] = args.pred_length 29 | self.args['save_name'] = args.save_name 30 | self.args['load_name'] = args.load_name 31 | 32 | # self.args['nll_only'] = True 33 | self.args["learning_rate"] = 1e-4 34 | self.args["w_decay"] = 1e-4 35 | 36 | self.args['name'] = 'test.tar' 37 | self.args["optim"] = 'Adam' 38 | self.args['train_loss'] = 'MSE' 39 | 40 | self.wandb = False 41 | if args.wandb: 42 | self.wandb = True 43 | print("Wandb is initialized...") 44 | wandb.init(project="vanilla_gru",\ 45 | # name="obs: {}, pred: {}".format(args.obs_length, args.pred_length), 46 | config=self.args) 47 | 48 | self.net = NNPred(args) 49 | if self.args['cuda']: 50 | self.net = self.net.cuda(self.args['device']) 51 | 52 | # for training// dataset 53 | self.train_dataset = DataSet(args, 'train') 54 | self.val_dataset = DataSet(args, 'val') 55 | self.test_dataset = DataSet(args, 'test') 56 | # for i in range(12300, 12500): 57 | # hist, hist_mask, fut, fut_mask, ref_pose, AgentInfo = self.test_dataset[0] 58 | # print(AgentInfo) 59 | # print(hist) 60 | # print(fut) 61 | # # A = fut+ref_pose#// These lines for 62 | # # A[:,0] = ((A[:,0]+1)/2)*(self.test_dataset.max_position_x - self.test_dataset.min_position_x)+ self.test_dataset.min_position_x 63 | # # A[:,1] = ((A[:,1]+1)/2)*(self.test_dataset.max_position_y - self.test_dataset.min_position_y)+ self.test_dataset.min_position_y 64 | # print(A) 65 | # # print(fut[:,1]+ref_pose[1]) 66 | # quit() 67 | 68 | self.trainDataloader = DataLoader(self.train_dataset, batch_size=self.args["batch_size"], shuffle=True, \ 69 | num_workers=6, collate_fn=self.train_dataset.GetBatch, drop_last=True) 70 | # print("trainDataloader completed!") 71 | self.valDataloader = DataLoader(self.val_dataset, batch_size=self.args["batch_size"], shuffle=True, \ 72 | num_workers=6, collate_fn=self.val_dataset.GetBatch, drop_last=True) 73 | # print("valDataloader completed!") 74 | self.testDataloader = DataLoader(self.test_dataset, batch_size=self.args["batch_size"], shuffle=True, \ 75 | num_workers=6, collate_fn=self.test_dataset.GetBatch, drop_last=True) 76 | # print("testDataloader completed!") 77 | # time.sleep(300) 78 | 79 | def train(self): 80 | 81 | total_epoch = self.args['epoch'] 82 | avg_trn_loss = 0 83 | avg_val_loss = 0 84 | if self.args["optim"] == 'Adam': 85 | optim = torch.optim.Adam(self.net.parameters(),lr=self.args['learning_rate'],weight_decay=self.args["w_decay"]) 86 | else: 87 | print("Undefined optimizer.") 88 | return 89 | criterion = torch.nn.MSELoss() 90 | print("Start training....") 91 | for epoch in range(total_epoch): 92 | tr_count = 0 93 | val_count = 0 94 | self.net.train_flag = True 95 | # for data in self.testDataloader:#for check 96 | for data in self.trainDataloader:# TRAINING PHASE 97 | # print("count! ", tr_count) 98 | hist_batch, fut_batch, fut_mask_batch, _, _ = data 99 | if self.args['cuda']: 100 | hist_batch = hist_batch.cuda(self.args['device']) 101 | fut_batch = fut_batch.cuda(self.args['device']) 102 | fut_mask_batch = fut_mask_batch.cuda(self.args['device']) 103 | 104 | fut_pred = self.net(hist_batch, fut_batch) 105 | # fut_mask_batch = torch.cat((fut_mask_batch, fut_mask_batch),2)#if network output has sigX 106 | 107 | fut_pred = fut_mask_batch*fut_pred 108 | 109 | loss = criterion(fut_pred, fut_batch) 110 | # loss = maskedMSE(fut_pred, fut_batch, fut_mask_batch, device=self.args['device']) 111 | # loss = maskedLastPositionLoss(fut_pred, fut_batch, fut_mask_batch, device=self.args['device']) 112 | optim.zero_grad() 113 | loss.backward() 114 | 115 | # nn.utils.clip_grad_norm_(self.net.parameters(), 1) 116 | optim.step() 117 | 118 | avg_trn_loss += loss.item() 119 | tr_count+=1 120 | if self.wandb and self.args['dataset']=='Lyft': 121 | # print("logging!") 122 | wandb.log({'Avg Train Loss per step': loss.item()}) 123 | 124 | avg_trn_loss /= tr_count 125 | print("Epoch: {} AvgTrainLoss: {}".format(epoch, avg_trn_loss)) 126 | 127 | ############################################################################################ 128 | self.net.train_flag = False 129 | for data in self.valDataloader:# VALIDATION PHASE 130 | hist_batch, fut_batch, fut_mask_batch, _, _ = data 131 | if self.args['cuda']: 132 | hist_batch = hist_batch.cuda(self.args['device']) 133 | fut_batch = fut_batch.cuda(self.args['device']) 134 | fut_mask_batch = fut_mask_batch.cuda(self.args['device']) 135 | 136 | fut_pred = self.net(hist_batch, fut_batch) 137 | 138 | # fut_mask_batch = torch.cat((fut_mask_batch, fut_mask_batch),2)#if network output has sigX 139 | 140 | fut_pred = fut_mask_batch*fut_pred 141 | 142 | # loss = maskedMSE(fut_pred, fut_batch, fut_mask_batch, device=self.args['device']) 143 | # loss = maskedLastPositionLoss(fut_pred, fut_batch, fut_mask_batch, device=self.args['device']) 144 | loss = criterion(fut_pred, fut_batch) 145 | 146 | avg_val_loss += loss.item() 147 | val_count+=1 148 | if self.wandb and self.args['dataset']=='Lyft': 149 | wandb.log({'Avg Val Loss per step': loss.item()}) 150 | 151 | avg_val_loss /= val_count 152 | print("Epoch: {} AvgValLoss: {}".format(epoch, avg_val_loss)) 153 | print("===================================================") 154 | if self.wandb: 155 | wandb.log({'Avg Train Loss per Epoch': avg_trn_loss, 'Avg Val Loss per Epoch': avg_val_loss}) 156 | self.saveModel(epoch) 157 | 158 | def saveModel(self, epoch): 159 | save_dir = os.path.join("./weight/", str(epoch)+"."+self.args['save_name']) 160 | # name = os.path.join(self.args['modelLoc'], "epochs.{}.".format(engine.state.epoch)+self.args['name']) 161 | torch.save(self.net.state_dict(), save_dir) 162 | print("Model saved {}.".format(save_dir)) 163 | 164 | 165 | def load(self): 166 | load_dir = os.path.join("./weight/", self.args['load_name']) 167 | if os.path.exists(load_dir): 168 | self.net.load_state_dict(torch.load(load_dir)) 169 | print(Fore.YELLOW) 170 | print("\n[INFO]: model {} loaded, successfully!\n".format(load_dir)) 171 | print(Style.RESET_ALL) 172 | 173 | else: 174 | print(Fore.RED) 175 | print("\n[INFO]: CAN NOT FIND MODEL AT {}".format(load_dir)) 176 | print(Style.RESET_ALL) 177 | 178 | def test(self): 179 | self.net.train_flag = False 180 | # (m_x, m_y, sigX, sigY) = self.test_dataset.scale_param# for gaussian scailing 181 | (min_position_x, max_position_x, min_position_y, max_position_y) = self.test_dataset.scale_param 182 | for data in self.testDataloader: 183 | hist_batch, fut_batch, fut_mask_batch, ref_pose_batch, AgentsInfo = data 184 | if self.args['cuda']: 185 | hist_batch = hist_batch.cuda(self.args['device']) 186 | 187 | fut_pred_batch = self.net(hist_batch, fut_batch) 188 | 189 | # print(fut_pred_batch.shape) 190 | hist_batch = hist_batch.cpu().detach().numpy() 191 | fut_pred_batch = fut_pred_batch[:,:,:2].cpu().detach().numpy() 192 | fut_batch = fut_batch.numpy() 193 | fut_mask_batch = fut_mask_batch.numpy() 194 | ref_pose_batch = ref_pose_batch.numpy() 195 | 196 | for ind, agent_traj in enumerate(fut_batch): 197 | print("dataset: {}, AgentID: {}, AgentType: {}"\ 198 | .format(AgentsInfo[ind][0], AgentsInfo[ind][2], class_objtype(AgentsInfo[ind][3], self.args['dataset']))) 199 | 200 | # masking point helps for cutting non-enough gt case when visualizing 201 | masking_point = np.where(fut_mask_batch[ind]==0)[0][0] if np.where(fut_mask_batch[ind]==0)[0].shape[0] != 0 else len(fut_mask_batch[ind]) 202 | 203 | agent_traj += ref_pose_batch[ind] 204 | fut_pred_batch[ind] += ref_pose_batch[ind] 205 | hist_batch[ind] += ref_pose_batch[ind] 206 | hist_x = ((hist_batch[ind][:,0]+1)/2)*(max_position_x - min_position_x) + min_position_x 207 | hist_y = ((hist_batch[ind][:,1]+1)/2)*(max_position_y - min_position_y) + min_position_y 208 | 209 | gt_x = ((fut_batch[ind][:masking_point,0]+1)/2)*(max_position_x - min_position_x) + min_position_x 210 | gt_y = ((fut_batch[ind][:masking_point,1]+1)/2)*(max_position_y - min_position_y) + min_position_y 211 | # print(gt_x[0], gt_y[0]) 212 | # print(len(gt_x)) 213 | pred_x = ((fut_pred_batch[ind][:masking_point,0]+1)/2)*(max_position_x - min_position_x) + min_position_x 214 | pred_y = ((fut_pred_batch[ind][:masking_point,1]+1)/2)*(max_position_y - min_position_y) + min_position_y 215 | 216 | plt.plot(hist_x, hist_y, "ko-", alpha=0.5) 217 | plt.plot(hist_x[-1], hist_y[-1], "yo", alpha=0.8) 218 | plt.plot(gt_x, gt_y, "g^-") 219 | plt.plot(pred_x, pred_y, "r^-") 220 | plt.axis('equal') 221 | 222 | plt.show() 223 | 224 | 225 | 226 | 227 | -------------------------------------------------------------------------------- /temp_.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | # this script is for generating track file for Lyft dataset 3 | # dataset = os.listdir(raw_data_dir) 4 | 5 | # for ind, raw_file_name in enumerate(dataset): 6 | # file_path = os.path.join(raw_data_dir, raw_file_name)# Each .txt or .csv file 7 | # read = np.load(file_path)#Lyft 8 | 9 | # agents = np.unique(read[:,id_ind]) 10 | # track_data = None# Shape: agent x (frame, x, y) 11 | # for agent in agents: 12 | # a_track = read[read[:,id_ind]==agent][:,[id_ind, x_ind, y_ind]] 13 | # track_data[agent] = a_track 14 | # track_data_list[dataset_name] = track_data 15 | frame_idx = 20 16 | interval = 6 17 | obs_length = 6 18 | a = np.arange(100) 19 | # print(frame_idx+1 - obs_length*(interval)+(interval-1)) 20 | obs_length = int(frame_idx/interval + 1) if frame_idx - (obs_length-1)*(interval) < 0 else obs_length 21 | # print("minus!: ", obs_length) 22 | start_idx = np.maximum(0, frame_idx - (obs_length-1)*(interval)) 23 | end_idx = frame_idx + 1 24 | print(a[start_idx:end_idx:interval]) 25 | 26 | start = frame_idx-len(a) - (obs_length-1) 27 | end = frame_idx-len(a)+1 28 | print(a[start:end:-interval]) 29 | -------------------------------------------------------------------------------- /temp_rnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.optim as optim 4 | import numpy as np 5 | 6 | input_str = 'apple' 7 | label_str = 'pple!' 8 | char_vocab = sorted(list(set(input_str+label_str))) 9 | vocab_size = len(char_vocab) 10 | print ('문자 집합의 크기 : {}'.format(vocab_size)) 11 | 12 | input_size = vocab_size # 입력의 크기는 문자 집합의 크기 13 | hidden_size = 5 14 | output_size = 5 15 | learning_rate = 0.1 16 | 17 | char_to_index = dict((c, i) for i, c in enumerate(char_vocab)) # 문자에 고유한 정수 인덱스 부여 18 | print(char_to_index) 19 | 20 | index_to_char={} 21 | for key, value in char_to_index.items(): 22 | index_to_char[value] = key 23 | print(index_to_char) 24 | 25 | x_data = [char_to_index[c] for c in input_str] 26 | y_data = [char_to_index[c] for c in label_str] 27 | print(x_data) 28 | print(y_data) 29 | 30 | x_data = [x_data] 31 | y_data = [y_data] 32 | 33 | x_one_hot = [np.eye(vocab_size)[x] for x in x_data] 34 | print(x_one_hot) 35 | 36 | X = torch.FloatTensor(x_one_hot) 37 | Y = torch.LongTensor(y_data) 38 | 39 | class Net(torch.nn.Module): 40 | def __init__(self, input_size, hidden_size, output_size): 41 | super(Net, self).__init__() 42 | self.rnn = torch.nn.RNN(input_size, hidden_size, batch_first=True) 43 | self.fc = torch.nn.Linear(hidden_size, output_size, bias=True) 44 | 45 | def forward(self, x): 46 | x, _status = self.rnn(x) 47 | x = self.fc(x) 48 | return x 49 | 50 | net = Net(input_size, hidden_size, output_size) 51 | outputs = net(X) 52 | print(outputs.shape) 53 | 54 | criterion = torch.nn.CrossEntropyLoss() 55 | optimizer = optim.Adam(net.parameters(), learning_rate) 56 | for i in range(100): 57 | optimizer.zero_grad() 58 | outputs = net(X) 59 | loss = criterion(outputs.view(-1, input_size), Y.view(-1)) 60 | loss.backward() # 기울기 계산 61 | # print(loss.grad) 62 | optimizer.step() # 아까 optimizer 선언 시 넣어둔 파라미터 업데이트 63 | 64 | # 아래 세 줄은 모델이 실제 어떻게 예측했는지를 확인하기 위한 코드. 65 | result = outputs.data.numpy().argmax(axis=2) # 최종 예측값인 각 time-step 별 5차원 벡터에 대해서 가장 높은 값의 인덱스를 선택 66 | result_str = ''.join([index_to_char[c] for c in np.squeeze(result)]) 67 | print(i, "loss: ", loss.item(), "prediction: ", result, "true Y: ", y_data, "prediction str: ", result_str) -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import logging 3 | import os 4 | import pickle 5 | import random 6 | import torch 7 | import numpy as np 8 | import time 9 | import matplotlib.pyplot as plt 10 | 11 | from torch.utils.data import Dataset 12 | 13 | # torch.autograd.set_detect_anomaly(True)# for detecting abnormal 14 | 15 | def class_objtype(object_type, dataset='Apol'):# APOL 16 | if dataset == 'Apol': 17 | if object_type == 1 or object_type == 2:#Vehicle 18 | return 'Vehicle' 19 | elif object_type == 3:#Pedestrian 20 | return 'Pedestrian' 21 | elif object_type == 4:#Bicycle 22 | return 'Bicycle' 23 | else: 24 | return 'Unknown' 25 | 26 | def GetDatasetIndInfo(dataset): 27 | #return f_ind, id_ind, x_ind, y_ind, yaw_ind, type_ind 28 | if dataset=='Apol': 29 | return 0, 1, 3, 4, 5, 2 30 | elif dataset=='Lyft': 31 | return 0, 1, 2, 3, 4, 5 32 | else: 33 | return 0, 1, 3, 4, 5, 2 34 | 35 | def ExtractData(raw_data_dir, data_files, dataset): 36 | 37 | full_data_list = np.array([]) 38 | track_data_list = {} 39 | min_position_x = 1000 40 | max_position_x = -1000 41 | min_position_y = 1000 42 | max_position_y = -1000 43 | 44 | gaussian_scaling = False 45 | x_gather = list() 46 | y_gather = list() 47 | # Generate datasetf 48 | print(data_files) 49 | f_ind, id_ind, x_ind, y_ind, yaw_ind, type_ind = GetDatasetIndInfo(dataset) 50 | 51 | for ind_directory, raw_file_name in enumerate(data_files): 52 | 53 | # for Lyft 54 | file_path = os.path.join(raw_data_dir, raw_file_name)# Each .txt or .csv file 55 | tmp0 = file_path.split('/')[-1].split('_')[1] 56 | tmp1 = file_path.split('/')[-1].split('_')[2] 57 | dataset_name = int(tmp0+tmp1.zfill(2))#only for Apolloscape dataset 58 | 59 | read = np.loadtxt(file_path, delimiter=' ')#APOL 60 | # read = np.load('file_path')#Lyft 61 | min_position_x = min(min_position_x, min(read[:, x_ind])) 62 | max_position_x = max(max_position_x, max(read[:, x_ind])) 63 | min_position_y = min(min_position_y, min(read[:, y_ind])) 64 | max_position_y = max(max_position_y, max(read[:, y_ind])) 65 | 66 | if gaussian_scaling: 67 | x_gather += read[:, x_ind].tolist() 68 | y_gather += read[:, y_ind].tolist() 69 | relevant_data = read[:, [f_ind, id_ind, type_ind, x_ind, y_ind, yaw_ind]]#(frame, id, type, x, y, heading) 70 | dataset_name_list = np.full([relevant_data.shape[0], 1], dataset_name) 71 | # (dsId, frame, agnetId, type, x, y, heading) 72 | relevant_data = np.concatenate((dataset_name_list, relevant_data), axis=1) 73 | if full_data_list.size == 0: 74 | full_data_list = relevant_data 75 | else: 76 | full_data_list = np.concatenate((full_data_list, relevant_data), axis=0) 77 | 78 | agents = np.unique(read[:,id_ind]) 79 | track_data = {}# Shape: agent x (frame, x, y) 80 | for agent in agents: 81 | a_track = read[read[:,id_ind]==agent][:,[id_ind, x_ind, y_ind]] 82 | track_data[agent] = a_track 83 | track_data_list[dataset_name] = track_data 84 | 85 | # Scaling! 86 | if gaussian_scaling:# Scale 'standard normal distribution' 87 | x_gather = np.array(x_gather) 88 | y_gather = np.array(y_gather) 89 | mean_x = np.mean(x_gather) 90 | mean_y = np.mean(y_gather) 91 | std_x = np.std(x_gather) 92 | std_y = np.std(y_gather) 93 | 94 | scale_param = (min_position_x, max_position_x, min_position_y, max_position_y) 95 | full_data_list[:, 4] = (full_data_list[:, 4] - mean_x) / std_x 96 | full_data_list[:, 5] = (full_data_list[:, 5] - mean_y) / std_y 97 | 98 | for dataset_name in track_data_list.keys(): 99 | for agent in track_data_list[dataset_name].keys(): 100 | track_data_list[dataset_name][agent][:, 1] = (track_data_list[dataset_name][agent][:, 1] - mean_x) / std_x 101 | track_data_list[dataset_name][agent][:, 2] = (track_data_list[dataset_name][agent][:, 2] - mean_y) / std_y 102 | scale_param = (mean_x, std_x, mean_y, std_y) 103 | 104 | else:# Scale range [-1, 1] 105 | full_data_list[:, 4] = ( 106 | (full_data_list[:, 4] - min_position_x) / (max_position_x - min_position_x) 107 | ) * 2 - 1 108 | full_data_list[:, 5] = ( 109 | (full_data_list[:, 5] - min_position_y) / (max_position_y - min_position_y) 110 | ) * 2 - 1 111 | 112 | for dataset_name in track_data_list.keys(): 113 | for agent in track_data_list[dataset_name].keys(): 114 | track_data_list[dataset_name][agent][:, 1] = ( 115 | (track_data_list[dataset_name][agent][:, 1] - min_position_x) / (max_position_x - min_position_x) 116 | ) * 2 - 1 117 | track_data_list[dataset_name][agent][:, 2] = ( 118 | (track_data_list[dataset_name][agent][:, 2] - min_position_y) / (max_position_y - min_position_y) 119 | ) * 2 - 1 120 | scale_param = (min_position_x, max_position_x, min_position_y, max_position_y) 121 | 122 | return (full_data_list, track_data_list, scale_param) 123 | 124 | def GetScaleParam(raw_data_dir, data_files, dataset_name): 125 | min_position_x = 10000 126 | max_position_x = -10000 127 | min_position_y = 10000 128 | max_position_y = -10000 129 | 130 | # Generate datasetf 131 | f_ind, id_ind, x_ind, y_ind, yaw_ind, type_ind = GetDatasetIndInfo(dataset_name) 132 | 133 | for ind_directory, raw_file_name in enumerate(data_files): 134 | file_path = os.path.join(raw_data_dir, raw_file_name)# Each .txt or .csv file 135 | read = np.load(file_path) if dataset_name=='Lyft' else np.genfromtxt(file_path)#Apol 136 | 137 | min_position_x = min(min_position_x, min(read[:, x_ind])) 138 | max_position_x = max(max_position_x, max(read[:, x_ind])) 139 | min_position_y = min(min_position_y, min(read[:, y_ind])) 140 | max_position_y = max(max_position_y, max(read[:, y_ind])) 141 | 142 | return (min_position_x, max_position_x, min_position_y, max_position_y) 143 | 144 | def PreprocessData(args): 145 | 146 | random.seed(args.seed) 147 | np.random.seed(args.seed) 148 | 149 | tr_ratio, val_ratio, te_ratio = args.train_val_test_ratio 150 | train_fraction = tr_ratio + val_ratio 151 | val_fraction = val_ratio 152 | print("Data fraction is (train: {}%, val: {}%, test: {}%)"\ 153 | .format(tr_ratio*100, val_ratio*100, te_ratio*100)) 154 | 155 | # List of data directories where raw data resides 156 | raw_data_dir = args.raw_data_dir 157 | dataset_cnt = len(os.listdir(raw_data_dir)) 158 | dataset_idx = sorted(os.listdir(raw_data_dir)) 159 | np.random.shuffle(dataset_idx)# By random seed. 160 | 161 | # Divide the datasets to {train, val, test} 162 | data_dir_train = dataset_idx[: int(dataset_cnt * tr_ratio)] 163 | data_dir_val = dataset_idx[int(dataset_cnt * tr_ratio): int(dataset_cnt * train_fraction)] 164 | data_dir_test = dataset_idx[int(dataset_cnt * train_fraction) :] 165 | train_scale_param = GetScaleParam(raw_data_dir, data_dir_train, args.dataset_name) 166 | val_scale_param = GetScaleParam(raw_data_dir, data_dir_val, args.dataset_name) 167 | test_scale_param = GetScaleParam(raw_data_dir, data_dir_test, args.dataset_name) 168 | 169 | # Save dataset path corresponding that seed as the pickle file 170 | f = open(args.data_file, "wb") 171 | pickle.dump( 172 | ((data_dir_train,train_scale_param), (data_dir_val, val_scale_param), (data_dir_test, test_scale_param)), 173 | f, 174 | protocol=2, 175 | ) 176 | f.close() 177 | 178 | # Dataset 상속 179 | ''' 180 | For Lyft dataset:: 181 | Data config. of each row --> [frame_idx, agent_id, x, y, yaw, agent_type] 182 | We divde the whole dataset into 20 minute each. (Total driving time is 112 hours at train_0 dataset) 183 | Using agent_type: 184 | | 1: PEDESTRIAN 185 | | 2: BICYCLE, MOTORCYCLE, CYCLIST, MOTORCYCLIST 186 | | 3: CAR, VAN, TRAM, OTHER_VEHICLE 187 | | 4: BUS, TRUCK, EMERGENCY_VEHICLE 188 | | 5: UNKNOWN 189 | ''' 190 | class DataSet(Dataset): 191 | def __init__(self, args, dtype): 192 | 193 | # Store the arguments 194 | self.batch_size = args.batch_size 195 | self.obs_length = args.obs_length 196 | self.pred_length = args.pred_length 197 | self.use_cuda = args.cuda 198 | self.device = args.device 199 | 200 | self.dtype = dtype 201 | self.raw_data_dir = args.raw_data_dir 202 | self.total_data_num = 0# the total number of data 203 | self.dataset_len_list = None# cumulative data length for all data 204 | self.args = args 205 | 206 | self.min_position_x, self.max_position_x, self.min_position_y, self.max_position_y = None, None, None, None 207 | self.load_pickle(args.data_file) 208 | os_path = os.path.join("../data/", "{}-dataset-seed{}-{}_index.npy".format(self.args.dataset_name, self.args.seed, self.dtype))#pickle contains 'train', 'val', 'test' file name and 'scale param' 209 | if not os.path.exists(os_path): 210 | self.generate_dataset_ind() 211 | self.dataset_len_list = np.load(os_path).astype(int) 212 | self.total_data_num = self.dataset_len_list[-1] 213 | 214 | self.f_ind, self.id_ind, self.x_ind, self.y_ind, self.yaw_ind, self.type_ind = GetDatasetIndInfo(args.dataset_name) 215 | 216 | self.current_dataset = None 217 | self.current_data = None 218 | self.f_interval = args.frame_interval 219 | 220 | 221 | def load_pickle(self, data_dir): 222 | f = open(data_dir, "rb") 223 | self.raw_data_path = pickle.load(f) 224 | f.close() 225 | # Get 'dtype' data from the pickle file 226 | if self.dtype == 'train': 227 | self.full_data_path, self.scale_param = self.raw_data_path[0] 228 | elif self.dtype == 'val': 229 | self.full_data_path, self.scale_param = self.raw_data_path[1] 230 | else:# test 231 | self.full_data_path, self.scale_param = self.raw_data_path[2] 232 | self.min_position_x, self.max_position_x, self.min_position_y, self.max_position_y = self.scale_param 233 | print("Load the dataset... (total#: {})".format(len(self.full_data_path))) 234 | 235 | def generate_dataset_ind(self): 236 | self.dataset_len_list = np.zeros(len(self.full_data_path)) 237 | for ind, raw_file_name in enumerate(self.full_data_path): 238 | file_path = os.path.join(self.raw_data_dir, raw_file_name)# Each .txt or .csv file 239 | read = np.load(file_path) if self.args.dataset_name == 'Lyft' else np.genfromtxt(file_path)#Lyft 240 | self.total_data_num += len(read) 241 | self.dataset_len_list[ind] = self.total_data_num 242 | 243 | with open("../data/{}-dataset-seed{}-{}_index.npy".format(self.args.dataset_name, self.args.seed, self.dtype), 'wb') as f: 244 | np.save(f, self.dataset_len_list) 245 | print("Successfully, generate ind Info. and save the data") 246 | print("# of data: ", self.total_data_num) 247 | 248 | 249 | def GetDatasetFile(self, idx): 250 | prev_dataset_len = 0 251 | dataset_ind = 0 252 | data_ind = 0 253 | for ind, dataset_len in enumerate(self.dataset_len_list): 254 | if idx >= dataset_len: 255 | prev_dataset_len = dataset_len 256 | continue 257 | else: 258 | dataset_ind = ind 259 | data_ind = idx - prev_dataset_len 260 | break 261 | dataset_path = os.path.join(self.raw_data_dir, self.full_data_path[dataset_ind]) 262 | 263 | # get current dataset 264 | self.current_dataset = np.load(dataset_path) if self.args.dataset_name == 'Lyft' else np.genfromtxt(dataset_path)#Apol 265 | self.normalize_data()# Normalized a position in a range of [-1, +1] 266 | 267 | self.current_data = self.current_dataset[data_ind] 268 | 269 | return dataset_path, data_ind 270 | 271 | # 총 데이터의 개수를 리턴 (여기서는 each frame of each agent가 하나의 data point) 272 | def __len__(self): 273 | return self.total_data_num 274 | 275 | # 인덱스를 입력받아 그에 맵핑되는 입출력 데이터를 파이토치의 Tensor 형태로 리턴 276 | def __getitem__(self, idx): 277 | dataset_path, data_ind = self.GetDatasetFile(idx) 278 | 279 | dsId = dataset_path# dataset Id 280 | frame = self.current_data[self.f_ind].astype(int)# frame in the dataset Id 281 | agentId = self.current_data[self.id_ind].astype(int)# unique agentID in the dataset 282 | agentType = self.current_data[self.type_ind].astype(int) 283 | # pose = self.current_data[[self.x_ind, self.y_ind, self.yaw_ind]]# [x, y, yaw] 284 | AgentInfo = (dsId, frame, agentId, agentType) 285 | # if agentType == 3: 286 | # print(AgentInfo) 287 | hist, ref_pose, hist_mask = self.get_history(agentId, frame, agentId, dsId) 288 | # print(hist) 289 | fut = self.get_future(agentId, frame, dsId) 290 | fut_mask = len(fut) 291 | # print("DEL!") 292 | del self.current_dataset 293 | 294 | return hist, hist_mask, fut, fut_mask, ref_pose, AgentInfo 295 | 296 | def get_history(self, agentId, frame, ref_agentId, dsId): 297 | 298 | # Based on the reference trajectory, get a relative trajectory 299 | ref_track = self.current_dataset[self.current_dataset[:,self.id_ind]==agentId].astype(float) 300 | ref_pose = ref_track[ref_track[:,0]==frame][0, [self.x_ind, self.y_ind]] 301 | 302 | agent_track = self.current_dataset[self.current_dataset[:,self.id_ind]==agentId].astype(float) 303 | 304 | # for setting interval 305 | obs_length = int(np.argwhere(agent_track[:, 0] == frame).item(0)/self.f_interval + 1) if np.argwhere(agent_track[:, 0] == frame).item(0) - (self.obs_length-1)*(self.f_interval) < 0 else self.obs_length 306 | start_idx = np.maximum(0, np.argwhere(agent_track[:, 0] == frame).item(0) - (obs_length-1)*self.f_interval) 307 | end_idx = np.argwhere(agent_track[:, 0] == frame).item(0) + 1 308 | # print("hist::::") 309 | # print("frame: ", agent_track[start_idx:end_idx:self.f_interval, 0]) 310 | hist = agent_track[start_idx:end_idx:self.f_interval,[self.x_ind, self.y_ind]] - ref_pose# Get only relative positions [m] 311 | reasonable_inds = np.where(agent_track[start_idx:end_idx:self.f_interval, 0]>=frame-self.f_interval*(self.obs_length-1))[0] 312 | # print(reasonable_inds) 313 | hist = hist[reasonable_inds] 314 | # print("HIST: ", hist) 315 | hist_mask = len(hist) 316 | if len(hist) < self.obs_length: 317 | tmp0 = np.full((self.obs_length,2), hist[0]) 318 | # tmp0 = np.full((self.obs_length,2), 1e-6) 319 | tmp0[tmp0.shape[0]-hist.shape[0]:,:] = hist 320 | return tmp0, ref_pose, hist_mask 321 | 322 | return hist, ref_pose, hist_mask 323 | 324 | def get_future(self, agentId, frame, dsId): 325 | agent_track = self.current_dataset[self.current_dataset[:,self.id_ind]==agentId].astype(float) 326 | ref_pose = agent_track[agent_track[:,0]==frame][0, [self.x_ind, self.y_ind]] 327 | 328 | start_idx = np.argwhere(agent_track[:, 0] == frame).item(0)+self.f_interval#t+1 frame 329 | end_idx = np.minimum(len(agent_track), np.argwhere(agent_track[:, 0] == frame).item(0) + self.pred_length*self.f_interval + 1)#t+future frame 330 | # start_idx = np.argwhere(agent_track[:, 0] == frame).item(0)+1#t+1 frame 331 | # end_idx = np.minimum(len(agent_track), np.argwhere(agent_track[:, 0] == frame).item(0) + self.pred_length + 1)#t+future frame 332 | # print("fut::::") 333 | # print("frame: ", agent_track[start_idx:end_idx:self.f_interval, 0]) 334 | fut = agent_track[start_idx:end_idx:self.f_interval,[self.x_ind, self.y_ind]] - ref_pose 335 | reasonable_inds = np.where(agent_track[start_idx:end_idx:self.f_interval, 0]<=frame+self.f_interval*self.pred_length)[0] 336 | fut = fut[reasonable_inds] 337 | # print("FUT: ", fut) 338 | return fut 339 | 340 | 341 | def GetBatch(self, samples): 342 | 343 | # Filtering bad data sample// It can be modified, for filtering data which you don't want to involve 344 | # samples = [sample_pt for sample_pt in samples if sample_pt[-1][-1] == 3] 345 | # while self.batch_size > len(samples): 346 | # new_sample = self[np.random.randint(0, len(self))] 347 | # if new_sample[-1][-1] == 3: 348 | # samples.append(new_sample) 349 | # print(len(samples)) 350 | 351 | # quit() 352 | # Initialization 353 | hist_batch = torch.zeros(len(samples), self.obs_length, 2) 354 | fut_batch = torch.zeros(len(samples), self.pred_length, 2) 355 | fut_mask_batch = torch.zeros(len(samples), self.pred_length, 2) 356 | ref_pose_batch = torch.zeros(len(samples), 2) 357 | AgentsInfo = [None]*len(samples) 358 | for sampleId, (hist, hist_mask, fut, fut_mask, ref_pose, AgentInfo) in enumerate(samples): 359 | hist_batch[sampleId, :, 0] = torch.from_numpy(hist[:, 0])# x 360 | hist_batch[sampleId, :, 1] = torch.from_numpy(hist[:, 1])# y 361 | fut_batch[sampleId, 0:len(fut), 0] = torch.from_numpy(fut[:, 0])# x 362 | fut_batch[sampleId, 0:len(fut), 1] = torch.from_numpy(fut[:, 1])# y 363 | fut_mask_batch[sampleId, 0:len(fut), :] = 1# future (x,y) exist or not? 364 | ref_pose_batch[sampleId, 0] = ref_pose[0] 365 | ref_pose_batch[sampleId, 1] = ref_pose[1] 366 | AgentsInfo[sampleId] = AgentInfo 367 | 368 | return hist_batch, fut_batch, fut_mask_batch, ref_pose_batch, AgentsInfo 369 | 370 | def normalize_data(self): 371 | self.current_dataset[:, self.x_ind] = ( 372 | (self.current_dataset[:, self.x_ind] - self.min_position_x) / (self.max_position_x - self.min_position_x) 373 | ) * 2 - 1 374 | self.current_dataset[:, self.y_ind] = ( 375 | (self.current_dataset[:, self.y_ind] - self.min_position_y) / (self.max_position_y - self.min_position_y) 376 | ) * 2 - 1 377 | 378 | ## Batchwise MSE loss, uses mask for variable output lengths 379 | def maskedMSE(y_pred, y_gt, mask, device='cpu'): 380 | acc = torch.zeros_like(mask, device=device) 381 | muX = y_pred[:,:,0] 382 | muY = y_pred[:,:,1] 383 | x = y_gt[:,:, 0] 384 | y = y_gt[:,:, 1] 385 | 386 | out = torch.pow(x-muX, 2) + torch.pow(y-muY, 2) 387 | acc[:,:,0] = out 388 | acc[:,:,1] = out 389 | acc = acc*mask 390 | 391 | # fde_weight = 1 392 | # tmp0 = mask.cpu().numpy() 393 | # for i in range(len(tmp0)): 394 | # if list(np.where(tmp0[i][:,0]==1))[0] != []: 395 | # acc[i,list(np.where(tmp0[i][:,0]==1))[0][-1]]*=fde_weight 396 | # else: 397 | # continue 398 | # acc[i,list(np.where(tmp0[i][:,0]==1))[0][-1]]*=fde_weight 399 | # print 400 | lossVal = torch.sum(acc)/torch.sum(mask) 401 | 402 | return lossVal 403 | 404 | 405 | def maskedLastPositionLoss(y_pred, y_gt, mask, device='cpu'): 406 | 407 | masking_count = 0 408 | lossVal = 0 409 | batch_size = y_pred.shape[0] 410 | mask = torch.zeros((mask.shape)).cuda() 411 | for batch_ind in range(batch_size): 412 | masked_pred = y_pred[batch_ind][y_pred[batch_ind][:,0]!=0] 413 | if masked_pred.shape[0] == 0:#no GT exists 414 | continue 415 | # Sampling 416 | mean = [masked_pred[-1][0].item(), masked_pred[-1][1].item()] 417 | cov = [[masked_pred[-1][2].item()**2, 0],\ 418 | [0, masked_pred[-1][3].item()**2]] 419 | 420 | sample_x, sample_y = np.random.multivariate_normal(mean, cov, 1).T 421 | y_pred[batch_ind][masked_pred.shape[0]-1] = torch.tensor([sample_x[0], sample_y[0], 0., 0.]).cuda() 422 | mask[batch_ind][masked_pred.shape[0]-1] = torch.tensor([1,1,1,1]).cuda()#only considers last position 423 | 424 | acc = torch.zeros_like(mask, device=device) 425 | muX = y_pred[:,:,0] 426 | muY = y_pred[:,:,1] 427 | x = y_gt[:,:, 0] 428 | y = y_gt[:,:, 1] 429 | 430 | out = torch.pow(x-muX, 2) + torch.pow(y-muY, 2) 431 | acc[:,:,0] = out 432 | acc = acc[:,:,:1] 433 | mask = mask[:,:,:1] 434 | acc = acc*mask 435 | 436 | lossVal = torch.sum(acc)/torch.sum(mask) 437 | 438 | return lossVal -------------------------------------------------------------------------------- /vanilla_gru.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import random 5 | import numpy as np 6 | 7 | class NNPred(nn.Module): 8 | 9 | def __init__(self, args, dropout=0.5): 10 | super(NNPred, self).__init__() 11 | 12 | torch.manual_seed(args.seed) 13 | torch.cuda.manual_seed(args.seed) 14 | 15 | self.batch_size = args.batch_size 16 | self.teacher_forcing_ratio = 0.5#if it increases, more label is used for training 17 | 18 | self.using_cuda = args.cuda 19 | self.device = args.device 20 | self.train_flag = args.train 21 | 22 | 23 | # self.hidden_size = hidden_size 24 | hidden_size = 256 25 | self.in_length = args.obs_length 26 | self.out_length = args.pred_length 27 | self.num_network_output = 2#mx,my,sx,sy,no correlation 28 | # self.num_layers = 2 29 | 30 | self.in2gru = nn.Linear(2, hidden_size)#Embedding 31 | self.bn_gru = torch.nn.BatchNorm1d(hidden_size) 32 | self.gru = nn.GRU(hidden_size, hidden_size, num_layers = 2, batch_first=True) 33 | self.gruCell0 = nn.GRUCell(hidden_size, hidden_size) 34 | self.gruCell1 = nn.GRUCell(hidden_size, hidden_size) 35 | # self.bilstm = nn.LSTM(hidden_size, hidden_size//2,num_layers=self.num_layers, batch_first=True,dropout = dropout) 36 | 37 | self.fc0 = nn.Linear(hidden_size,hidden_size//2) 38 | self.bn_lin0 = torch.nn.BatchNorm1d(hidden_size//2) 39 | self.fc1 = nn.Linear(hidden_size//2,self.num_network_output) 40 | 41 | self.leaky_relu = nn.LeakyReLU() 42 | self.elu = torch.nn.ELU() 43 | self.tanh = torch.nn.Tanh() 44 | # self.in2out = nn.Linear(input_size, 64) 45 | 46 | 47 | def forward(self, input, label): 48 | 49 | last_inputEmbedding, last_hidden = self.Encoder(input) 50 | output = self.Decoder(last_inputEmbedding, last_hidden, label) 51 | # print("output.shape: ", output.shape) 52 | return output 53 | 54 | def Encoder(self, input): 55 | input = self.in2gru(input) 56 | # print("encode1: ", input.shape) 57 | input = input.permute(0, 2, 1)# Becuase it is a sequence 58 | inputEmbedding = self.bn_gru(input).permute(0, 2, 1)#[batch_size, input_sequence, hidden_size(feature)] 59 | 60 | inputEmbedding = self.elu(inputEmbedding) 61 | # print("encode2: ", inputEmbedding.shape) 62 | _, hidden_n = self.gru(inputEmbedding)#[batch_size, input_sequence, hidden_size(feature)], [num_layers, batch_size, last_hidden(feature)] 63 | last_inputEmbedding = inputEmbedding[:,-1,:]#[batch_size, last_input_feature] 64 | # print("encode_last: ", last_inputEmbedding.shape) 65 | # print("last_hidden: ", hidden_n) 66 | 67 | return last_inputEmbedding, hidden_n 68 | 69 | def Decoder(self, input, hidden, label):#2layer decoder 70 | 71 | output = torch.zeros((self.batch_size, self.out_length, self.num_network_output)) 72 | if self.cuda: 73 | output = output.cuda(self.device) 74 | hidden0 = hidden[0] 75 | hidden1 = hidden[1] 76 | for i in range(self.out_length): 77 | if i != 0: 78 | if self.train_flag and (random.random() < self.teacher_forcing_ratio):#only training 79 | #only for self.num_network_ouput >=2 80 | # cat_label = torch.cat((label[:,i-1,:], pred_input[:,2:]), 1) 81 | # input = self.bn_gru(self.in2gru(cat_label)) 82 | input = self.bn_gru(self.in2gru(label[:,i-1,:])) 83 | input = self.elu(input)#[batch_size, output_sequence, hidden_size(feature)] 84 | else: 85 | input = self.bn_gru(self.in2gru(pred_input[:,:2])) 86 | input = self.elu(input) 87 | # print("decoder input: ", input.shape) 88 | # print("hidden0: ", hidden0.shape) 89 | # print("hidden1: ", hidden1.shape) 90 | hidden0 = self.gruCell0(input, hidden0) 91 | hidden1 = self.gruCell1(hidden0, hidden1) 92 | pred_input = self.FClayers(hidden1)#Get predicted 93 | 94 | output[:,i,:] = pred_input 95 | # print(output) 96 | # quit() 97 | return output 98 | 99 | 100 | def FClayers(self, input): 101 | input = self.bn_lin0(self.fc0(input)) 102 | input = self.elu(input) 103 | pred_input = self.fc1(input)#Get predicted 104 | 105 | return pred_input 106 | --------------------------------------------------------------------------------