├── checkpoints └── best_model.pth ├── loss.py ├── config └── config.json ├── utils ├── ranking.py ├── writer.py ├── collision.py ├── plots.py ├── datainfo.py ├── maps.py └── utils.py ├── LICENSE ├── .gitignore ├── README.md ├── evaluate.py ├── loader.py ├── main.py └── models.py /checkpoints/best_model.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SeongjuLee/DCENet-PyTorch/HEAD/checkpoints/best_model.pth -------------------------------------------------------------------------------- /loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class DCENetLoss(nn.Module): 7 | def __init__(self, config): 8 | super(DCENetLoss, self).__init__() 9 | self.beta = config['beta'] 10 | self.pred_seq = config['pred_seq'] 11 | 12 | def forward(self, mu, log_var, y_pred, y_true): 13 | kl_loss = 0.5 * torch.sum(mu.pow(2) + log_var.exp() - log_var - 1) 14 | mse_loss = self.pred_seq * F.mse_loss(y_pred, y_true) 15 | 16 | loss = mse_loss * self.beta + kl_loss * (1 - self.beta) 17 | return loss 18 | -------------------------------------------------------------------------------- /config/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "o_drop": 0.2, 3 | "s_drop": 0.1, 4 | "z_drop": 0.15, 5 | "n_hidden": 512, 6 | "z_dim": 2, 7 | "encoder_dim": 16, 8 | "z_decoder_dim": 64, 9 | "hidden_size": 32, 10 | "x_encoder_dim": 64, 11 | "y_encoder_dim": 64, 12 | "x_encoder_layers": 3, 13 | "y_encoder_layers": 3, 14 | "x_encoder_head": 8, 15 | "y_encoder_head": 8, 16 | "occu_encoder_x_dim": 32, 17 | "occu_encoder_y_dim": 32, 18 | "occu_encoder_x_layers": 3, 19 | "occu_encoder_y_layers": 3, 20 | "occu_encoder_x_head": 2, 21 | "occu_encoder_y_head": 2, 22 | 23 | "obs_seq": 8, 24 | "enviro_pdim": [32, 32, 3], 25 | "pred_seq": 12, 26 | 27 | "lr": 3e-5, 28 | "beta": 0.65, 29 | "split": 0.8, 30 | "max_epochs": 1, 31 | "batch_size": 192, 32 | "patience": 5, 33 | 34 | "num_pred": 25, 35 | 36 | "neptune": false, 37 | "neptune_token": "" 38 | } -------------------------------------------------------------------------------- /utils/ranking.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import numpy as np 4 | from scipy.stats import multivariate_normal 5 | 6 | 7 | def gauss_rank(pred_trajs, addnoise=False): 8 | ''' 9 | pred_trajs: numberofpredictions*length*[x, y] 10 | ''' 11 | # Swap time axis to the first 12 | pred_trajs_t = np.swapaxes(pred_trajs, 1, 0) 13 | rank = np.zeros((0, pred_trajs.shape[0])) 14 | for pred_poss in pred_trajs_t: 15 | 16 | # pred_poss is the sampled positions at each time step 17 | # pred_poss will be used to fit a bivariable gaussian distribution 18 | if addnoise == True: 19 | pred_poss = pred_poss + np.random.normal(0, 1, pred_poss.shape) 20 | mu = np.mean(pred_poss, axis=0) 21 | covariance = np.cov(pred_poss.T) 22 | pos_pdf = multivariate_normal.pdf(pred_poss, mean=mu, cov=covariance) 23 | rank = np.vstack((rank, pos_pdf)) 24 | rank = np.mean(np.log(rank), axis=0) 25 | return rank -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 SeongjuLee 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # Custom 132 | processed_data/ 133 | -------------------------------------------------------------------------------- /utils/writer.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import numpy as np 4 | from utils.ranking import gauss_rank 5 | from utils.maps import Maps 6 | 7 | 8 | def test(): 9 | path = "../WORLD H-H TRAJ/Train/crowds/crowds_zara03.txt" 10 | trajs = sort_data(path, 20) 11 | trajs = trajs[:, :8, :] 12 | print(trajs.shape) 13 | 14 | 15 | def sort_data(path, seq_length): 16 | data = np.genfromtxt(path, delimiter='') 17 | # challenge dataset have nan for prediction time steps 18 | data = data[~np.isnan(data).any(axis=1)] 19 | datamaps = Maps(data) 20 | trajs = np.reshape(datamaps.sorted_data, (-1, seq_length, 4)) 21 | return trajs 22 | 23 | 24 | def write_pred_txt(obs_trajs, pred_trajs, dataname, folder): 25 | with open("../%s/%s.txt"%(folder, dataname), 'w+') as myfile: 26 | for i, obs_traj in enumerate(obs_trajs): 27 | start_frame, pedestrian, _, _ = obs_traj[0] 28 | step = obs_traj[-1, 0] - obs_traj[-2, 0] 29 | predictions = pred_trajs[i] 30 | 31 | if predictions.shape[0]>1: 32 | # Get the ranks and select the most likely prediction 33 | ranks = gauss_rank(predictions) 34 | ranked_predicion = predictions[np.argmax(ranks)] 35 | else: 36 | ranked_predicion = predictions[0] 37 | 38 | traj = np.concatenate((obs_traj[:, 2:4], ranked_predicion), axis=0) 39 | for j, pos in enumerate(traj): 40 | myfile.write(str(int(start_frame+j*step))+' '+ 41 | str(round(pedestrian, 0))+' '+ 42 | str(round(pos[0], 3))+' '+ 43 | str(round(pos[1], 3))+' '+'\n') 44 | myfile.close() 45 | 46 | 47 | def get_index(obs_trajs, pred_trajs): 48 | 49 | trajectories = [] 50 | 51 | for i, obs_traj in enumerate(obs_trajs): 52 | start_frame, pedestrian, _, _ = obs_traj[0] 53 | step = obs_traj[-1, 0] - obs_traj[-2, 0] 54 | predictions = pred_trajs[i] 55 | 56 | if predictions.shape[0]>1: 57 | # Get the ranks and select the most likely prediction 58 | ranks = gauss_rank(predictions) 59 | ranked_predicion = predictions[np.argmax(ranks)] 60 | else: 61 | ranked_predicion = predictions[0] 62 | 63 | traj = np.concatenate((obs_traj[:, 2:4], ranked_predicion), axis=0) 64 | for j, pos in enumerate(traj): 65 | trajectories.append([start_frame+j*step, 66 | pedestrian, 67 | pos[0], 68 | pos[1]]) 69 | 70 | return np.asarray(trajectories) 71 | if __name__ == "__main__": 72 | test() -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DCENet-PyTorch 2 | PyTorch Implementation of DCENet (https://arxiv.org/abs/2010.16267) for Trajectory Forecasting 3 | 4 | 5 | ## Requirements 6 | * `python 3.8` 7 | * `pytorch 1.7.1` 8 | * `matplotlib` 9 | * `scipy` 10 | * `neptune-client` if you need 11 | 12 | 13 | ## Dataset 14 | Processed data (`./processed_data/train/train_merged0.npz`, `./processed_data/train/train_merged1.npz`, `./processed_data/train/train_merged2.npz`, `./processed_data/train/biwi_hotel.npz`) is requried. 15 | 16 | You can obtain the processed data from the original repository (https://github.com/tanjatang/DCENet). 17 | 18 | ## Train 19 | * Command for training 20 | 21 | `python main.py --gpu $GPU_NUMS --config $CONFIG_FILENAME` 22 | 23 | * Example 24 | 25 | `python main.py --gpu 0 --config config.json` 26 | 27 | ## Test 28 | * Command for evaluation 29 | 30 | `python evaluate.py --gpu $GPU_NUMS --config $CONFIG_FILENAME --resume-name #CHECKPOINT_FILENAME` 31 | 32 | * Example 33 | 34 | `python evaluate.py --gpu 0 --config config.json --resume-name best_model.pth` 35 | 36 | ## Performance on Biwi Hotel 37 | ### Evaluation Results @Top25 38 | | Criteria | Original Implementation (Tensorflow) | My Implementation (PyTorch) | 39 | |:------------------: |:------------------------------------: |:---------------------------: | 40 | | ADE | 0.37 m | 0.36 m | 41 | | FDE | 0.76 m | 0.67 m | 42 | | Hausdorff Distance | 0.75 m | 0.67 m | 43 | | Speed Deviation | 0.06 m/s | 0.05 m/s | 44 | | Heading Error | 25.60 degree | 24.67 degree | 45 | 46 | ### Evaluation Results for Most-likely Predictions 47 | | Criteria | Original Implementation (Tensorflow) | My Implementation (PyTorch) | 48 | |:------------------: |:------------------------------------: |:---------------------------: | 49 | | ADE | 0.39 m | 0.42 m | 50 | | FDE | 0.78 m | 0.79 m | 51 | | Hausdorff Distance | 0.77 m | 0.78 m | 52 | | Speed Deviation | 0.06 m/s | 0.05 m/s | 53 | | Heading Error | 30.98 degree | 30.62 degree | 54 | 55 | 56 | ## License 57 | Model details and most of utility functions are from from the origianl DCENet repository (https://github.com/tanjatang/DCENet). 58 | 59 | Codes for progress bar came from https://github.com/AaronHeee/MEAL. 60 | 61 | Codes for early stopping came from https://github.com/Bjarten/early-stopping-pytorch. 62 | 63 | ## Who Am I? 64 | I am on Ph.D course in Artificial Intelligence Lab. ([Homepage](https://ailab.gist.ac.kr/ailab/)), Gwangju Institute of Science and Technology (GIST, [Homepage](https://www.gist.ac.kr/kr/event_2st/index.html)), Korea. 65 | -------------------------------------------------------------------------------- /utils/collision.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | # from utils import writer 6 | from utils.maps import Maps 7 | 8 | 9 | def check_collision(data, thred=0.1, n=1, obs_seq=8, pred_seq=12): 10 | """ 11 | data: frameId, pedId, x, y 12 | data.shape: num_trajs*seq_length*4, 13 | """ 14 | print("The shape of the input data", data.shape) 15 | if len((data.shape))==2: 16 | datamaps = Maps(data) 17 | data = np.reshape(datamaps.sorted_data, (-1, obs_seq+pred_seq, 4)) 18 | data = data [:, obs_seq:, :] 19 | print("The shape of the new data", data.shape) 20 | 21 | 22 | count_collisions = 0 23 | encounters = 0 24 | traj_data = data.reshape(-1, 4) 25 | 26 | for ped_traj in data: 27 | 28 | ego_pedid = ped_traj[0, 1] 29 | ped_frameIds = ped_traj[:, 0] 30 | 31 | co_traj_data = traj_data[traj_data[:, 0]>=np.min(ped_frameIds)] 32 | co_traj_data = co_traj_data[co_traj_data[:, 0]<=np.max(ped_frameIds)] 33 | co_pedids = np.unique(co_traj_data[:, 1]) 34 | 35 | for co_pedid in co_pedids: 36 | if co_pedid != ego_pedid: 37 | con_ped_traj = co_traj_data[co_traj_data[:, 1]==co_pedid] 38 | if con_ped_traj.size != 0: 39 | encounters += 1 40 | count = count_collision(ped_traj, con_ped_traj, thred, n) 41 | count_collisions += count 42 | 43 | print("Total trajectories %.0f, Total encounters %.0f, collisions %.0f, collision rate %.2f"% 44 | (len(data), encounters, count_collisions, count_collisions/encounters)) 45 | return encounters, count_collisions 46 | 47 | 48 | 49 | def count_collision(ego_ped_traj, co_ped_traj, thred, n): 50 | ego_ped_traj = ego_ped_traj[ego_ped_traj[:, 0]>=np.min(co_ped_traj[:, 0])] 51 | ego_ped_traj = ego_ped_traj[ego_ped_traj[:, 0]<=np.max(co_ped_traj[:, 0])] 52 | 53 | # Interpolation 54 | ego_x, ego_y = linear_interp(ego_ped_traj[:, 2], ego_ped_traj[:, 3], n) 55 | co_x, co_y = linear_interp(co_ped_traj[:, 2], co_ped_traj[:, 3], n) 56 | 57 | count = eucl_dis(np.vstack((ego_x, ego_y)).T, np.vstack((co_x, co_y)).T, n, thred) 58 | return count 59 | 60 | 61 | def linear_interp(x, y, n): 62 | ''' 63 | This is the one-D linear interpolation in each dimension 64 | ''' 65 | xvals_ = np.empty((0)) 66 | yvals_ = np.empty((0)) 67 | for i, j in enumerate(x): 68 | if i 1 else DCENet(config).to(device) 47 | 48 | # ================= Data Loader ================ # 49 | datalist = DataInfo() 50 | test_datalist = datalist.train_biwi 51 | print('Test data list', test_datalist) 52 | 53 | test_offsets, test_trajs, test_occupancy = load_data(config, test_datalist, datatype="test") 54 | test_x = test_offsets[:, :config['obs_seq'] - 1, 4:6] 55 | test_occu = test_occupancy[:, :config['obs_seq'] - 1, ..., :config['enviro_pdim'][-1]] 56 | last_obs_test = test_offsets[:, config['obs_seq'] - 2, 2:4] 57 | y_truth = test_offsets[:, config['obs_seq'] - 1:, :4] 58 | xy_truth = test_offsets[:, :, :4] 59 | 60 | print('test_trajs', test_trajs.shape) 61 | 62 | print("%.0f trajectories for testing" % (test_x.shape[0])) 63 | 64 | # Test 65 | model.load_state_dict(torch.load(os.path.join('checkpoints', args.resume_name))) 66 | model.eval() 67 | torch.manual_seed(10) 68 | with torch.no_grad(): 69 | test_x, test_occu = input2tensor(test_x, test_occu, device) 70 | x_latent = model.encoder_x(test_x, test_occu) 71 | predictions = [] 72 | for i, x_ in enumerate(x_latent): 73 | last_pos = last_obs_test[i] 74 | x_ = x_.view(1, -1) 75 | for i in range(config['num_pred']): 76 | y_p = model.decoder(x_, train=False) 77 | y_p_ = np.concatenate(([last_pos], np.squeeze(y_p.cpu().numpy())), axis=0) 78 | y_p_sum = np.cumsum(y_p_, axis=0) 79 | predictions.append(y_p_sum[1:, :]) 80 | 81 | predictions = np.reshape(predictions, [-1, config['num_pred'], config['pred_seq'], 2]) 82 | 83 | print('Predicting done!') 84 | print(predictions.shape) 85 | plot_pred(xy_truth, predictions) 86 | # Get the errors for ADE, DEF, Hausdorff distance, speed deviation, heading error 87 | print("\nEvaluation results @top%.0f" % config['num_pred']) 88 | errors = get_errors(y_truth, predictions) 89 | check_collision(y_truth) 90 | 91 | ## Get the first time prediction by g 92 | ranked_prediction = [] 93 | for prediction in predictions: 94 | ranks = gauss_rank(prediction) 95 | ranked_prediction.append(prediction[np.argmax(ranks)]) 96 | ranked_prediction = np.reshape(ranked_prediction, [-1, 1, config['pred_seq'], 2]) 97 | print("\nEvaluation results for most-likely predictions") 98 | ranked_errors = get_errors(y_truth, ranked_prediction) 99 | 100 | 101 | if __name__ == "__main__": 102 | main() 103 | -------------------------------------------------------------------------------- /loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import numpy as np 4 | from torch.utils.data import Dataset 5 | 6 | 7 | class TrajDataset(Dataset): 8 | def __init__(self, x, x_occu, y, y_occu, mode='train'): 9 | self.x = x 10 | self.x_occu = x_occu 11 | self.y = y 12 | self.y_occu = y_occu 13 | self.mode = mode 14 | 15 | def __getitem__(self, index): 16 | # Shape of x, y: [seq_len, input_len(2)], Shape of x_occu, y_occu: [channel, seq_len, height, width] 17 | if self.mode == 'train' or self.mode == 'val': 18 | x = torch.tensor(self.x[index], dtype=torch.float) 19 | x_occu = torch.tensor(self.x_occu[index], dtype=torch.float).permute(0, 3, 1, 2) 20 | y = torch.tensor(self.y[index], dtype=torch.float) 21 | y_occu = torch.tensor(self.y_occu[index], dtype=torch.float).permute(0, 3, 1, 2) 22 | 23 | return x, x_occu, y, y_occu 24 | 25 | elif self.mode == 'test': 26 | x = torch.tensor(self.x[index], dtype=torch.float) 27 | x_occu = torch.tensor(self.x_occu[index], dtype=torch.float).permute(0, 3, 1, 2) 28 | y = torch.tensor(self.y[index], dtype=torch.float) 29 | 30 | return x, x_occu, y 31 | 32 | elif self.mode == 'challenge': 33 | x = torch.tensor(self.x[index], dtype=torch.float) 34 | x_occu = torch.tensor(self.x_occu[index], dtype=torch.float).permute(0, 3, 1, 2) 35 | 36 | return x, x_occu 37 | 38 | else: 39 | raise NotImplemented 40 | 41 | def __len__(self): 42 | return len(self.x) 43 | 44 | 45 | def load_data(config, dataset_list, datatype="train"): 46 | # Store the data across datasets 47 | # All the datasets are merged for training 48 | if datatype == "train" or datatype == "test": 49 | offsets = np.empty((0, config['obs_seq'] + config['pred_seq'] - 1, 8)) 50 | traj_data = np.empty((0, config['obs_seq'] + config['pred_seq'], 4)) 51 | occupancy = np.empty((0, config['obs_seq'] + config['pred_seq'] - 1, config['enviro_pdim'][0], config['enviro_pdim'][1], 3)) 52 | 53 | if dataset_list[0] == "train_merged": 54 | for i in range(3): 55 | data = np.load("./processed_data/train/%s.npz" % (dataset_list[0]+str(i))) 56 | _offsets, _traj_data, _occupancy = data["offsets"], data["traj_data"], data["occupancy"] 57 | 58 | 59 | offsets = np.concatenate((offsets, _offsets), axis=0) 60 | traj_data = np.concatenate((traj_data, _traj_data), axis=0) 61 | 62 | occupancy = np.concatenate((occupancy, _occupancy), axis=0) 63 | print(dataset_list[0], "contains %.0f trajectories" % len(offsets)) 64 | else: 65 | for i, dataset in enumerate(dataset_list): 66 | # Only take the orinal data 67 | # ToDo, here needs to be test if augumentation will boost the performance 68 | # if dataset != "train_merged": 69 | # ToDo chenge this to make compatible with linus 70 | data = np.load("./processed_data/train/%s.npz" % (dataset)) 71 | _offsets, _traj_data, _occupancy = data["offsets"], data["traj_data"], data["occupancy"] 72 | print(dataset, "contains %.0f trajectories" % len(_offsets)) 73 | offsets = np.concatenate((offsets, _offsets), axis=0) 74 | traj_data = np.concatenate((traj_data, _traj_data), axis=0) 75 | occupancy = np.concatenate((occupancy, _occupancy), axis=0) 76 | 77 | # NOTE: When load the challenge data, there is no need to merge them 78 | # The submission requires each challenge data set (in total 20) to be separated 79 | # Hence, each time only one challenge data set is called 80 | elif datatype == "challenge": 81 | offsets = np.empty((0, config['obs_seq'] - 1, 8)) 82 | traj_data = np.empty((0, config['obs_seq'], 4)) 83 | occupancy = np.empty((0, config['obs_seq'] - 1, config['enviro_pdim'][0], config['enviro_pdim'][1], 3)) 84 | for dataset in dataset_list: 85 | data = np.load("./processed_data/challenge/%s.npz" % (dataset)) 86 | _offsets, _traj_data, _occupancy = data["offsets"], data["traj_data"], data["occupancy"] 87 | offsets = np.concatenate((offsets, _offsets), axis=0) 88 | traj_data = np.concatenate((traj_data, _traj_data), axis=0) 89 | occupancy = np.concatenate((occupancy, _occupancy), axis=0) 90 | 91 | elif datatype == "test": 92 | assert len(dataset_list) == 1, print("Only one untouched dataset is left fot testing!") 93 | elif datatype == "challenge": 94 | assert len(dataset_list) == 1, print("predict one by one") 95 | if datatype == "train": 96 | if not os.path.exists("./processed_data/train/train_merged2.npz"): 97 | # Save the merged training data 98 | # sigle file storage more than 16G is not supported in linux system, so I tried to store them in 3 files. 99 | # it's not necessary for windows user to separate data into 3 files 100 | offsets_list = [offsets[:15000,:],offsets[15000:30000,:],offsets[30000:,:]] 101 | traj_data_list = [traj_data[:15000,:],traj_data[15000:30000,:],traj_data[30000:,:]] 102 | occupancy_list = [occupancy[:15000,:],occupancy[15000:30000,:],occupancy[30000:,:]] 103 | 104 | for i in range(len(offsets_list)): 105 | np.savez("./processed_data/train/train_merged%s.npz"%(i), 106 | offsets=offsets_list[i], 107 | traj_data=traj_data_list[i], 108 | occupancy=occupancy_list[i]) 109 | 110 | return offsets, traj_data, occupancy -------------------------------------------------------------------------------- /utils/maps.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | from scipy.ndimage.filters import gaussian_filter 6 | 7 | class Maps(): 8 | 9 | def __init__(self, data, scale=1, sigma=1, pad=0.5): 10 | self.scale = scale 11 | self.sigma = sigma 12 | self.pad = pad 13 | self.xmin, self.ymin = np.min(data[:, 2:4], axis=0) 14 | # Shift the x and y coordinates from zero centralized to positive 15 | self.data = np.subtract(data, [0, 0, self.xmin, self.ymin]) 16 | # Scale the x and y coordinates from meter to decimeter (default unit meters and scale=10) 17 | self.data = np.multiply(self.data, [1, 1, self.scale, self.scale]) 18 | # Compute the height and width of the maps 19 | self.dimensions = np.flip(np.ceil(np.max(self.data[:, 2:4], axis=0)).astype(int)) 20 | # Very important, the order of userId should be obtained 21 | ped_ids = [] 22 | for ped_id in data[:, 1]: 23 | if ped_id not in ped_ids: 24 | ped_ids.append(ped_id) 25 | self.ped_ids = np.asarray(ped_ids) 26 | 27 | fig, ax = plt.subplots() 28 | ax.set_aspect("equal") 29 | offsets = np.empty((0, 6)) 30 | sorted_data = np.empty((0, 4)) 31 | for ped_id in self.ped_ids: 32 | ped_traj = self.data[self.data[:, 1]==ped_id, :] 33 | sorted_data = np.vstack((sorted_data, ped_traj)) 34 | 35 | ped_offset = ped_traj[1:, 2:4] - ped_traj[:-1, 2:4] 36 | ped_offset = np.concatenate((ped_traj[1:, :], ped_offset), axis=1) 37 | offsets = np.vstack((offsets, ped_offset)) 38 | ax.plot(ped_traj[:, 2], ped_traj[:, 3]) 39 | theta = np.arctan2(offsets[:, 5], offsets[:, 4]) 40 | velocity = np.linalg.norm(offsets[:, 4:6], axis=1)/0.4 41 | 42 | offsets = np.concatenate((offsets, theta.reshape(-1, 1), velocity.reshape(-1, 1)), axis=1) 43 | ax.set_title("Trajectories") 44 | plt.gca().invert_yaxis() 45 | plt.show() 46 | plt.gcf().clear() 47 | plt.close() 48 | 49 | self.offsets = offsets 50 | self.min_speed = np.min(self.offsets[:, -1]) 51 | self.max_speed = np.max(self.offsets[:, -1]) 52 | 53 | sorted_data = np.divide(sorted_data, [1, 1, self.scale, self.scale]) 54 | self.sorted_data = np.add(sorted_data, [0, 0, self.xmin, self.ymin]) 55 | 56 | 57 | def trajectory_map(self, traj_data=None): 58 | ''' 59 | This is the main function to get the heatmap for scene context 60 | ''' 61 | # Gaussian Sigma 62 | if np.all(traj_data)!=None: 63 | x, y = traj_data[:, 2], traj_data[:, 3] 64 | else: 65 | x, y = self.data[:, 2], self.data[:, 3] 66 | 67 | # Here need to check if the filtered data is empty 68 | if x.size == 0 or y.size == 0: 69 | bins_real = [int(self.dimensions[1]), int(self.dimensions[0])] 70 | bin_x_min = 0 71 | bin_y_min = 0 72 | else: 73 | #### solve the offset problem 74 | bin_x_max, bin_x_min = int(np.ceil(min(np.max(x), self.dimensions[1]))), int(np.floor(max(np.min(x), 0))) 75 | bin_y_max, bin_y_min = int(np.ceil(min(np.max(y), self.dimensions[0]))), int(np.floor(max(np.min(y), 0))) 76 | bin_width = max((bin_x_max - bin_x_min), 1) 77 | bin_heiht = max((bin_y_max - bin_y_min), 1) 78 | # Get the actual heatmap where data can be seen 79 | bins_real = [bin_width, bin_heiht] 80 | 81 | def heatmap(x, y, s, bins=1000): 82 | heatmap, xedges, yedges = np.histogram2d(x, y, bins=bins) 83 | heatmap = gaussian_filter(heatmap, sigma=s) 84 | if heatmap.max() != 0: 85 | nor_heatmap = heatmap / heatmap.max() 86 | else: 87 | nor_heatmap = heatmap 88 | return heatmap.T, nor_heatmap.T 89 | 90 | img, nor_heatmap = heatmap(x, y, self.sigma, bins_real) 91 | # Align the data to the image size 92 | traj_map = np.zeros(self.dimensions) 93 | for r, row in enumerate(nor_heatmap): 94 | for c, column in enumerate(row): 95 | traj_map[r+bin_y_min, c+bin_x_min] = column 96 | assert np.max(traj_map)<=1 and np.max(traj_map)>=0, print("traj_map contains invalid values") 97 | return traj_map 98 | 99 | 100 | def motion_map(self, offsets=None, max_speed=50): 101 | """ 102 | This is the function to convert the motion in into a motion map 103 | Return speed map and orientation map 104 | """ 105 | speed_map = np.zeros((self.dimensions)) 106 | orient_map = np.zeros((self.dimensions)) 107 | count = np.zeros((self.dimensions)) 108 | if np.all(offsets) == None: 109 | offsets = self.offsets 110 | 111 | if max_speed: 112 | # clip the very unlikely speed 113 | offsets[:, -1] = np.clip(offsets[:, -1], a_min=0, a_max=max_speed) 114 | 115 | for offset in offsets: 116 | x, y, delta_x, delta_y = offset[2:6] 117 | xl, xr = sorted((x, x+delta_x)) 118 | xl, xr = max(0, int(xl-self.pad)), min(int(xr+self.pad), self.dimensions[1]) 119 | yl, yr = sorted((y, y+delta_y)) 120 | yl, yr = max(0, int(yl-self.pad)), min(int(yr+self.pad), self.dimensions[0]) 121 | # count the frequency for normalization 122 | count[yl:yr, xl:xr] += 1 123 | speed_map[yl:yr, xl:xr] += offset[-1] 124 | orient_map[yl:yr, xl:xr] += (offset[-2]+np.pi)/(2*np.pi) 125 | 126 | def normalize(count, data): 127 | norm_data = np.zeros_like(data) 128 | for i, row in enumerate(data): 129 | for j, value in enumerate(row): 130 | if count[i, j] != 0: 131 | norm_data[i, j] = value/count[i, j] 132 | return norm_data 133 | 134 | orient_map = normalize(count, orient_map) 135 | speed_map = normalize(count, speed_map) 136 | # speed_map = speed_map/self.max_speed 137 | speed_map = speed_map/max_speed 138 | 139 | assert np.max(orient_map)<=1 and np.max(orient_map)>=0, print("orient_map contains invalid values", 140 | np.max(orient_map), np.min(orient_map)) 141 | assert np.max(speed_map)<=1 and np.max(speed_map)>=0, print("speed_map contains invalid values", 142 | np.max(speed_map), np.min(speed_map)) 143 | return orient_map, speed_map -------------------------------------------------------------------------------- /utils/utils.py: -------------------------------------------------------------------------------- 1 | 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Sat Mar 2 21:31:32 2019 5 | @author: cheng 6 | """ 7 | import os 8 | import sys 9 | import time 10 | import torch 11 | import numpy as np 12 | from scipy.spatial.distance import directed_hausdorff 13 | 14 | 15 | class EarlyStopping: 16 | """Early stops the training if validation loss doesn't improve after a given patience.""" 17 | def __init__(self, patience=7, verbose=False, delta=0, path='checkpoints', filename='checkpoint.pt'): 18 | """ 19 | Args: 20 | patience (int): How long to wait after last time validation loss improved. 21 | Default: 7 22 | verbose (bool): If True, prints a message for each validation loss improvement. 23 | Default: False 24 | delta (float): Minimum change in the monitored quantity to qualify as an improvement. 25 | Default: 0 26 | path (str): Path for the checkpoint to be saved to. 27 | Default: 'checkpoint.pt' 28 | """ 29 | self.patience = patience 30 | self.verbose = verbose 31 | self.counter = 0 32 | self.best_score = None 33 | self.early_stop = False 34 | self.val_loss_min = np.Inf 35 | self.delta = delta 36 | self.path = path 37 | self.filename = filename 38 | os.makedirs(path, exist_ok=True) 39 | 40 | def __call__(self, val_loss, model): 41 | 42 | score = -val_loss 43 | 44 | if self.best_score is None: 45 | self.best_score = score 46 | self.save_checkpoint(val_loss, model) 47 | elif score < self.best_score + self.delta: 48 | self.counter += 1 49 | print(f'EarlyStopping counter: {self.counter} out of {self.patience}') 50 | if self.counter >= self.patience: 51 | self.early_stop = True 52 | else: 53 | self.best_score = score 54 | self.save_checkpoint(val_loss, model) 55 | self.counter = 0 56 | 57 | def save_checkpoint(self, val_loss, model): 58 | '''Saves model when validation loss decrease.''' 59 | if self.verbose: 60 | print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...') 61 | torch.save(model.state_dict(), os.path.join(self.path, self.filename)) 62 | self.val_loss_min = val_loss 63 | 64 | 65 | def get_errors(groundtruth, predictions): 66 | _, num_pred, pred_seq, _ = predictions.shape 67 | 68 | # Calculate the corresponding errors 69 | errors = get_evaluation(groundtruth, predictions, num_pred) 70 | # print('\nErrors by ADE and FDE\n', np.array_str(errors[0:2, 2], precision=2, suppress_small=True)) 71 | 72 | print("Avearge displacement error: %.2f [m]"%(errors[0, 2])) 73 | print("Final displacement error: %.2f [m]" % (errors[1, 2])) 74 | print("Hausdorff distance: %.2f [m]" % (errors[2, 2])) 75 | print("Speed deviation: %.2f [m/s]" % (errors[3, 2])) 76 | print("Heading error: %.2f [°]" % (errors[4, 2])) 77 | 78 | return errors 79 | 80 | 81 | def get_evaluation(groundtruth, predictions, num_pred, scale=1.0): 82 | # Evaluation 83 | evaluations = np.zeros([len(predictions), num_pred, 5]) 84 | for i, user_gt in enumerate(groundtruth): 85 | user_preds = predictions[i] 86 | for j, user_pred in enumerate(user_preds): 87 | evaluations[i, j, :] = get_eva_values(user_gt[:, 2:4]*scale, user_pred*scale) 88 | # Compute the average errors across all users and all predictions 89 | mean_evaluations = np.reshape(evaluations, [-1, 5]) 90 | mean_errors = np.mean(mean_evaluations, axis=0) 91 | mean_std = np.std(mean_evaluations, axis=0) 92 | # Comput the minimum errors across all users for the best prediction 93 | min_evaluations = np.min(evaluations, axis=1) 94 | min_errors = np.mean(min_evaluations, axis=0) 95 | min_std = np.std(min_evaluations, axis=0) 96 | # Save the evaluation results 97 | errors = np.concatenate((np.reshape(mean_errors, [-1, 1]), 98 | np.reshape(mean_std, [-1, 1]), 99 | np.reshape(min_errors, [-1, 1]), 100 | np.reshape(min_std, [-1, 1])), axis=1) 101 | return errors 102 | 103 | 104 | def get_eva_values(y_t, y_p): 105 | ''' 106 | y_t: 2d numpy array for true trajectory. Shape: steps*2 107 | y_p: 2d numpy array for predicted trajectory. Shape: steps*2 108 | ''' 109 | Euclidean = get_euclidean(y_t, y_p) 110 | last_disp = get_last_disp(y_t, y_p) 111 | Hausdorff = get_hausdorff(y_t, y_p) 112 | speed_dev = get_speeddev(y_t, y_p) 113 | heading_error = get_headerror(y_t, y_p) 114 | eva_values = [Euclidean, last_disp, Hausdorff, speed_dev, heading_error] 115 | return eva_values 116 | 117 | def get_euclidean(y_true, y_prediction): 118 | Euclidean = np.linalg.norm((y_true - y_prediction), axis=1) 119 | Euclidean = np.mean(Euclidean) 120 | return Euclidean 121 | 122 | def get_last_disp(y_true, y_prediction): 123 | last_disp = np.linalg.norm((y_true[-1, :] - y_prediction[-1, :])) 124 | return last_disp 125 | 126 | def get_hausdorff(y_true, y_prediction): 127 | ''' 128 | Here is the directed Hausdorff distance, but it computes both directions and output the larger value 129 | ''' 130 | Hausdorff = max(directed_hausdorff(y_true, y_prediction)[0], directed_hausdorff(y_prediction, y_true)[0]) 131 | return Hausdorff 132 | 133 | def get_speeddev(y_true, y_prediction): 134 | if len(y_true) == 1: 135 | return 0 136 | else: 137 | speed_dev = 0.0 138 | for t in range(len(y_true)-1): 139 | speed_t = np.linalg.norm(y_true[t+1] - y_true[t]) 140 | speed_p = np.linalg.norm(y_prediction[t+1] - y_prediction[t]) 141 | speed_dev += abs(speed_t - speed_p) 142 | speed_dev /= (len(y_true)-1) 143 | return speed_dev 144 | 145 | def get_headerror(y_true, y_prediction): 146 | if len(y_prediction) == 1: 147 | return 0 148 | else: 149 | heading_error = 0.0 150 | for t in range(len(y_true)-1): 151 | xcoor_t = y_true[t+1, 0] - y_true[t, 0] 152 | ycoor_t = y_true[t+1, 1] - y_true[t, 1] 153 | angle_t = np.arctan2(ycoor_t, xcoor_t) 154 | xcoor_p = y_prediction[t+1, 0] - y_prediction[t, 0] 155 | ycoor_p = y_prediction[t+1, 1] - y_prediction[t, 1] 156 | angle_p = np.arctan2(ycoor_p, xcoor_p) 157 | angle = np.rad2deg((abs(angle_t - angle_p)) % (np.pi)) 158 | heading_error += angle 159 | heading_error /= len(y_true)-1 160 | return heading_error 161 | 162 | _, term_width = os.popen('stty size', 'r').read().split() 163 | term_width = int(term_width) 164 | 165 | TOTAL_BAR_LENGTH = 65. 166 | last_time = time.time() 167 | begin_time = last_time 168 | 169 | def progress_bar(current, total, msg=None): 170 | global last_time, begin_time 171 | if current == 0: 172 | begin_time = time.time() # Reset for new bar. 173 | 174 | cur_len = int(TOTAL_BAR_LENGTH*current/total) 175 | rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1 176 | 177 | sys.stdout.write(' [') 178 | for i in range(cur_len): 179 | sys.stdout.write('=') 180 | sys.stdout.write('>') 181 | for i in range(rest_len): 182 | sys.stdout.write('.') 183 | sys.stdout.write(']') 184 | 185 | cur_time = time.time() 186 | step_time = cur_time - last_time 187 | last_time = cur_time 188 | tot_time = cur_time - begin_time 189 | 190 | L = [] 191 | L.append(' Step: %s' % format_time(step_time)) 192 | L.append(' | Tot: %s' % format_time(tot_time)) 193 | if msg: 194 | L.append(' | ' + msg) 195 | 196 | msg = ''.join(L) 197 | sys.stdout.write(msg) 198 | for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3): 199 | sys.stdout.write(' ') 200 | 201 | # Go back to the center of the bar. 202 | for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2): 203 | sys.stdout.write('\b') 204 | sys.stdout.write(' %d/%d ' % (current+1, total)) 205 | 206 | if current < total-1: 207 | sys.stdout.write('\r') 208 | else: 209 | sys.stdout.write('\n') 210 | sys.stdout.flush() 211 | 212 | def format_time(seconds): 213 | days = int(seconds / 3600/24) 214 | seconds = seconds - days*3600*24 215 | hours = int(seconds / 3600) 216 | seconds = seconds - hours*3600 217 | minutes = int(seconds / 60) 218 | seconds = seconds - minutes*60 219 | secondsf = int(seconds) 220 | seconds = seconds - secondsf 221 | millis = int(seconds*1000) 222 | 223 | f = '' 224 | i = 1 225 | if days > 0: 226 | f += str(days) + 'D' 227 | i += 1 228 | if hours > 0 and i <= 2: 229 | f += str(hours) + 'h' 230 | i += 1 231 | if minutes > 0 and i <= 2: 232 | f += str(minutes) + 'm' 233 | i += 1 234 | if secondsf > 0 and i <= 2: 235 | f += str(secondsf) + 's' 236 | i += 1 237 | if millis > 0 and i <= 2: 238 | f += str(millis) + 'ms' 239 | i += 1 240 | if f == '': 241 | f = '0ms' 242 | return f 243 | 244 | 245 | def input2tensor(traj, occu, device): 246 | traj = torch.tensor(traj, dtype=torch.float) 247 | occu = torch.tensor(occu, dtype=torch.float).permute(0, 1, 4, 2, 3) 248 | 249 | return traj.to(device), occu.to(device) 250 | 251 | 252 | def get_lr(optimizer): 253 | for param_group in optimizer.param_groups: 254 | return param_group['lr'] -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | '''Train DCENet with PyTorch''' 2 | # from __future__ import print_function 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.optim as optim 7 | 8 | from torch.utils.data import DataLoader 9 | 10 | import os 11 | import json 12 | import neptune 13 | import argparse 14 | import numpy as np 15 | 16 | from loader import * 17 | from utils.plots import * 18 | from utils.utils import * 19 | from utils.collision import * 20 | from utils.datainfo import DataInfo 21 | from utils.ranking import gauss_rank 22 | 23 | from models import DCENet 24 | from loss import DCENetLoss 25 | 26 | 27 | 28 | def main(): 29 | 30 | # ================= Arguments ================ # 31 | parser = argparse.ArgumentParser(description='PyTorch Knowledge Distillation') 32 | parser.add_argument('--gpu', type=str, default="4", help='gpu id') 33 | parser.add_argument('--config', type=str, default="config", help='.json') 34 | args = parser.parse_args() 35 | 36 | # ================= Device Setup ================ # 37 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 38 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 39 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 40 | 41 | # ================= Config Load ================ # 42 | with open('config/' + args.config) as config_file: 43 | config = json.load(config_file) 44 | 45 | # ================= Neptune Setup ================ # 46 | if config['neptune']: 47 | neptune.init('seongjulee/DCENet', api_token=config["neptune_token"]) # username/project-name, api_token=token from neptune 48 | neptune.create_experiment(name='EXP', params=config) # name=project name (anything is ok), params=parameter list (json format) 49 | neptune.append_tag(args.config) # neptune tag (str or string list) 50 | 51 | # ================= Model Setup ================ # 52 | model = nn.DataParallel(DCENet(config)).to(device) if len(args.gpu.split(',')) > 1 else DCENet(config).to(device) 53 | 54 | # ================= Loss Function ================ # 55 | criterion = DCENetLoss(config) 56 | 57 | # ================= Optimizer Setup ================ # 58 | optimizer = optim.Adam(model.parameters(), lr=config['lr'], betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-6, amsgrad=False) 59 | 60 | # ================= Data Loader ================ # 61 | datalist = DataInfo() 62 | train_datalist = datalist.train_merged 63 | print('Train data list', train_datalist) 64 | 65 | test_datalist = datalist.train_biwi 66 | print('Test data list', test_datalist) 67 | 68 | np.random.seed(10) 69 | offsets, traj_data, occupancy = load_data(config, train_datalist, datatype="train") 70 | trainval_split = np.random.rand(len(offsets)) < config['split'] 71 | 72 | train_x = offsets[trainval_split, :config['obs_seq'] - 1, 4:6] 73 | train_occu = occupancy[trainval_split, :config['obs_seq'] - 1, ..., :config['enviro_pdim'][-1]] 74 | train_y = offsets[trainval_split, config['obs_seq'] - 1:, 4:6] 75 | train_y_occu = occupancy[trainval_split, config['obs_seq'] - 1:, ..., :config['enviro_pdim'][-1]] 76 | 77 | val_x = offsets[~trainval_split, :config['obs_seq'] - 1, 4:6] 78 | val_occu = occupancy[~trainval_split, :config['obs_seq'] - 1, ..., :config['enviro_pdim'][-1]] 79 | val_y = offsets[~trainval_split, config['obs_seq'] - 1:, 4:6] 80 | val_y_occu = occupancy[~trainval_split, config['obs_seq'] - 1:, ..., :config['enviro_pdim'][-1]] 81 | 82 | print("%.0f trajectories for training\n %.0f trajectories for valiadation" %(train_x.shape[0], val_x.shape[0])) 83 | 84 | test_offsets, test_trajs, test_occupancy = load_data(config, test_datalist, datatype="test") 85 | test_x = test_offsets[:, :config['obs_seq'] - 1, 4:6] 86 | test_occu = test_occupancy[:, :config['obs_seq'] - 1, ..., :config['enviro_pdim'][-1]] 87 | last_obs_test = test_offsets[:, config['obs_seq'] - 2, 2:4] 88 | y_truth = test_offsets[:, config['obs_seq'] - 1:, :4] 89 | xy_truth = test_offsets[:, :, :4] 90 | 91 | print('test_trajs', test_trajs.shape) 92 | 93 | print("%.0f trajectories for testing" % (test_x.shape[0])) 94 | 95 | train_dataset = TrajDataset(x=train_x, x_occu=train_occu, y=train_y, y_occu=train_y_occu, mode='train') 96 | train_loader = DataLoader(dataset=train_dataset, batch_size=config["batch_size"], shuffle=True, num_workers=4) 97 | 98 | val_dataset = TrajDataset(x=val_x, x_occu=val_occu, y=val_y, y_occu=val_y_occu, mode='val') 99 | val_loader = DataLoader(dataset=val_dataset, batch_size=config["batch_size"], shuffle=False, num_workers=4) 100 | 101 | # test_dataset = TrajDataset(x=test_x, x_occu=test_occu, y=y_truth, y_occu=None, mode='test') 102 | # test_loader = DataLoader(dataset=test_dataset, batch_size=config["batch_size"], shuffle=False, num_workers=4) 103 | 104 | # ================= Training Loop ================ # 105 | early_stopping = EarlyStopping(patience=config['patience'], verbose=True, filename=args.config.split('/')[-1].replace('.json', '.pth')) 106 | for epoch in range(config['max_epochs']): 107 | train_one_epoch(config, epoch, device, model, optimizer, criterion, train_loader) 108 | val_loss = evaluate(config, device, model, optimizer, criterion, val_loader) 109 | early_stopping(val_loss, model) 110 | if early_stopping.early_stop: 111 | print("Early stopping") 112 | break 113 | 114 | # ================= Test ================ # 115 | model.load_state_dict(torch.load(os.path.join('checkpoints', args.config.split('/')[-1].replace('.json', '.pth')))) 116 | model.eval() 117 | with torch.no_grad(): 118 | test_x, test_occu = input2tensor(test_x, test_occu, device) 119 | x_latent = model.encoder_x(test_x, test_occu) 120 | predictions = [] 121 | for i, x_ in enumerate(x_latent): 122 | last_pos = last_obs_test[i] 123 | x_ = x_.view(1, -1) 124 | for i in range(config['num_pred']): 125 | y_p = model.decoder(x_, train=False) 126 | y_p_ = np.concatenate(([last_pos], np.squeeze(y_p.cpu().numpy())), axis=0) 127 | y_p_sum = np.cumsum(y_p_, axis=0) 128 | predictions.append(y_p_sum[1:, :]) 129 | 130 | predictions = np.reshape(predictions, [-1, config['num_pred'], config['pred_seq'], 2]) 131 | 132 | print('Predicting done!') 133 | print(predictions.shape) 134 | plot_pred(xy_truth, predictions) 135 | # Get the errors for ADE, DEF, Hausdorff distance, speed deviation, heading error 136 | print("\nEvaluation results @top%.0f" % config['num_pred']) 137 | errors = get_errors(y_truth, predictions) 138 | check_collision(y_truth) 139 | 140 | ## Get the first time prediction by g 141 | ranked_prediction = [] 142 | for prediction in predictions: 143 | ranks = gauss_rank(prediction) 144 | ranked_prediction.append(prediction[np.argmax(ranks)]) 145 | ranked_prediction = np.reshape(ranked_prediction, [-1, 1, config['pred_seq'], 2]) 146 | print("\nEvaluation results for most-likely predictions") 147 | ranked_errors = get_errors(y_truth, ranked_prediction) 148 | 149 | 150 | # Function for one epoch training 151 | def train_one_epoch(config, epoch, device, model, optimizer, criterion, loader): 152 | print('\nEpoch: %d' % epoch) 153 | model.train() 154 | train_total, train_loss = 0, 0 155 | 156 | for batch_idx, (x, x_occu, y, y_occu) in enumerate(loader): 157 | x, x_occu, y, y_occu = x.to(device), x_occu.to(device), y.to(device), y_occu.to(device) 158 | optimizer.zero_grad() 159 | 160 | y_pred, mu, log_var = model(x, x_occu, y, y_occu, train=True) 161 | loss = criterion(mu, log_var, y_pred, y) 162 | 163 | loss.backward() 164 | optimizer.step() 165 | 166 | # train_ade += ade * x.size(0) 167 | # train_fde += fde * x.size(0) 168 | train_total += x.size(0) 169 | train_loss += loss.item() * x.size(0) 170 | 171 | if config['neptune']: 172 | # neptune.log_metric('train_batch_ADE', ade) 173 | # neptune.log_metric('train_batch_FDE', fde) 174 | neptune.log_metric('train_batch_Loss', loss.item()) 175 | 176 | # progress_bar(batch_idx, len(loader), 'Lr: %.4e | Loss: %.3f | ADE[m]: %.3f | FDE[m]: %.3f' 177 | # % (get_lr(optimizer), train_loss / train_total, train_ade / train_total, train_fde / train_total)) 178 | progress_bar(batch_idx, len(loader), 'Lr: %.4e | Loss: %.3f' % (get_lr(optimizer), train_loss / train_total)) 179 | 180 | 181 | # Function for validation 182 | @torch.no_grad() 183 | def evaluate(config, device, model, optimizer, criterion, loader): 184 | model.eval() 185 | # eval_ade, eval_fde, eval_total = 0, 0, 0 186 | eval_total, eval_loss = 0, 0 187 | 188 | for batch_idx, (x, x_occu, y, y_occu) in enumerate(loader): 189 | x, x_occu, y, y_occu = x.to(device), x_occu.to(device), y.to(device), y_occu.to(device) 190 | 191 | y_pred, mu, log_var = model(x, x_occu, y, y_occu, train=True) 192 | loss = criterion(mu, log_var, y_pred, y) 193 | 194 | eval_total += x.size(0) 195 | eval_loss += loss.item() * x.size(0) 196 | 197 | progress_bar(batch_idx, len(loader), 'Lr: %.4e | Loss: %.3f' % (get_lr(optimizer), eval_loss / eval_total)) 198 | # progress_bar(batch_idx, len(loader), 'Lr: %.4e | ADE[m]: %.3f | FDE[m]: %.3f' 199 | # % (get_lr(optimizer), eval_ade / eval_total, eval_fde / eval_total)) 200 | 201 | if config['neptune']: 202 | neptune.log_metric('val_Loss', eval_loss / eval_total) 203 | 204 | # neptune.log_metric('{}_ADE'.format(loader.dataset.mode), eval_ade / eval_total) 205 | # neptune.log_metric('{}_FDE'.format(loader.dataset.mode), eval_fde / eval_total) 206 | 207 | return eval_loss / eval_total 208 | 209 | 210 | if __name__ == "__main__": 211 | main() 212 | -------------------------------------------------------------------------------- /models.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | 7 | class DCENet(nn.Module): 8 | def __init__(self, config): 9 | super(DCENet, self).__init__() 10 | 11 | self.encoder_x = Encoder(config) 12 | self.encoder_y = Encoder(config) 13 | self.decoder = CVAE(config) 14 | 15 | def forward(self, x, x_occu, y=None, y_occu=None, train=True): 16 | x_encoded_dense = self.encoder_x(x, x_occu) 17 | y_encoded_dense = None 18 | 19 | if train: 20 | y_encoded_dense = self.encoder_y(y, y_occu) 21 | 22 | # out for train: pred_traj, mu, log_var | out for test: pred_traj 23 | out = self.decoder(x_encoded_dense, y_encoded_dense, train=train) 24 | 25 | return out 26 | 27 | 28 | class Encoder(nn.Module): 29 | def __init__(self, config, x_or_y='x'): 30 | super(Encoder, self).__init__() 31 | 32 | self.config = config 33 | 34 | # For trajectory 35 | self.traj_conv1 = nn.Conv1d(in_channels=2, out_channels=config['n_hidden'] // 16, kernel_size=3, stride=1, padding=1) 36 | self.traj_fc = nn.Sequential( 37 | nn.Linear(in_features=config['n_hidden'] // 16, out_features=config['n_hidden'] // 8), 38 | nn.ReLU() 39 | ) 40 | self.traj_pos_encode = PositionalEncoding( 41 | d_in=config['n_hidden'] // 8, 42 | d_model=config['{}_encoder_dim'.format(x_or_y)] 43 | ) 44 | self.traj_transformer_encoder_layer = nn.TransformerEncoderLayer( 45 | d_model=config['{}_encoder_dim'.format(x_or_y)], 46 | nhead=config['{}_encoder_head'.format(x_or_y)], 47 | dim_feedforward=config['{}_encoder_dim'.format(x_or_y)] 48 | ) 49 | self.traj_transformer_encoder = nn.TransformerEncoder( 50 | encoder_layer=self.traj_transformer_encoder_layer, 51 | num_layers=config['{}_encoder_layers'.format(x_or_y)] 52 | ) 53 | self.traj_avg_pool = nn.Sequential( 54 | nn.AdaptiveAvgPool1d(1), 55 | nn.Flatten() 56 | ) 57 | 58 | # For dynamic map 59 | self.occu_model = nn.Sequential( 60 | nn.ZeroPad2d((0, 1, 0, 1)), # Tensorflow padding 'SAME' option 61 | nn.Conv2d(in_channels=3, out_channels=6, kernel_size=2, stride=1, padding=0), 62 | nn.ReLU(), 63 | nn.ZeroPad2d((0, 1, 0, 1)), # Tensorflow padding 'SAME' option 64 | nn.MaxPool2d(kernel_size=(2, 2), stride=1, padding=0), 65 | nn.Dropout(p=config['o_drop']), 66 | nn.Flatten() 67 | ) 68 | self.occu_time_distributed = TimeDistributed(self.occu_model, tdim=1) 69 | self.occu_pos_encode = PositionalEncoding( 70 | d_in=6144, 71 | d_model=config['occu_encoder_{}_dim'.format(x_or_y)] 72 | ) 73 | self.occu_transformer_encoder_layer = nn.TransformerEncoderLayer( 74 | d_model=config['occu_encoder_{}_dim'.format(x_or_y)], 75 | nhead=config['occu_encoder_{}_head'.format(x_or_y)], 76 | dim_feedforward=config['occu_encoder_{}_dim'.format(x_or_y)] 77 | ) 78 | self.occu_transformer_encoder = nn.TransformerEncoder( 79 | encoder_layer=self.occu_transformer_encoder_layer, 80 | num_layers=config['occu_encoder_{}_layers'.format(x_or_y)] 81 | ) 82 | self.occu_lstm = nn.LSTM( 83 | input_size=config['occu_encoder_{}_dim'.format(x_or_y)], 84 | hidden_size=config['hidden_size'], 85 | num_layers=1, 86 | batch_first=True, 87 | ) 88 | self.occu_dropout = nn.Dropout(p=config['s_drop']) 89 | 90 | # For encoding 91 | self.encode_fc = nn.Sequential( 92 | nn.Linear(in_features=config['{}_encoder_dim'.format(x_or_y)] + config['hidden_size'], out_features=config['encoder_dim']), 93 | nn.ReLU() 94 | ) 95 | 96 | # print(sum(p.numel() for p in self.occu_time_distributed.parameters() if p.requires_grad)) 97 | 98 | def init_hidden(self, x): 99 | h0 = torch.zeros((1, x.size(0), self.config['hidden_size'])).cuda() 100 | c0 = torch.zeros((1, x.size(0), self.config['hidden_size'])).cuda() 101 | 102 | return h0, c0 103 | 104 | def forward(self, traj, dmap): 105 | traj = traj.transpose(1, 2) 106 | traj = self.traj_conv1(traj) 107 | traj = traj.transpose(1, 2) 108 | traj = self.traj_fc(traj) 109 | traj = traj.transpose(0, 1) # (L, B, H) 110 | traj = self.traj_pos_encode(traj) 111 | traj = self.traj_transformer_encoder(traj) 112 | traj = traj.transpose(0, 1).transpose(1, 2) 113 | traj = self.traj_avg_pool(traj) 114 | 115 | dmap = self.occu_time_distributed(dmap) 116 | dmap = dmap.transpose(0, 1) 117 | dmap = self.occu_pos_encode(dmap) 118 | dmap = self.occu_transformer_encoder(dmap) 119 | dmap = dmap.transpose(0, 1) 120 | dmap_hidden = self.init_hidden(dmap) 121 | dmap_out, dmap_hidden = self.occu_lstm(dmap, dmap_hidden) 122 | dmap = dmap_out[:, -1, :] 123 | dmap = self.occu_dropout(dmap) 124 | 125 | out = torch.cat((traj, dmap), dim=1) 126 | out = self.encode_fc(out) 127 | 128 | return out 129 | 130 | 131 | class PositionalEncoding(nn.Module): 132 | 133 | def __init__(self, d_in, d_model, dropout=0.1, max_len=5000): 134 | super(PositionalEncoding, self).__init__() 135 | self.dropout = nn.Dropout(p=dropout) 136 | self.fc = nn.Sequential( 137 | nn.Linear(in_features=d_in, out_features=d_model), 138 | nn.ReLU() 139 | ) 140 | pe = torch.zeros(max_len, d_model) 141 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 142 | div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) 143 | pe[:, 0::2] = torch.sin(position * div_term) 144 | pe[:, 1::2] = torch.cos(position * div_term) 145 | pe = pe.unsqueeze(0).transpose(0, 1) 146 | self.register_buffer('pe', pe) 147 | 148 | def forward(self, x): 149 | x = self.fc(x) 150 | x = x + self.pe[:x.size(0), :] 151 | return self.dropout(x) 152 | 153 | 154 | class TimeDistributed(nn.Module): 155 | "Applies a module over tdim identically for each step" 156 | def __init__(self, module, low_mem=False, tdim=1): 157 | super(TimeDistributed, self).__init__() 158 | self.module = module 159 | self.low_mem = low_mem 160 | self.tdim = tdim 161 | 162 | def forward(self, *args, **kwargs): 163 | "input x with shape:(bs,seq_len,channels,width,height)" 164 | if self.low_mem or self.tdim!=1: 165 | return self.low_mem_forward(*args) 166 | else: 167 | #only support tdim=1 168 | inp_shape = args[0].shape 169 | bs, seq_len = inp_shape[0], inp_shape[1] 170 | out = self.module(*[x.reshape(bs*seq_len, *x.shape[2:]) for x in args], **kwargs) 171 | out_shape = out.shape 172 | return out.view(bs, seq_len,*out_shape[1:]) 173 | 174 | def low_mem_forward(self, *args, **kwargs): 175 | "input x with shape:(bs,seq_len,channels,width,height)" 176 | tlen = args[0].shape[self.tdim] 177 | args_split = [torch.unbind(x, dim=self.tdim) for x in args] 178 | out = [] 179 | for i in range(tlen): 180 | out.append(self.module(*[args[i] for args in args_split]), **kwargs) 181 | return torch.stack(out,dim=self.tdim) 182 | def __repr__(self): 183 | return f'TimeDistributed({self.module})' 184 | 185 | 186 | class CVAE(nn.Module): 187 | def __init__(self, config): 188 | super(CVAE, self).__init__() 189 | self.z_dim = config['z_dim'] 190 | self.pred_seq = config['pred_seq'] 191 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 192 | 193 | # For concatenated input 194 | self.xy_encoded_fc1 = nn.Sequential( 195 | nn.Linear(in_features=config['encoder_dim'] * 2, out_features=config['n_hidden']), 196 | nn.ReLU() 197 | ) 198 | self.xy_encoded_fc2 = nn.Sequential( 199 | nn.Linear(in_features=config['n_hidden'], out_features=config['n_hidden'] // 2), 200 | nn.ReLU() 201 | ) 202 | self.mu = nn.Linear(config['n_hidden'] // 2, config['z_dim']) 203 | self.log_var = nn.Linear(config['n_hidden'] // 2, config['z_dim']) 204 | 205 | # decoder part 206 | self.z_fc = nn.Sequential( 207 | nn.Linear(config['z_dim'] + config['encoder_dim'], config['n_hidden'] // 2), 208 | nn.ReLU() 209 | ) 210 | self.z_lstm = nn.LSTM(input_size=config['n_hidden'] // 2, hidden_size=config['z_decoder_dim'], num_layers=1, batch_first=True) 211 | self.z_dropout = nn.Dropout(p=config['z_drop']) 212 | 213 | self.y_decoder_model = nn.Conv1d(in_channels=1, out_channels=2, kernel_size=config['z_decoder_dim'], stride=config['z_decoder_dim']) 214 | self.y_decoder = TimeDistributed(self.y_decoder_model, tdim=1) 215 | 216 | def encoder(self, x, y): 217 | concat_input = torch.cat([x, y], 1) 218 | h = self.xy_encoded_fc1(concat_input) 219 | h = self.xy_encoded_fc2(h) 220 | return self.mu(h), self.log_var(h) 221 | 222 | def sampling(self, mu, log_var): 223 | std = torch.exp(0.5 * log_var) 224 | eps = torch.randn_like(std) 225 | return eps.mul(std).add(mu) # return z sample 226 | 227 | def decoder(self, z, x): 228 | concat_input = torch.cat([z, x], 1) 229 | h = self.z_fc(concat_input) 230 | h = h.view(h.size(0), 1, -1).repeat(1, self.pred_seq, 1) # 12 for predicted sequence length 231 | h, _ = self.z_lstm(h) 232 | h = self.z_dropout(torch.tanh(h)) 233 | h = h.view(h.size(0), h.size(1), 1, -1) 234 | out = self.y_decoder(h).view(h.size(0), h.size(1), -1) 235 | return out 236 | 237 | def forward(self, x, y=None, train=True): 238 | if train: 239 | mu, log_var = self.encoder(x, y) 240 | z = self.sampling(mu, log_var) 241 | return self.decoder(z, x), mu, log_var 242 | else: 243 | z = torch.randn((x.size(0), self.z_dim)).to(self.device) 244 | return self.decoder(z, x) 245 | --------------------------------------------------------------------------------