def predict(paths, predict_all=False): multimodal_outputs = {} neighbours_tracks = [] ## Primary Prediction if not predict_all: paths = paths[0:1] for i, path in enumerate(paths): path = paths[i] initial_state_mean = [path[0].x, 0, path[0].y, 0] transition_matrix = [[1, 1, 0, 0], [0, 1, 0, 0], [0, 0, 1, 1], [0, 0, 0, 1]] observation_matrix = [[1, 0, 0, 0], [0, 0, 1, 0]] kf = pykalman.KalmanFilter(transition_matrices=transition_matrix, observation_matrices=observation_matrix, transition_covariance=1e-5 * np.eye(4), observation_covariance=0.05**2 * np.eye(2), initial_state_mean=initial_state_mean) # kf.em([(r.x, r.y) for r in path[:9]], em_vars=['transition_matrices', # 'observation_matrices']) kf.em([(r.x, r.y) for r in path[:9]]) observed_states, _ = kf.smooth([(r.x, r.y) for r in path[:9]]) # prepare predictions frame_diff = path[1].frame - path[0].frame first_frame = path[8].frame + frame_diff ped_id = path[8].pedestrian # sample predictions (first sample corresponds to last state) # average 5 sampled predictions predictions = None for _ in range(5): _, pred = kf.sample(12 + 1, initial_state=observed_states[-1]) if predictions is None: predictions = pred else: predictions += pred predictions /= 5.0 #write if i == 0: primary_track = [ trajnettools.TrackRow(first_frame + j * frame_diff, ped_id, x, y) for j, (x, y) in enumerate(predictions[1:]) ] else: neighbours_tracks.append([ trajnettools.TrackRow(first_frame + j * frame_diff, ped_id, x, y) for j, (x, y) in enumerate(predictions[1:]) ]) ## Unimodal Ouput multimodal_outputs[0] = primary_track, neighbours_tracks return multimodal_outputs
def __call__(self, paths, n_predict=5, modes=1): self.model.eval() observed_path = paths[0] ped_id = observed_path[0].pedestrian ped_id_ = [] for j in range(len(paths)): ped_id_.append(paths[j][0].pedestrian) frame_diff = observed_path[1].frame - observed_path[0].frame first_frame = observed_path[4].frame + frame_diff with torch.no_grad(): xy = trajnettools.Reader.paths_to_xy(paths) xy = drop_distant(xy, r=10.0) xy = torch.Tensor(xy) #.to(self.device) multimodal_outputs = {} for np in range(modes): _, output_scenes = self.model(xy[:5], n_predict=n_predict) outputs = output_scenes[-n_predict:, 0] output_scenes = output_scenes[-n_predict:] output_primary = [ trajnettools.TrackRow(first_frame + i * frame_diff, ped_id, outputs[i, 0], outputs[i, 1], 0) for i in range(len(outputs)) ] output_all = [[ trajnettools.TrackRow(first_frame + i * frame_diff, ped_id_[j], output_scenes[i, j, 0], output_scenes[i, j, 1], 0) for i in range(len(outputs)) ] for j in range(1, output_scenes.shape[1])] multimodal_outputs[np] = [output_primary, output_all] return multimodal_outputs
def rotate_path(path, theta): ct = math.cos(theta) st = math.sin(theta) return [ trajnettools.TrackRow(r.frame, r.pedestrian, ct * r.x + st * r.y, -st * r.x + ct * r.y) for r in path ]
def main(args): ## List of test files (.json) inside the test folder (waiting to be predicted by the prediction model) datasets = sorted([f for f in os.listdir(args.data.replace('_pred', '')) if not f.startswith('.') and f.endswith('.ndjson')]) ## Handcrafted Baselines (if required to compare) if args.kf: args.output.append('/kf.pkl') if args.sf: args.output.append('/sf.pkl') args.output.append('/sf_opt.pkl') if args.orca: args.output.append('/orca.pkl') args.output.append('/orca_opt.pkl') ## Extract Model names from arguments and create its own folder in 'test_pred' for storing predictions ## WARNING: If Model predictions already exist from previous run, this process SKIPS WRITING for model in args.output: model_name = model.split('/')[-1].replace('.pkl', '') ## Check if model predictions already exist if not os.path.exists(args.data): os.makedirs(args.data) if not os.path.exists(args.data + model_name): os.makedirs(args.data + model_name) else: continue ## Start writing predictions in dataset/test_pred for dataset in datasets: # Model's name name = dataset.replace(args.data.replace('_pred', '') + 'test/', '') # Copy observations from test folder into test_pred folder shutil.copyfile(args.data.replace('_pred', '') + name, args.data + '{}/{}'.format(model_name, name)) print('processing ' + name) # Read Scenes from 'test' folder reader = trajnettools.Reader(args.data.replace('_pred', '') + dataset, scene_type='paths') scenes = [s for s in reader.scenes()] # Loading the APPROPRIATE model ## Keep Adding Different Models to this List print("Model Name: ", model_name) if model_name == 'kf': print("Kalman") predictor = trajnetbaselines.classical.kalman.predict elif model_name == 'sf' or model_name == 'sf_opt': print("Social Force") predictor = trajnetbaselines.classical.socialforce.predict elif model_name == 'orca' or model_name == 'orca_opt': print("ORCA") predictor = trajnetbaselines.classical.orca.predict elif 'sgan' in model_name: print("SGAN") predictor = trajnetbaselines.sgan.SGANPredictor.load(model) # On CPU device = torch.device('cpu') predictor.model.to(device) else: print("LSTM") predictor = trajnetbaselines.lstm.LSTMPredictor.load(model) # On CPU device = torch.device('cpu') predictor.model.to(device) # Get the model prediction and write them in corresponding test_pred file """ VERY IMPORTANT: Prediction Format The predictor function should output a dictionary. The keys of the dictionary should correspond to the prediction modes. ie. predictions[0] corresponds to the first mode. predictions[m] corresponds to the m^th mode.... Multimodal predictions! Each modal prediction comprises of primary prediction and neighbour (surrrounding) predictions i.e. predictions[m] = [primary_prediction, neigh_predictions] Note: Return [primary_prediction, []] if model does not provide neighbour predictions Shape of primary_prediction: Tensor of Shape (Prediction length, 2) Shape of Neighbour_prediction: Tensor of Shape (Prediction length, n_tracks - 1, 2). (See LSTMPredictor.py for more details) """ with open(args.data + '{}/{}'.format(model_name, name), "a") as myfile: for scene_id, paths in scenes: ## Extract 1) first_frame, 2) frame_diff 3) ped_ids for writing predictions observed_path = paths[0] frame_diff = observed_path[1].frame - observed_path[0].frame first_frame = observed_path[args.obs_length-1].frame + frame_diff ped_id = observed_path[0].pedestrian ped_id_ = [] for j, _ in enumerate(paths[1:]): ## Only need neighbour ids ped_id_.append(paths[j+1][0].pedestrian) ## For each scene, get predictions if model_name == 'sf_opt': predictions = predictor(paths, sf_params=[0.5, 1.0, 0.1], n_predict=args.pred_length, obs_length=args.obs_length) ## optimal sf_params elif model_name == 'orca_opt': predictions = predictor(paths, orca_params=[0.25, 1.0, 0.3], n_predict=args.pred_length, obs_length=args.obs_length) ## optimal orca_params else: predictions = predictor(paths, n_predict=args.pred_length, obs_length=args.obs_length) for m in range(len(predictions)): prediction, neigh_predictions = predictions[m] ## Write Primary for i in range(len(prediction)): # print(i) track = trajnettools.TrackRow(first_frame + i * frame_diff, ped_id, prediction[i, 0].item(), prediction[i, 1].item(), m, scene_id) myfile.write(trajnettools.writers.trajnet(track)) myfile.write('\n') ## Write Neighbours (if non-empty) for n in range(neigh_predictions.shape[1]): neigh = neigh_predictions[:, n] for j in range(len(neigh)): track = trajnettools.TrackRow(first_frame + j * frame_diff, ped_id_[n], neigh[j, 0].item(), neigh[j, 1].item(), m, scene_id) myfile.write(trajnettools.writers.trajnet(track)) myfile.write('\n') print('')
def predict(input_paths, predict_all=False): def init_states(input_paths, start_frame): initial_state = [] for i in range(len(input_paths)): path = input_paths[i] past_path = [t for t in path if t.frame <= start_frame] past_frames = [t.frame for t in path if t.frame <= start_frame] l = len(past_path) ## To consider agent or not consider. if (start_frame in past_frames) and l >= 4: curr = past_path[-1] prev = past_path[-4] [v_x, v_y] = vel_state(prev, curr, 3) if np.linalg.norm([v_x, v_y]) < 1e-6: continue [d_x, d_y] = dest_state(past_path, l-1) if np.linalg.norm([curr.x - d_x, curr.y - d_y]) < 1e-6: continue initial_state.append([curr.x, curr.y, v_x, v_y, d_x, d_y]) return np.array(initial_state) def vel_state(prev, curr, l): diff = np.array([curr.x - prev.x, curr.y - prev.y]) theta = np.arctan2(diff[1], diff[0]) speed = np.linalg.norm(diff) / (l * 0.4) return [speed*np.cos(theta), speed*np.sin(theta)] def dest_state(path, l): x = [t.x for t in path] y = [t.y for t in path] time = [i for i in range(l+1)] f = interp1d(x=time, y=[x, y], fill_value='extrapolate') return f(time[-1] + 12) multimodal_outputs = {} primary = input_paths[0] neighbours_tracks = [] frame_diff = primary[1].frame - primary[0].frame start_frame = primary[8].frame first_frame = primary[8].frame + frame_diff # initialize initial_state = init_states(input_paths, start_frame) # run s = socialforce.Simulator(initial_state) states = np.stack([s.step().state.copy() for _ in range(12)]) ## states : 12 x num_ped x 7 # predictions for i in range(states.shape[1]): ped_id = input_paths[i][0].pedestrian if i == 0: primary_track = [trajnettools.TrackRow(first_frame + j * frame_diff, ped_id, x, y) for j, (x, y) in enumerate(states[:, i, 0:2])] else: neighbours_tracks.append([trajnettools.TrackRow(first_frame + j * frame_diff, ped_id, x, y) for j, (x, y) in enumerate(states[:, i, 0:2])]) ## Primary Prediction if not predict_all: neighbours_tracks = [] # Unimodal Prediction multimodal_outputs[0] = primary_track, neighbours_tracks return multimodal_outputs
def main(): parser = argparse.ArgumentParser() parser.add_argument('--data', default='trajdata', help='directory of data to test') parser.add_argument('--output', required=True, nargs='+', help='relative path to saved model') parser.add_argument('--obs_length', default=9, type=int, help='observation length') parser.add_argument('--pred_length', default=12, type=int, help='prediction length') parser.add_argument('--disable-write', action='store_true', help='disable writing new files') parser.add_argument('--disable-collision', action='store_true', help='disable collision metrics') parser.add_argument('--labels', required=False, nargs='+', help='labels of models') args = parser.parse_args() ## Path to the data folder name to predict args.data = 'DATA_BLOCK/' + args.data + '/' ## Test_pred: Folders for saving model predictions args.data = args.data + 'test_pred/' ## Writes to Test_pred ## Does this overwrite existing predictions? No. ### datasets = sorted([ f for f in os.listdir(args.data.replace('_pred', '')) if not f.startswith('.') and f.endswith('.ndjson') ]) ## Model names are passed as arguments for model in args.output: model_name = model.split('/')[-1].replace('.pkl', '') # Loading the appropriate model (currently written only for LSTMs) print("Model Name: ", model_name) predictor = trajnetbaselines.lstm.LSTMPredictor.load(model) # On CPU device = torch.device('cpu') predictor.model.to(device) total_scenes = 0 average = 0 final = 0 topk_average = 0 topk_final = 0 ## Start writing in dataset/test_pred for dataset in datasets: # Model's name name = dataset.replace( args.data.replace('_pred', '') + 'test/', '') # Copy file from test into test/train_pred folder print('processing ' + name) if 'collision_test' in name: continue # Read file from 'test' reader = trajnettools.Reader(args.data.replace('_pred', '') + dataset, scene_type='paths') scenes = [s for _, s in reader.scenes()] reader_gt = trajnettools.Reader( args.data.replace('_pred', '_private') + dataset, scene_type='paths') scenes_gt = [s for _, s in reader_gt.scenes()] scenes_id_gt = [i for i, _ in reader_gt.scenes()] total_scenes += len(scenes_gt) for i, paths in enumerate(scenes): ground_truth = scenes_gt[i] predictions = predictor(paths, n_predict=args.pred_length, obs_length=args.obs_length) ## Considers only the First MODE prediction, neigh_predictions = predictions[0] ## Convert numpy array to Track Rows ## ## Extract 1) first_frame, 2) frame_diff 3) ped_ids for writing predictions observed_path = paths[0] frame_diff = observed_path[1].frame - observed_path[0].frame first_frame = observed_path[args.obs_length - 1].frame + frame_diff ped_id = observed_path[0].pedestrian ## make Track Rows prediction = [ trajnettools.TrackRow(first_frame + i * frame_diff, ped_id, prediction[i, 0], prediction[i, 1], 0) for i in range(len(prediction)) ] primary_tracks = [ t for t in prediction if t.prediction_number == 0 ] # frame_gt = [t.frame for t in ground_truth[0]][-args.pred_length:] frame_gt = [ t.frame for t in ground_truth[0] ][args.obs_length:args.obs_length + args.pred_length] frame_pred = [t.frame for t in primary_tracks] ## To verify if same scene if frame_gt != frame_pred: raise Exception('frame numbers are not consistent') average_l2 = trajnettools.metrics.average_l2( ground_truth[0][args.obs_length:args.obs_length + args.pred_length], primary_tracks, n_predictions=args.pred_length) final_l2 = trajnettools.metrics.final_l2( ground_truth[0][args.obs_length:args.obs_length + args.pred_length], primary_tracks) # aggregate FDE and ADE average += average_l2 final += final_l2 ## TODO # if len(predictions) > 2: # # print(predictions) # primary_tracks_all = [t for mode in predictions for t in predictions[mode][0]] # # import pdb # # pdb.set_trace() # topk_ade, topk_fde = trajnettools.metrics.topk(primary_tracks_all, ground_truth[0][args.obs_length:args.obs_length+args.pred_length], n_predictions=args.pred_length) # topk_average += topk_ade # topk_final += topk_fde ## Average ADE and FDE average /= total_scenes final /= total_scenes # ##Adding value to dict print('ADE: ', average) print('FDE: ', final)
def test_final_l2(): path1 = [trajnettools.TrackRow(0, 0, x, 0.0) for x in range(21)] path2 = [trajnettools.TrackRow(0, 0, x, 1.0) for x in range(21)] assert trajnettools.metrics.final_l2(path1, path2) == 1.0
def test_collision(): path1 = [trajnettools.TrackRow(x, 0, x, x) for x in range(21)] path2 = [trajnettools.TrackRow(x, 0, x, 21 - x) for x in range(21)] assert trajnettools.metrics.collision(path1, path2) == 1
def predict(paths, n_obs, file_name=None, sample_rate=None, pixel_scale=None, scene_funcs=None, timestamp=None, store_image=0, center_line=None): constant_vel = 1 n_frames = len(paths[0]) path = paths[0] # prepare predictions frame_diff = path[1].frame - path[0].frame first_frame = path[n_obs].frame ped_id = path[-1].pedestrian if (constant_vel): path_np = np.zeros([15, 2]) j = 0 for r1 in path: path_np[j, 0] = r1.x path_np[j, 1] = r1.y j += 1 pos_temp = path_np[4, :] last_vel = path_np[4, :] - path_np[3, :] predictions = np.zeros([10, 2]) for idx in range(10): pos_temp = last_vel + pos_temp predictions[idx:idx + 1] = pos_temp scale = float(scene_funcs.pixel_scale_dict[file_name]) offset = scene_funcs.offset_dict[file_name] preds = predictions.copy() preds[:, 1] = scale * ( preds[:, 1] - offset[0] ) #second dimension is the longer axes, horizontal one preds[:, 0] = -scale * (preds[:, 0] - offset[1]) scene_violation = offroad_detector(preds, file_name, scene_funcs.image) return [ trajnettools.TrackRow(first_frame + i * frame_diff, ped_id, x, y) for i, (x, y) in enumerate(predictions) ], 0, scene_violation, 0 else: initial_state_mean = [path[0].x, 0, path[0].y, 0] transition_matrix = [[1, 1, 0, 0], [0, 1, 0, 0], [0, 0, 1, 1], [0, 0, 0, 1]] observation_matrix = [[1, 0, 0, 0], [0, 0, 1, 0]] kf = pykalman.KalmanFilter(transition_matrices=transition_matrix, observation_matrices=observation_matrix, transition_covariance=1e-5 * np.eye(4), observation_covariance=0.05**2 * np.eye(2), initial_state_mean=initial_state_mean) # kf.em([(r.x, r.y) for r in path[:9]], em_vars=['transition_matrices', # 'observation_matrices']) kf.em([(r.x, r.y) for r in path[:n_obs]]) observed_states, _ = kf.smooth([(r.x, r.y) for r in path[:n_obs]]) # sample predictions (first sample corresponds to last state) # average 5 sampled predictions predictions = None for _ in range(5): _, pred = kf.sample( n_frames - n_obs + 1, initial_state=observed_states[-1] ) # I don't know why but sven didn't consider the first sample so we have one more sample and in the last line we start from the sample 1. if predictions is None: predictions = pred else: predictions += pred predictions /= 5.0 preds = predictions.data[1:] scale = float(scene_funcs.pixel_scale_dict[file_name]) offset = scene_funcs.offset_dict[file_name] preds = preds.copy() preds[:, 1] = scale * ( preds[:, 1] - offset[0] ) #second dimension is the longer axes, horizontal one preds[:, 0] = -scale * (preds[:, 0] - offset[1]) scene_violation = offroad_detector(preds, file_name, scene_funcs.image) return [ trajnettools.TrackRow(first_frame + i * frame_diff, ped_id, x, y) for i, (x, y) in enumerate(predictions[1:]) ], 0, scene_violation, 0
def main(args, kf=False, sf=False): ## List of .json file inside the args.data (waiting to be predicted by the testing model) datasets = sorted([ f for f in os.listdir(args.data.replace('_pred', '')) if not f.startswith('.') and f.endswith('.ndjson') ]) if kf: args.output.append('/kf.pkl') if sf: args.output.append('/sf.pkl') ## Model names are passed as arguments for model in args.output: model_name = model.split('/')[-1].replace('.pkl', '') ## Make a directory in DATA_BLOCK which will contain the model outputs ## If model is already written, you skip writing if not os.path.exists(args.data): os.makedirs(args.data) if not os.path.exists(args.data + model_name): os.makedirs(args.data + model_name) else: continue ## Start writing in dataset/test_pred for dataset in datasets: # Model's name name = dataset.replace( args.data.replace('_pred', '') + 'test/', '') # Copy file from test into test/train_pred folder shutil.copyfile( args.data.replace('_pred', '') + name, args.data + '{}/{}'.format(model_name, name)) print('processing ' + name) # Read file from 'test' reader = trajnettools.Reader(args.data.replace('_pred', '') + dataset, scene_type='paths') scenes = [s for s in reader.scenes()] print("Model Name: ", model_name) # Load the model if model_name == 'kf': predictor = trajnetbaselines.kalman.predict elif model_name == 'sf': predictor = trajnetbaselines.socialforce.predict else: predictor = trajnetbaselines.lstm.LSTMPredictor.load(model) # On CPU device = torch.device('cpu') predictor.model.to(device) # Write the prediction with open(args.data + '{}/{}'.format(model_name, name), "a") as myfile: for scene_id, paths in scenes: predictions = predictor(paths) for m in range(len(predictions)): prediction, neigh_predictions = predictions[m] ## Write Primary for i in range(len(prediction)): track = trajnettools.TrackRow( prediction[i].frame, prediction[i].pedestrian, prediction[i].x.item(), prediction[i].y.item(), m, scene_id) myfile.write(trajnettools.writers.trajnet(track)) myfile.write('\n') ## Write Neighbours for n in range(len(neigh_predictions)): neigh = neigh_predictions[n] for j in range(len(neigh)): track = trajnettools.TrackRow( neigh[j].frame, neigh[j].pedestrian, neigh[j].x.item(), neigh[j].y.item(), m, scene_id) myfile.write( trajnettools.writers.trajnet(track)) myfile.write('\n') print('')
def __call__(self, paths, n_obs=None, file_name=None, sample_rate=None, pixel_scale=None, scene_funcs=None, store_image=0, center_line=None): if (n_obs is None): n_obs = self.n_obs n_pred = self.n_pred start_time = time.time() self.model.eval() self.model.resampling_dim = (38, 74) device = self.model.device self.model.to(device) ped_id = paths[0][0].pedestrian frame_diff = paths[0][1].frame - paths[0][0].frame first_frame = paths[0][n_obs].frame torch.backends.cudnn.enabled = False with torch.no_grad(): xy = trajnettools.Reader.paths_to_xy(paths) xy = torch.tensor(xy, dtype=torch.float32, device=device).unsqueeze(0) scale = float(scene_funcs.pixel_scale_dict[file_name]) offset = scene_funcs.offset_dict[file_name] xy_copy = xy.clone() xy_copy[:, :, :, 1] = scale * ( xy[:, :, :, 1] - offset[0]) # second dimension is the longer axes, horizontal one xy_copy[:, :, :, 0] = -scale * (xy[:, :, :, 0] - offset[1]) xy_copy_unrotated = xy_copy.clone() rotated_scene, resampled_scene, xy_copy, theta = scene_preprocess(xy_copy, [file_name], n_obs, self.model.resampling_dim, scene_funcs) pixel_scale = pixel_scale.to(device) center_line_rotated = augmentation.rotate_all_path_by_theta( center_line[file_name + '.txt'].to(device=device), xy_copy[:, n_obs - 1:n_obs, 0:1], theta, centerline=1) if 'RRB' in self.model.__class__.__name__: prediction_nn, prediction_kd, prob, prediction_rrb, prediction_speed = self.model( obs=xy_copy[:, :n_obs, :, :], scene=resampled_scene, sample_rate=torch.tensor([sample_rate], device=device), pixel_scale=pixel_scale, center_line_dict=center_line_rotated, rotated_scene=rotated_scene[0], file_name=file_name, margin=scene_funcs.return_margin([file_name]), prediction_truth=xy_copy[:, n_obs:, :, :]) elif 'EDN' in self.model.__class__.__name__: prediction_v, prediction_nn, prob, prediction_speed = self.model(obs=xy_copy[:, :n_obs, :, :], scene=resampled_scene, sample_rate=torch.tensor([sample_rate], device=device), pixel_scale=pixel_scale, center_line_dict=center_line_rotated, rotated_scene=rotated_scene[0], file_name=file_name, margin=scene_funcs.return_margin( [file_name]), prediction_truth=xy_copy[:, n_obs:, :, :]) prediction_kd = prediction_nn prediction_rrb = prediction_nn prob, prediction_nn, prediction_kd, prediction_rrb = prob[0], prediction_nn[0, :, :, :2].to( device), prediction_kd[0, :, :, :2].to(device), prediction_rrb[0, :, :, :2].to( device) # as we have one batch prob = torch.nn.functional.softmax(prob, dim=0).cpu().numpy() best_mode_real_nn = l2_dist(xy_copy[0, n_obs:], prediction_nn) best_mode_real_rrb = l2_dist(xy_copy[0, n_obs:], prediction_rrb) best_mode_real_kd = l2_dist(xy_copy[0, n_obs:], prediction_kd) if (use_mpc): for i in range(prediction_rrb.shape[0]): prediction_rrb[i] = mpc_fun(prediction_rrb[i], sample_rate, pixel_scale, n_pred, xy_copy[:, :n_obs, :, :])[1:] # since the output has 11 best_mode_prediction_nn = prediction_nn[best_mode_real_nn, :, :].reshape(-1, 2) # torch.Size([n_pred, 2]) best_mode_prediction_rrb = prediction_rrb[best_mode_real_rrb, :, :].reshape(-1, 2) # torch.Size([n_pred, 2]) best_mode_prediction_kd = prediction_kd[best_mode_real_kd, :, :].reshape(-1, 2) # torch.Size([n_pred, 2]) best_mode_prediction_rrb_unrotated = best_mode_prediction_rrb.clone() best_mode_prediction_kd_unrotated = best_mode_prediction_kd.clone() best_mode_prediction_nn_unrotated = best_mode_prediction_nn.clone() prediction_nn_rotated = prediction_nn.clone() prediction_kd_rotated = prediction_kd.clone() prediction_rrb_rotated = prediction_rrb.clone() rotation_enabled = 1 # if 1, it draws when we normalized the scene (for scene models) and if 0, draws without rotation. best_mode_prediction_nn = augmentation.rotate_path_by_theta(best_mode_prediction_nn, center=xy_copy[:, self.n_obs - 1, 0, :], n_pred=n_pred, theta=-theta) # rotate back to original scene best_mode_prediction_kd = augmentation.rotate_path_by_theta(best_mode_prediction_kd, center=xy_copy[:, self.n_obs - 1, 0, :], n_pred=n_pred, theta=-theta) # rotate back to original scene best_mode_prediction_rrb = augmentation.rotate_path_by_theta(best_mode_prediction_rrb, center=xy_copy[:, self.n_obs - 1, 0, :], n_pred=n_pred, theta=-theta) # rotate back to original scene for i in range(prediction_nn.size(0)): # num modes prediction_nn_rotated[i] = augmentation.rotate_path_by_theta(prediction_nn[i], center=xy_copy[:, self.n_obs - 1, 0, :], n_pred=n_pred, theta=-theta) # rotate back to original scene prediction_kd_rotated[i] = augmentation.rotate_path_by_theta(prediction_kd[i], center=xy_copy[:, self.n_obs - 1, 0, :], n_pred=n_pred, theta=-theta) # rotate back to original scene prediction_rrb_rotated[i] = augmentation.rotate_path_by_theta(prediction_rrb[i], center=xy_copy[:, self.n_obs - 1, 0, :], n_pred=n_pred, theta=-theta) # rotate back to original scene test_time = start_time - time.time() scene_violation = offroad_detector(prediction_rrb_rotated, file_name, scene_funcs.image, prob) projected_back_traj_rrb = best_mode_prediction_rrb.clone() projected_back_traj_rrb[:, 1] = best_mode_prediction_rrb[:, 1] / scale + offset[0] projected_back_traj_rrb[:, 0] = -best_mode_prediction_rrb[:, 0] / scale + offset[1] projected_back_traj_nn = best_mode_prediction_nn.clone() projected_back_traj_nn[:, 1] = best_mode_prediction_nn[:, 1] / scale + offset[0] projected_back_traj_nn[:, 0] = -best_mode_prediction_nn[:, 0] / scale + offset[1] projected_back_traj_kd = best_mode_prediction_kd.clone() projected_back_traj_kd[:, 1] = best_mode_prediction_kd[:, 1] / scale + offset[0] projected_back_traj_kd[:, 0] = -best_mode_prediction_kd[:, 0] / scale + offset[1] flag = 0 if ('Merging' in file_name): pass else: scene_violation_truth = offroad_detector(xy_copy_unrotated[:, 5:, 0], file_name, scene_funcs.image) if (scene_violation_truth > 1): # and (error_kd_old-error_kd)>1): flag = 1 if visualize: draw_scene(scene_funcs, xy_copy, file_name, rotated_scene, n_obs, n_pred, prob, best_mode_prediction_rrb_unrotated.unsqueeze(0), rotation_enabled, ped_id, first_frame, 'res', best_mode_prediction_nn_unrotated.unsqueeze(0), best_mode_prediction_kd_unrotated.unsqueeze(0)) return [trajnettools.TrackRow(first_frame + i * frame_diff, ped_id, x, y) for i, (x, y) in enumerate(projected_back_traj_rrb)], test_time, scene_violation, flag