def generate_world_coordinates_bikers(path_in): for folder in os.listdir(path_in): for sub_folder in os.listdir(os.path.join(path_in, folder)): file_name = os.path.join(path_in, folder, sub_folder, 'annotations.txt') scene_name = folder +'_' + sub_folder[-1] if scene_name in ('bookstore_6', 'bookstore_5', 'bookstore_4', 'coupa_2', 'nexus_0', 'hyang_2', 'hyang_3', 'hyang_4', 'hyang_9', 'quad_2', 'quad_3'): continue data = pd.read_csv(file_name, sep=" ", header=None) data.columns = ["agentID", "xmin", "ymin", "xmax", "ymax", "frameID", "lost", "occluded", "generated", "label"] # Transform h_matrix = pd.read_csv(get_root_dir() + '/data/SDD/' + scene_name + '/' + scene_name + "_homography.txt", delim_whitespace=True, header=None).values pixels = np.transpose(np.vstack((data.xmax.values + data.xmin.values, data.ymax.values + data.ymin.values))) / 2 coordinates = get_world_from_pixels(pixels, h_matrix) # Filter idx = data.label == "Biker" data = data[idx] coordinates = coordinates[idx] idx_down = data.frameID % SAMPLE_RATE == 0 data = data[idx_down] coordinates = coordinates[idx_down] data_sdd_original = np.transpose(np.vstack((data.frameID.values, data.agentID.values, coordinates[:, 0], coordinates[:, 1]))) np.savetxt("/home/q392358/Documents/FLORA/data/SDD_ALL/sdd_all/Training/train/{}.txt".format(scene_name), data_sdd_original, delimiter=' ', fmt='%.3f')
def save_sample(fig, ax, directory, model_name, scene, model_variant, frame, sample): extent = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) fig.savefig(get_root_dir() + directory + '{}/scene_{}_variant_{}_frame_{}_sample_{}.png'.format( model_name, scene, model_variant, frame, sample), bbox_inches=extent)
def load_pickle(name, scene, num, data_set, model_name): path_name = '{}/{}/{}/scene_{}_batch_{}_{}.pkl'.format( get_root_dir() + '/results/trajectories/', data_set, model_name, scene, str(num), name) with open(path_name, 'rb') as handle: list = pickle.load(handle) return list
def get_homography_and_map(dset, annotated_points_name = '/world_points_boundary.npy'): directory = get_root_dir() + '/data/' path_group = os.path.join(directory, get_dset_group_name(dset)) path = os.path.join(path_group, dset) h_matrix = pd.read_csv(path + '/{}_homography.txt'.format(dset), delim_whitespace=True, header=None).values if 'txt' in annotated_points_name: map = np.loadtxt(path + annotated_points_name, delimiter=' ') elif 'jpg' in annotated_points_name: map = load_bin_map(path + annotated_points_name) else: map = np.load(path + annotated_points_name) return map, h_matrix
def generate_world_coordinates(training_path, path_in, path_out): files = sorted(os.listdir(get_root_dir() + training_path)) for file in files: dataset_name = file[:-10] if dataset_name in ('hyang_3', 'hyang_4', 'hyang_9', 'nexus_0'): return data_trajnet_pixel = pd.read_csv(path_in, sep=" ", header=None) data_trajnet = np.transpose(np.vstack((data_trajnet_pixel.loc[:, 1], data_trajnet_pixel.loc[:, 2], data_trajnet_pixel.loc[:, 3], data_trajnet_pixel.loc[:, 4]))) h = get_homography(dataset_name) data_trajnet_world_from_original = get_world_from_pixels(data_trajnet[:, 2:4], h) data_trajnet_original = np.transpose(np.vstack((data_trajnet_pixel.loc[:, 1], data_trajnet_pixel.loc[:, 2], data_trajnet_world_from_original[:, 0], data_trajnet_world_from_original[:, 1]))) np.savetxt(path_out, data_trajnet_original, delimiter=' ', fmt='%.3f')
def set_dset_list(self, data_dir, down_sampling=True, down_samples=200): directory = get_root_dir() + '/datasets/safegan_dataset/' self.list_data_files = sorted([get_dset_name(os.path.join(data_dir, _path).split("/")[-1]) for _path in os.listdir(data_dir)]) for name in self.list_data_files: path_group = os.path.join(directory, get_dset_group_name(name)) """ The inputs are the boundary points between the traversable and non-traversable areas. It is possible to take all points or just a sample""" path = os.path.join(path_group, name) map = np.load(path + "/world_points_boundary.npy") if down_samples != -1 and down_sampling and map.shape[0] > down_samples: down_sampling = (map.shape[0] // down_samples) sampled = map[::down_sampling] map = sampled[:down_samples] self.scene_information[name] = torch.from_numpy(map).type(torch.float).to(device)
def main(): model_path = os.path.join(get_root_dir(), 'models_sdd/temp/checkpoint_with_model.pt') if True: # load checkpoint of first model and arguments checkpoint1 = torch.load(model_path) args1 = AttrDict(checkpoint1['args']) generator1 = get_generator(checkpoint1, args1) encoder_out = torch.randn(1, 64, 32).cuda() curr_hidden = torch.randn(1, 32).cuda() embed_info = torch.randn(1, 4).cuda() _, attention_weights = generator1.pooling.pooling_list[ 1].attention_decoder.forward( encoder_out, curr_hidden, embed_info) #torch.zeros(1, 256*256).cuda() visualize_attention_weights('gates_8', 8, attention_weights, torch.tensor([25, 55]).unsqueeze(0))
def main(args): data_set = 'ALL' model_path = os.path.join( get_root_dir(), 'results/models/{}/{}'.format(data_set, args.model_folder)) if os.path.isdir(model_path): filenames = sorted(os.listdir(model_path)) paths = [os.path.join(model_path, file_) for file_ in filenames] paths = [path for path in paths if 'no_model' not in path] # load checkpoint of first model and arguments checkpoint1 = torch.load(paths[0]) print('Loading model from path: ' + paths[0]) # load checkpoing of second model checkpoint2 = torch.load(paths[1]) print('Loading model from path: ' + paths[1]) evaluate_training_metric(checkpoint1, checkpoint2, args.metric, 'val') print('Check folder name {}'.format(os.path.join(model_path)))
def set_dset_list(self, data_dir): """ Fill scene_information with the static environment features that will be used as part of the input of Static Scene Feature Extractor module in SafeGAN""" directory = get_root_dir() + '/datasets/safegan_dataset/' self.list_data_files = sorted([ get_dset_name(os.path.join(data_dir, _path).split("/")[-1]) for _path in os.listdir(data_dir) ]) for name in self.list_data_files: path_group = os.path.join(directory, get_dset_group_name(name)) if self.pool_static_type == "physical_attention_no_encoder": """ In this case the features are the one extracted by one of Segmentation Networks I trained on the new dataset I created. The features are taken before the last upsample layers.""" path = os.path.join(path_group + "/segmented_features", name) features = np.load(path + "_segmentation_features.npy") features = torch.from_numpy(features).type( torch.float).to(device) elif self.pool_static_type == "physical_attention_with_encoder": """ In this case the input is the raw image or the segmented one (by one of the Segmentation Networks I trained on the new dataset I created). This image is then encoded by a Deep Network like ResNet""" path = os.path.join(path_group + "/segmented_scenes", name) image = plt.imread(path + ".jpg") image = torch.from_numpy(image).type(torch.float).to(device) # Images fed to the model must be a Float tensor of dimension N, 3, 256, 256, where N is the batch size. # PyTorch follows the NCHW convention, which means the channels dimension (C) must precede the size dimensions image = image.permute(2, 0, 1) # Normalize the image image = self.transform(image) features = self.attention_encoder(image.unsqueeze(0)) else: print( "ERROR in recognizing physical attention pool static type") exit() self.scene_information[name] = features
def set_dset_list(self, data_dir, down_sampling=True): """ Fill scene_information with the static environment features that will be used as part of the input of Static Scene Feature Extractor module in SafeGAN""" directory = get_root_dir() + '/data/' self.list_data_files = sorted([ get_dset_name(os.path.join(data_dir, _path).split("/")[-1]) for _path in os.listdir(data_dir) ]) for name in self.list_data_files: path_group = os.path.join(directory, get_dset_group_name(name)) """ The inputs are the boundary points between the traversable and non-traversable areas. It is possible to take all points or just a sample""" path = os.path.join(path_group, name) map = np.load(path + "/world_points_boundary.npy") if self.down_samples != -1 and down_sampling and map.shape[ 0] > self.down_samples: down_sampling = (map.shape[0] // self.down_samples) sampled = map[::down_sampling] map = sampled[:self.down_samples] self.scene_information[name] = torch.from_numpy(map).type( torch.float).to(device)
def save_pickle(list, name, scene, num, data_set, model_name): path_name = '{}/{}/{}/scene_{}_batch_{}_{}.pkl'.format( get_root_dir() + '/results/trajectories/', data_set, model_name, scene, str(num), name) with open(path_name, 'wb') as fp: pickle.dump(list, fp)
import imageio import numpy as np import os import datetime import pandas as pd from scripts.data_processing.generate_world_coordinates import generate_world_coordinates from scripts.data_processing.generate_world_points_boundary import generate_boundary_points from sgan.model.folder_utils import get_root_dir, get_sdd_dir from sgan.model.models_static_scene import get_homography, get_pixels_from_world, get_world_from_pixels from sgan.model.models_static_scene import load_bin_map files = sorted( os.listdir(get_root_dir() + '/data/TRAJNETPIXEL/trajnetpixel/Training/train')) frame_rate_sdd = 30 frame_rate_tn = 12 def convert_image_to_world(path, dataset_name): occupancy_map = load_bin_map(path) coordinates = [] for row in range(occupancy_map.shape[0]): for col in range(occupancy_map.shape[1]): if occupancy_map[row][col] == 0: coordinates.append(np.array([col, row])) coordinates = np.asarray(coordinates) h = get_homography(dataset_name)
def get_path(dset): directory = get_root_dir() + '/data/' path_group = os.path.join(directory, get_dset_group_name(dset)) path = os.path.join(path_group, dset) return path
def main(args): data_set = 'TRAJNET' model_path1 = os.path.join( get_root_dir(), 'results/models/{}/{}/{}'.format(data_set, args.model_folder, args.model_name1)) model_path2 = os.path.join( get_root_dir(), 'results/models/{}/{}/{}'.format(data_set, args.model_folder, args.model_name2)) data_dir = get_test_data_path(data_set.lower()) # load checkpoint of first model and arguments checkpoint1 = torch.load(model_path1) args1 = AttrDict(checkpoint1['args']) print('Loading model from path: ' + model_path1) generator1 = get_generator(checkpoint1, args1) # load checkpoing of second model checkpoint2 = torch.load(model_path2) args2 = AttrDict(checkpoint2['args']) print('Loading model from path: ' + model_path2) generator2 = get_generator(checkpoint2, args2) if args.precompute_required: collect_generated_samples(args1, generator1, generator2, data_dir, data_set, args.model_folder, selected_scene=args.scene, selected_batch=-1) m1, m2, counter = 0, 0, 0 for batch in range(0, 16): print('testing data_set = {}, scene = {}, batch = {}'.format( data_set, args.scene, batch)) if args.metric == 'ade': out1, out2 = evaluate_test_ade(data_set, args.model_folder, selected_scene=args.scene, selected_batch=batch) elif args.metric == 'fde': out1, out2 = evaluate_test_fde(data_set, args.model_folder, selected_scene=args.scene, selected_batch=batch) elif args.metric == 'ade_pixel': out1, out2 = evaluate_test_pixel_ade(data_set, args.model_folder, selected_scene=args.scene, selected_batch=batch) elif args.metric == 'fde_pixel': out1, out2 = evaluate_test_pixel_fde(data_set, args.model_folder, selected_scene=args.scene, selected_batch=batch) elif args.metric == 'cols': out1, out2 = evaluate_test_cols(data_set, args.model_folder, selected_scene=args.scene, selected_batch=batch) elif args.metric == 'quality': evaluate_trajectory_quality(data_set, args.model_folder, selected_scene=args.scene, selected_batch=-1) if out1 > 0. and out2 > 0.: print('{} model 1: {:.6f} model 2: {:.6f}'.format( args.metric, out1, out2)) m1 += out1 m2 += out2 counter += 1 print('Avarage {} model 1: {:.6f} model 2: {:.6f}'.format( args.metric, m1 / counter, m2 / counter))
def get_homography(dset): directory = get_root_dir() + '/data/' path_group = os.path.join(directory, get_dset_group_name(dset)) path = os.path.join(path_group, dset) h_matrix = pd.read_csv(path + '/{}_homography.txt'.format(dset), delim_whitespace=True, header=None).values return h_matrix
def main(args): if args.summary_writer_name is not None: writer = SummaryWriter(args.summary_writer_name) os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_num train_path = get_dset_path(args.dataset_path, args.dataset_name, 'train') val_path = get_dset_path(args.dataset_path, args.dataset_name, 'val') long_dtype, float_dtype = get_dtypes(args) logger.info("Initializing val dataset") val_dset, val_loader = data_loader(args, val_path, shuffle=False) logger.info("Initializing train dataset") train_dset, train_loader = data_loader(args, train_path, shuffle=True) print(len(train_loader)) steps = max(args.g_steps, args.c_steps) steps = max(steps, args.d_steps) iterations_per_epoch = math.ceil(len(train_dset) / args.batch_size / steps) if args.num_epochs: args.num_iterations = int(iterations_per_epoch * args.num_epochs) logger.info('There are {} iterations per epoch, prints {} plots {}'.format( iterations_per_epoch, args.print_every, args.checkpoint_every)) generator = helper_get_generator(args, train_path) generator.apply(init_weights) generator.type(float_dtype).train() logger.info('Here is the generator:') logger.info(generator) g_loss_fn = gan_g_loss optimizer_g = optim.Adam(filter(lambda x: x.requires_grad, generator.parameters()), lr=args.g_learning_rate) # build trajectory discriminator = TrajectoryDiscriminator( obs_len=args.obs_len, pred_len=args.pred_len, embedding_dim=args.embedding_dim, h_dim=args.encoder_h_dim_d, mlp_dim=args.mlp_dim, num_layers=args.num_layers, dropout=args.dropout, activation=args.activation, batch_norm=args.batch_norm, grid_size=args.grid_size, neighborhood_size=args.neighborhood_size) discriminator.apply(init_weights) discriminator.type(float_dtype).train() logger.info('Here is the discriminator:') logger.info(discriminator) d_loss_fn = gan_d_loss optimizer_d = optim.Adam(discriminator.parameters(), lr=args.d_learning_rate) critic = helper_get_critic(args, train_path) critic.apply(init_weights) critic.type(float_dtype).train() logger.info('Here is the critic:') logger.info(critic) c_loss_fn = gan_d_loss optimizer_c = optim.Adam(filter(lambda x: x.requires_grad, critic.parameters()), lr=args.c_learning_rate) trajectory_evaluator = TrajectoryGeneratorEvaluator() if args.d_loss_weight > 0: logger.info('Discrimintor loss') trajectory_evaluator.add_module(discriminator, gan_g_loss, args.d_loss_weight) if args.c_loss_weight > 0: logger.info('Critic loss') trajectory_evaluator.add_module(critic, g_critic_loss_function, args.c_loss_weight) # Maybe restore from checkpoint restore_path = None if args.checkpoint_start_from is not None: restore_path = os.path.join(get_root_dir(), args.output_dir, args.checkpoint_start_from) elif args.restore_from_checkpoint == 1: restore_path = os.path.join(get_root_dir(), args.output_dir, '%s_with_model.pt' % args.checkpoint_name) if restore_path is not None and os.path.isfile(restore_path): logger.info('Restoring from checkpoint {}'.format(restore_path)) checkpoint = torch.load(restore_path) generator.load_state_dict(checkpoint['g_state']) # discriminator.load_state_dict(checkpoint['d_state']) optimizer_g.load_state_dict(checkpoint['g_optim_state']) # optimizer_d.load_state_dict(checkpoint['d_optim_state']) t = checkpoint['counters']['t'] epoch = checkpoint['counters']['epoch'] checkpoint['restore_ts'].append(t) else: # Starting from scratch, so initialize checkpoint data structure t, epoch = 0, -1 checkpoint = { 'args': args.__dict__, 'G_losses': defaultdict(list), 'D_losses': defaultdict(list), 'C_losses': defaultdict(list), 'losses_ts': [], 'metrics_val': defaultdict(list), 'metrics_train': defaultdict(list), 'sample_ts': [], 'restore_ts': [], 'norm_g': [], 'norm_d': [], 'norm_c': [], 'counters': { 't': None, 'epoch': None, }, 'g_state': None, 'g_optim_state': None, 'd_state': None, 'd_optim_state': None, 'c_state': None, 'c_optim_state': None, 'g_best_state': None, 'd_best_state': None, 'c_best_state': None, 'best_t': None, 'g_best_nl_state': None, 'd_best_state_nl': None, 'best_t_nl': None, } t0 = None # Number of times a generator, discriminator and critic steps are done in 1 epoch num_d_steps = ((len(train_dset) / args.batch_size) / (args.g_steps + args.d_steps + args.c_steps)) * args.d_steps num_c_steps = ((len(train_dset) / args.batch_size) / (args.g_steps + args.d_steps + args.c_steps)) * args.c_steps num_g_steps = ((len(train_dset) / args.batch_size) / (args.g_steps + args.d_steps + args.c_steps)) * args.g_steps while t < args.num_iterations: if epoch == args.num_epochs: break gc.collect() d_steps_left = args.d_steps g_steps_left = args.g_steps c_steps_left = args.c_steps epoch += 1 # Average losses over all batches in the training set for 1 epoch avg_losses_d = {} avg_losses_c = {} avg_losses_g = {} logger.info('Starting epoch {} - [{}]'.format( epoch, time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))) for batch_num, batch in enumerate(train_loader): if args.timing == 1: torch.cuda.synchronize() t1 = time.time() # Decide whether to use the batch for stepping on discriminator or # generator; an iteration consists of args.d_steps steps on the # discriminator followed by args.g_steps steps on the generator. if d_steps_left > 0: step_type = 'd' losses_d = discriminator_step(args, batch, generator, discriminator, d_loss_fn, optimizer_d) checkpoint['norm_d'].append( get_total_norm(discriminator.parameters())) d_steps_left -= 1 if len(avg_losses_d) == 0: for k, v in sorted(losses_d.items()): avg_losses_d[k] = v / num_d_steps else: for k, v in sorted(losses_d.items()): avg_losses_d[k] += v / num_d_steps elif c_steps_left > 0: step_type = 'c' losses_c = critic_step(args, batch, generator, critic, c_loss_fn, optimizer_c) checkpoint['norm_c'].append(get_total_norm( critic.parameters())) c_steps_left -= 1 if len(avg_losses_c) == 0: for k, v in sorted(losses_c.items()): avg_losses_c[k] = v / num_c_steps else: for k, v in sorted(losses_c.items()): avg_losses_c[k] += v / num_c_steps elif g_steps_left > 0: step_type = 'g' losses_g = generator_step(args, batch, generator, optimizer_g, trajectory_evaluator) checkpoint['norm_g'].append( get_total_norm(generator.parameters())) g_steps_left -= 1 if len(avg_losses_g) == 0: for k, v in sorted(losses_g.items()): avg_losses_g[k] = v / num_g_steps else: for k, v in sorted(losses_g.items()): avg_losses_g[k] += v / num_g_steps if args.timing == 1: torch.cuda.synchronize() t2 = time.time() logger.info('{} step took {}'.format(step_type, t2 - t1)) # Skip the rest if we are not at the end of an iteration if d_steps_left > 0 or g_steps_left > 0 or c_steps_left > 0: continue if args.timing == 1: if t0 is not None: logger.info('Iteration {} took {}'.format( t - 1, time.time() - t0)) t0 = time.time() t += 1 d_steps_left = args.d_steps g_steps_left = args.g_steps c_steps_left = args.c_steps if epoch % args.print_every == 0 and epoch > 0: # Save losses logger.info('t = {} / {}'.format(t + 1, args.num_iterations)) if args.d_steps > 0: for k, v in sorted(avg_losses_d.items()): logger.info(' [D] {}: {:.3f}'.format(k, v)) checkpoint['D_losses'][k].append(v) if args.summary_writer_name is not None: writer.add_scalar('Train/' + k, v, epoch) for k, v in sorted(avg_losses_g.items()): logger.info(' [G] {}: {:.3f}'.format(k, v)) checkpoint['G_losses'][k].append(v) if args.summary_writer_name is not None: writer.add_scalar('Train/' + k, v, epoch) if args.c_steps > 0: for k, v in sorted(avg_losses_c.items()): logger.info(' [C] {}: {:.3f}'.format(k, v)) checkpoint['C_losses'][k].append(v) if args.summary_writer_name is not None: writer.add_scalar('Train/' + k, v, epoch) checkpoint['losses_ts'].append(t) if epoch % args.checkpoint_every == 0 and epoch > 0: # Maybe save a checkpoint if t > 0: checkpoint['counters']['t'] = t checkpoint['counters']['epoch'] = epoch checkpoint['sample_ts'].append(t) metrics_train, metrics_val = {}, {} if args.g_steps > 0: logger.info('Checking G stats on train ...') metrics_train = check_accuracy_generator( 'train', epoch, args, train_loader, generator, True) logger.info('Checking G stats on val ...') metrics_val = check_accuracy_generator( 'val', epoch, args, val_loader, generator, True) if args.c_steps > 0: logger.info('Checking C stats on train ...') metrics_train_c = check_accuracy_critic( args, train_loader, generator, critic, c_loss_fn, True) metrics_train.update(metrics_train_c) logger.info('Checking C stats on val ...') metrics_val_c = check_accuracy_critic( args, val_loader, generator, critic, c_loss_fn, True) metrics_val.update(metrics_val_c) if args.d_steps > 0: logger.info('Checking D stats on train ...') metrics_train_d = check_accuracy_discriminator( args, train_loader, generator, discriminator, d_loss_fn, True) metrics_train.update(metrics_train_d) logger.info('Checking D stats on val ...') metrics_val_d = check_accuracy_discriminator( args, val_loader, generator, discriminator, d_loss_fn, True) metrics_val.update(metrics_val_d) for k, v in sorted(metrics_val.items()): logger.info(' [val] {}: {:.3f}'.format(k, v)) checkpoint['metrics_val'][k].append(v) if args.summary_writer_name is not None: writer.add_scalar('Validation/' + k, v, epoch) for k, v in sorted(metrics_train.items()): logger.info(' [train] {}: {:.3f}'.format(k, v)) checkpoint['metrics_train'][k].append(v) if args.summary_writer_name is not None: writer.add_scalar('Train/' + k, v, epoch) # Save another checkpoint with model weights and # optimizer state checkpoint['g_state'] = generator.state_dict() checkpoint['g_optim_state'] = optimizer_g.state_dict() checkpoint['d_state'] = discriminator.state_dict() checkpoint['d_optim_state'] = optimizer_d.state_dict() checkpoint['c_state'] = critic.state_dict() checkpoint['c_optim_state'] = optimizer_c.state_dict() checkpoint_path = os.path.join( get_root_dir(), args.output_dir, '{}_{}_with_model.pt'.format(args.checkpoint_name, epoch)) logger.info('Saving checkpoint to {}'.format(checkpoint_path)) torch.save(checkpoint, checkpoint_path) logger.info('Done.') # Save a checkpoint with no model weights by making a shallow # copy of the checkpoint excluding some items checkpoint_path = os.path.join( get_root_dir(), args.output_dir, '{}_{}_no_model.pt'.format(args.checkpoint_name, epoch)) logger.info('Saving checkpoint to {}'.format(checkpoint_path)) key_blacklist = [ 'g_state', 'd_state', 'c_state', 'g_best_state', 'g_best_nl_state', 'g_optim_state', 'd_optim_state', 'd_best_state', 'd_best_nl_state', 'c_optim_state', 'c_best_state' ] small_checkpoint = {} for k, v in checkpoint.items(): if k not in key_blacklist: small_checkpoint[k] = v torch.save(small_checkpoint, checkpoint_path) logger.info('Done.') if args.g_steps < 1: continue min_ade = min(checkpoint['metrics_val']['ade']) min_ade_nl = min(checkpoint['metrics_val']['ade_nl']) if metrics_val['ade'] == min_ade: logger.info('New low for avg_disp_error') checkpoint['best_t'] = t checkpoint['g_best_state'] = generator.state_dict() checkpoint['d_best_state'] = discriminator.state_dict() if args.c_steps > 0: checkpoint['c_best_state'] = critic.state_dict() if metrics_val['ade_nl'] == min_ade_nl: logger.info('New low for avg_disp_error_nl') checkpoint['best_t_nl'] = t checkpoint['g_best_nl_state'] = generator.state_dict() checkpoint['d_best_nl_state'] = discriminator.state_dict() if args.summary_writer_name is not None: writer.close()
def visualize_attention_weights(scene_name, encoded_image_size, attention_weights, curr_end_pos, ax1, ax2, counter=0): """ Function to visualize the attention weights on their relative scene image during inference time (training or testing). :param scene_name: the name of the SDD scene from which the attention weights were computed :param encoded_image_size: the width/height dimension of the scene image used as input for the Attention Encoder. The image should be a squared image, thus the dimension should be (encoded_image_size, encoded_image_size) :param attention_weights: the weights computed by the attention module :param curr_end_pos: the current positions of all agents in a scene """ ped_id = 0 grid_size = 8.0 # 'upscaling_factor' is used to increment the size of the scene image (to make it better visualizable) as well as # the agents' positions coordinates do adapt them to the new image size upscaling_factor = 1 ax1.cla() ax2.cla() # 'dataset_path' represents the path with the SDD scene folders inside dataset_path = get_root_dir() + "/datasets/safegan_dataset/SDD/" # Load the raw scene image on which the attention weights will be plotted. # Here I suppose they are inside a folder called "segmented_scenes" image_original = Image.open(get_root_dir() + "/datasets/safegan_dataset/SDD/" + scene_name + "/reference.jpg") original_image_size = Image.open(dataset_path + scene_name + "/annotated_boundaries.jpg").size # Increase the dimension of the raw scene image image_original = image_original.resize([ original_image_size[0] * upscaling_factor, original_image_size[1] * upscaling_factor ], Image.LANCZOS) # In order to plot the agents's coordinates on the scene image it is necessary to load the homography matrix of that scene # and then to convert the world coordinates into pixel values h_matrix = pd.read_csv(dataset_path + scene_name + '/{}_homography.txt'.format(scene_name), delim_whitespace=True, header=None).values pixels = get_pixels_from_world(curr_end_pos, h_matrix, True) #pixels = pixels * (encoded_image_size * upscaling_factor / original_image_size[0], # encoded_image_size * upscaling_factor / original_image_size[1]) # Here it is necessary to resize also the pixel coordinates of the agents' positions, according to the upscaling # factor and the original dimension of the scene image (that I take from the image with the annotated boundary points) original_image_size = Image.open(dataset_path + scene_name + "/annotated_boundaries.jpg").size # w, h = image_original.size col = np.round(pixels[ped_id][0]) row = np.round(pixels[ped_id][1]) grid_left_upper_corner = curr_end_pos - torch.tensor( [grid_size / 2.0, grid_size / 2.0]).expand_as(curr_end_pos).to(device) pixels_grid = get_pixels_from_world(grid_left_upper_corner, h_matrix, True) col_grid = (col - np.round(pixels_grid[ped_id][0])) row_grid = (row - np.round(pixels_grid[ped_id][1])) if row - row_grid > 0 and row + row_grid < h and col - col_grid > 0 and col + col_grid < w: image = image_original.crop( (col - col_grid, row - row_grid, col + col_grid, row + row_grid)) else: image = image_original.crop((w // 2 - col_grid, h // 2 - row_grid, w // 2 + col_grid, h // 2 + row_grid)) #image = image.resize([20 * upscaling_factor, 20 * upscaling_factor], Image.LANCZOS) # Resize the attention weights dimension to match it with the dimension of the upscaled raw scene image. # To expand the attention weights I use skimage that allows us to also smooth the pixel values during the expansion attention_weights = attention_weights.view( -1, encoded_image_size, encoded_image_size).detach().cpu().numpy() upscaling_factor = image.size[0] / encoded_image_size alpha = skimage.transform.pyramid_expand(attention_weights[ped_id], upscale=upscaling_factor, sigma=8) #pixels = np.expand_dims( np.array([w//2, h//2]), axis=0) rect = patches.Rectangle((pixels_grid[ped_id, 0], pixels_grid[ped_id, 1]), 2 * col_grid, 2 * row_grid, linewidth=1, edgecolor='white', facecolor='none') # Plot raw scene image, the agents' positions and the attention weights ax1.imshow(image_original) ax1.scatter(pixels[:, 0], pixels[:, 1], marker='.', color="b") ax1.scatter(pixels[ped_id, 0], pixels[ped_id, 1], marker='X', color="r") ax1.add_patch(rect) ax1.axis('off') ax2.imshow(image) ax2.imshow(np.flipud(alpha), alpha=0.7) plt.set_cmap(cm.Greys_r) ax2.axis('off') directory = get_root_dir() + '/results/plots/SDD/safeGAN_DP/attention' files = len(os.listdir(directory)) plt.savefig(directory + '/frame_{}.png'.format(files + 1)) plt.draw() plt.waitforbuttonpress()