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()
