def write_video_summary(vid_dataset, model, model_input, gt, model_output, writer, total_steps, prefix='train_'):
resolution = vid_dataset.shape
frames = [0, 60, 120, 200]
Nslice = 10
with torch.no_grad():
coords = [dataio.get_mgrid((1, resolution[1], resolution[2]), dim=3)[None,...].cpu() for f in frames]
for idx, f in enumerate(frames):
coords[idx][..., 0] = (f / (resolution[0] - 1) - 0.5) * 2
coords = torch.cat(coords, dim=0)
output = torch.zeros(coords.shape)
split = int(coords.shape[1] / Nslice)
for i in range(Nslice):
pred = model({'coords':coords[:, i*split:(i+1)*split, :]})['model_out']
output[:, i*split:(i+1)*split, :] = pred.cpu()
pred_vid = output.view(len(frames), resolution[1], resolution[2], 3) / 2 + 0.5
pred_vid = torch.clamp(pred_vid, 0, 1)
gt_vid = torch.from_numpy(vid_dataset.vid[frames, :, :, :])
psnr = 10*torch.log10(1 / torch.mean((gt_vid - pred_vid)**2))
pred_vid = pred_vid.permute(0, 3, 1, 2)
gt_vid = gt_vid.permute(0, 3, 1, 2)
output_vs_gt = torch.cat((gt_vid, pred_vid), dim=-2)
writer.add_image(prefix + 'output_vs_gt', make_grid(output_vs_gt, scale_each=False, normalize=True),
global_step=total_steps)
min_max_summary(prefix + 'coords', model_input['coords'], writer, total_steps)
min_max_summary(prefix + 'pred_vid', pred_vid, writer, total_steps)
writer.add_scalar(prefix + "psnr", psnr, total_steps)
def val_fn(model, ckpt_dir, epoch):
# Time values at which the function needs to be plotted
times = [0., 0.5 * (opt.tMax - 0.1), (opt.tMax - 0.1)]
num_times = len(times)
# Theta slices to be plotted
thetas = [-math.pi, -0.5 * math.pi, 0., 0.5 * math.pi, math.pi]
num_thetas = len(thetas)
# Create a figure
fig = plt.figure(figsize=(5 * num_times, 5 * num_thetas))
# Get the meshgrid in the (x, y) coordinate
sidelen = 200
mgrid_coords = dataio.get_mgrid(sidelen)
# Start plotting the results
for i in range(num_times):
time_coords = torch.ones(mgrid_coords.shape[0], 1) * times[i]
for j in range(num_thetas):
theta_coords = torch.ones(mgrid_coords.shape[0], 1) * thetas[j]
theta_coords = theta_coords / (opt.angle_alpha * math.pi)
coords = torch.cat((time_coords, mgrid_coords, theta_coords),
dim=1)
if torch.cuda.is_available():
model_in = {'coords': coords.cuda()}
else:
model_in = {'coords': coords.cpu()}
model_out = model(model_in)['model_out']
# Detatch model ouput and reshape
model_out = model_out.detach().cpu().numpy()
model_out = model_out.reshape((sidelen, sidelen))
# Unnormalize the value function
norm_to = 0.02
mean = 0.25
var = 0.5
model_out = (model_out * var / norm_to) + mean
# Plot the zero level sets
model_out = (model_out <= 0.001) * 1.
# Plot the actual data
ax = fig.add_subplot(num_times, num_thetas,
(j + 1) + i * num_thetas)
ax.set_title('t = %0.2f, theta = %0.2f' % (times[i], thetas[j]))
s = ax.imshow(model_out.T,
cmap='bwr',
origin='lower',
extent=(-1., 1., -1., 1.))
fig.colorbar(s)
fig.savefig(
os.path.join(ckpt_dir, 'BRS_validation_plot_epoch_%04d.png' % epoch))
def write_sdf_summary(model,
model_input,
gt,
model_output,
writer,
total_steps,
prefix='train_'):
slice_coords_2d = dataio.get_mgrid(512)
with torch.no_grad():
yz_slice_coords = torch.cat(
(torch.zeros_like(slice_coords_2d[:, :1]), slice_coords_2d),
dim=-1)
yz_slice_model_input = {'coords': yz_slice_coords.cuda()[None, ...]}
yz_model_out = model(yz_slice_model_input)
sdf_values = yz_model_out['model_out']
sdf_values = dataio.lin2img(sdf_values).squeeze().cpu().numpy()
fig = make_contour_plot(sdf_values)
writer.add_figure(prefix + 'yz_sdf_slice',
fig,
global_step=total_steps)
xz_slice_coords = torch.cat(
(slice_coords_2d[:, :1], torch.zeros_like(
slice_coords_2d[:, :1]), slice_coords_2d[:, -1:]),
dim=-1)
xz_slice_model_input = {'coords': xz_slice_coords.cuda()[None, ...]}
xz_model_out = model(xz_slice_model_input)
sdf_values = xz_model_out['model_out']
sdf_values = dataio.lin2img(sdf_values).squeeze().cpu().numpy()
fig = make_contour_plot(sdf_values)
writer.add_figure(prefix + 'xz_sdf_slice',
fig,
global_step=total_steps)
xy_slice_coords = torch.cat(
(slice_coords_2d[:, :2],
-0.75 * torch.ones_like(slice_coords_2d[:, :1])),
dim=-1)
xy_slice_model_input = {'coords': xy_slice_coords.cuda()[None, ...]}
xy_model_out = model(xy_slice_model_input)
sdf_values = xy_model_out['model_out']
sdf_values = dataio.lin2img(sdf_values).squeeze().cpu().numpy()
fig = make_contour_plot(sdf_values)
writer.add_figure(prefix + 'xy_sdf_slice',
fig,
global_step=total_steps)
min_max_summary(prefix + 'model_out_min_max',
model_output['model_out'], writer, total_steps)
min_max_summary(prefix + 'coords', model_input['coords'], writer,
total_steps)
def write_wave_summary(model, model_input, gt, model_output, writer, total_steps, prefix='train_'):
sl = 256
def scale_percentile(pred, min_perc=1, max_perc=99):
min = np.percentile(pred.cpu().numpy(),1)
max = np.percentile(pred.cpu().numpy(),99)
pred = torch.clamp(pred, min, max)
return (pred - min) / (max-min)
with torch.no_grad():
frames = [0.0, 0.05, 0.1, 0.15, 0.25]
coords = [dataio.get_mgrid((1, sl, sl), dim=3)[None,...].cpu() for f in frames]
for idx, f in enumerate(frames):
coords[idx][..., 0] = f
coords = torch.cat(coords, dim=0)
Nslice = 10
output = torch.zeros(coords.shape[0], coords.shape[1], 1)
split = int(coords.shape[1] / Nslice)
for i in range(Nslice):
pred = model({'coords':coords[:, i*split:(i+1)*split, :]})['model_out']
output[:, i*split:(i+1)*split, :] = pred.cpu()
min_max_summary(prefix + 'pred', pred, writer, total_steps)
pred = output.view(len(frames), 1, sl, sl)
plt.switch_backend('agg')
fig = plt.figure()
plt.subplot(2,2,1)
data = pred[0, :, sl//2, :].numpy().squeeze()
plt.plot(np.linspace(-1, 1, sl), data)
plt.ylim([-0.01, 0.02])
plt.subplot(2,2,2)
data = pred[1, :, sl//2, :].numpy().squeeze()
plt.plot(np.linspace(-1, 1, sl), data)
plt.ylim([-0.01, 0.02])
plt.subplot(2,2,3)
data = pred[2, :, sl//2, :].numpy().squeeze()
plt.plot(np.linspace(-1, 1, sl), data)
plt.ylim([-0.01, 0.02])
plt.subplot(2,2,4)
data = pred[3, :, sl//2, :].numpy().squeeze()
plt.plot(np.linspace(-1, 1, sl), data)
plt.ylim([-0.01, 0.02])
writer.add_figure(prefix + 'center_slice', fig, global_step=total_steps)
pred = torch.clamp(pred, -0.002, 0.002)
writer.add_image(prefix + 'pred_img', make_grid(pred, scale_each=False, normalize=True),
global_step=total_steps)
def val_fn_BRS_posspace(model, model_1P):
# Create a figure
fig = plt.figure(figsize=(5 * num_slices, 5))
fig_error = plt.figure(figsize=(5 * num_slices, 5))
fig_valfunc = plt.figure(figsize=(5 * num_slices, 5))
# Get the meshgrid in the (x, y) coordinate
sidelen = 200
mgrid_coords = dataio.get_mgrid(sidelen)
# Time coordinates
time_coords = torch.ones(mgrid_coords.shape[0], 1) * time
# Start plotting the results
for i in range(num_slices):
coords = torch.cat((time_coords, mgrid_coords), dim=1)
pairwise_coords = {}
# Setup the X-Y coordinates
for j in range(2):
evader_key = '%i' % (j + 1) + 'E'
# X-Y coordinates of the evaders for the full game
xcoords = torch.ones(mgrid_coords.shape[0],
1) * poss[evader_key][i][0]
ycoords = torch.ones(mgrid_coords.shape[0],
1) * poss[evader_key][i][1]
coords = torch.cat((coords, xcoords, ycoords), dim=1)
# X-Y coordinates of the evaders for the pairwise game
pairwise_coords[evader_key] = torch.cat(
(time_coords, xcoords, ycoords, mgrid_coords), dim=1)
# Setup the theta coordinates
coords_ego_theta = ego_vehicle_theta[i] * torch.ones(
mgrid_coords.shape[0], 1) / (math.pi * angle_alpha)
coords = torch.cat((coords, coords_ego_theta), dim=1)
for j in range(2):
evader_key = '%i' % (j + 1) + 'E'
# Theta coordinates of the evaders for the full game
tcoords = torch.ones(
mgrid_coords.shape[0],
1) * thetas[evader_key][i] / (math.pi * angle_alpha)
coords = torch.cat((coords, tcoords), dim=1)
# Theta coordinates of the evaders for the pairwise game
pairwise_coords[evader_key] = torch.cat(
(pairwise_coords[evader_key], tcoords, coords_ego_theta),
dim=1)
model_in = {'coords': coords[:, None, :].cuda()}
model_out = model(model_in)
# Detatch model ouput and reshape
model_out = model_out['model_out'].detach().cpu().numpy()
model_out = model_out.reshape((sidelen, sidelen))
# Unnormalize the value function
norm_to = 0.02
mean = 0.25
var = 0.5
model_out = (model_out * var / norm_to) + mean
# Plot the zero level sets
valfunc = model_out * 1.
model_out = (model_out <= level) * 1.
# Plot the actual data and small aircrafts
ax = fig.add_subplot(1, num_slices, i + 1)
ax_valfunc = fig_valfunc.add_subplot(1, num_slices, i + 1)
ax_error = fig_error.add_subplot(1, num_slices, i + 1)
aircraft_size = 0.2
sA = {}
for j in range(2):
evader_key = '%i' % (j + 1) + 'E'
aircraft_image = scipy.ndimage.rotate(
plt.imread('resources/ego_aircraft.png'),
180.0 * thetas[evader_key][i] / math.pi)
sA[evader_key] = ax.imshow(
aircraft_image,
extent=(poss[evader_key][i][0] - aircraft_size,
poss[evader_key][i][0] + aircraft_size,
poss[evader_key][i][1] - aircraft_size,
poss[evader_key][i][1] + aircraft_size))
ax.plot(poss[evader_key][i][0], poss[evader_key][i][1], "o")
s = ax.imshow(model_out.T,
cmap='bwr_r',
alpha=0.5,
origin='lower',
vmin=-1.,
vmax=1.,
extent=(-1., 1., -1., 1.))
sV1 = ax_valfunc.imshow(valfunc.T,
cmap='bwr_r',
alpha=0.8,
origin='lower',
vmin=-0.2,
vmax=0.2,
extent=(-1., 1., -1., 1.))
sV2 = ax_valfunc.contour(valfunc.T,
cmap='bwr_r',
alpha=0.5,
origin='lower',
vmin=-0.2,
vmax=0.2,
levels=30,
extent=(-1., 1., -1., 1.))
plt.clabel(sV2, levels=30, colors='k')
fig_valfunc.colorbar(sV1)
# Compute and plot pairwise collision sets
sP = {}
model_out_pairwise_sofar = None
valfunc_pairwise = None
for j in range(2):
evader_key = '%i' % (j + 1) + 'E'
model_in_pairwise = {'coords': pairwise_coords[evader_key].cuda()}
model_out_pairwise = model_1P(
model_in_pairwise)['model_out'].detach().cpu().numpy()
model_out_pairwise = model_out_pairwise.reshape((sidelen, sidelen))
norm_to_pairwise = 0.02
mean_pairwise = 0.25
var_pairwise = 0.5
model_out_pairwise = (model_out_pairwise * var_pairwise /
norm_to_pairwise) + mean_pairwise
if model_out_pairwise_sofar is None:
model_out_pairwise_sofar = (model_out_pairwise <= level) * 1.
valfunc_pairwise = model_out_pairwise * 1.
else:
model_out_pairwise_sofar = np.clip(
(model_out_pairwise <= level) * 1. +
model_out_pairwise_sofar, 0., 1.)
valfunc_pairwise = np.minimum(valfunc_pairwise,
model_out_pairwise * 1.0)
s2 = ax.imshow(model_out_pairwise_sofar.T,
cmap='seismic',
alpha=0.5,
origin='lower',
vmin=-1.,
vmax=1.,
extent=(-1., 1., -1., 1.))
# Error plot
error = np.clip(model_out - model_out_pairwise_sofar, 0., 1.)
ax_error.imshow(error.T,
cmap='bwr',
origin='lower',
vmin=-1.,
vmax=1.,
extent=(-1., 1., -1., 1.))
## Append the value functions
val_functions['pairwise'].append(valfunc_pairwise)
val_functions['full'].append(valfunc)
return fig, fig_error, fig_valfunc, val_functions
def write_helmholtz_summary(model,
model_input,
gt,
model_output,
writer,
total_steps,
prefix='train_'):
sl = 256
coords = dataio.get_mgrid(sl)[None, ...].cuda()
def scale_percentile(pred, min_perc=1, max_perc=99):
min = np.percentile(pred.cpu().numpy(), 1)
max = np.percentile(pred.cpu().numpy(), 99)
pred = torch.clamp(pred, min, max)
return (pred - min) / (max - min)
with torch.no_grad():
if 'coords_sub' in model_input:
summary_model_input = {
'coords':
coords.repeat(min(2, model_input['coords_sub'].shape[0]), 1, 1)
}
summary_model_input['coords_sub'] = model_input['coords_sub'][:2,
...]
summary_model_input['img_sub'] = model_input['img_sub'][:2, ...]
pred = model(summary_model_input)['model_out']
else:
pred = model({'coords': coords})['model_out']
if 'pretrain' in gt:
gt['squared_slowness_grid'] = pred[..., -1, None].clone() + 1.
if torch.all(gt['pretrain'] == -1):
gt['squared_slowness_grid'] = torch.clamp(
pred[..., -1, None].clone(), min=-0.999) + 1.
gt['squared_slowness_grid'] = torch.where(
(torch.abs(coords[..., 0, None]) > 0.75) |
(torch.abs(coords[..., 1, None]) > 0.75),
torch.ones_like(gt['squared_slowness_grid']),
gt['squared_slowness_grid'])
pred = pred[..., :-1]
pred = dataio.lin2img(pred)
pred_cmpl = pred[..., 0::2, :, :].cpu().numpy(
) + 1j * pred[..., 1::2, :, :].cpu().numpy()
pred_angle = torch.from_numpy(np.angle(pred_cmpl))
pred_mag = torch.from_numpy(np.abs(pred_cmpl))
min_max_summary(prefix + 'coords', model_input['coords'], writer,
total_steps)
min_max_summary(prefix + 'pred_real', pred[..., 0::2, :, :], writer,
total_steps)
min_max_summary(
prefix + 'pred_abs',
torch.sqrt(pred[..., 0::2, :, :]**2 + pred[..., 1::2, :, :]**2),
writer, total_steps)
min_max_summary(prefix + 'squared_slowness',
gt['squared_slowness_grid'], writer, total_steps)
pred = scale_percentile(pred)
pred_angle = scale_percentile(pred_angle)
pred_mag = scale_percentile(pred_mag)
pred = pred.permute(1, 0, 2, 3)
pred_mag = pred_mag.permute(1, 0, 2, 3)
pred_angle = pred_angle.permute(1, 0, 2, 3)
writer.add_image(prefix + 'pred_real',
make_grid(pred[0::2, :, :, :],
scale_each=False,
normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'pred_imaginary',
make_grid(pred[1::2, :, :, :],
scale_each=False,
normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'pred_angle',
make_grid(pred_angle, scale_each=False, normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'pred_mag',
make_grid(pred_mag, scale_each=False, normalize=True),
global_step=total_steps)
if 'gt' in gt:
gt_field = dataio.lin2img(gt['gt'])
gt_field_cmpl = gt_field[..., 0, :, :].cpu().numpy(
) + 1j * gt_field[..., 1, :, :].cpu().numpy()
gt_angle = torch.from_numpy(np.angle(gt_field_cmpl))
gt_mag = torch.from_numpy(np.abs(gt_field_cmpl))
gt_field = scale_percentile(gt_field)
gt_angle = scale_percentile(gt_angle)
gt_mag = scale_percentile(gt_mag)
writer.add_image(prefix + 'gt_real',
make_grid(gt_field[..., 0, :, :],
scale_each=False,
normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'gt_imaginary',
make_grid(gt_field[..., 1, :, :],
scale_each=False,
normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'gt_angle',
make_grid(gt_angle, scale_each=False, normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'gt_mag',
make_grid(gt_mag, scale_each=False, normalize=True),
global_step=total_steps)
min_max_summary(prefix + 'gt_real', gt_field[..., 0, :, :], writer,
total_steps)
velocity = torch.sqrt(1 / dataio.lin2img(gt['squared_slowness_grid']))[:1]
min_max_summary(prefix + 'velocity', velocity[..., 0, :, :], writer,
total_steps)
velocity = scale_percentile(velocity)
writer.add_image(prefix + 'velocity',
make_grid(velocity[..., 0, :, :],
scale_each=False,
normalize=True),
global_step=total_steps)
if 'squared_slowness_grid' in gt:
writer.add_image(prefix + 'squared_slowness',
make_grid(dataio.lin2img(
gt['squared_slowness_grid'])[:2, :1],
scale_each=False,
normalize=True),
global_step=total_steps)
if 'img_sub' in model_input:
writer.add_image(prefix + 'img',
make_grid(dataio.lin2img(
model_input['img_sub'])[:2, :1],
scale_each=False,
normalize=True),
global_step=total_steps)
if isinstance(model, meta_modules.NeuralProcessImplicit2DHypernetBVP):
hypernet_activation_summary(model, model_input, gt, model_output,
writer, total_steps, prefix)
def plot_sds_with_gradients(gt_contour, gt_normals, mgrid_sds, mgrid_grads):
'''Plot signed distances with gradients.
gt_contour: np array of shape (-1, 2) with x,y coordinates of gt contour.
mgrid_sds: signed distances evaluated on square meshgrid.
mgrid_grads: grads evaluated on square meshgrid.
'''
# Images are square, but flattened - compute the sidelength.
num_pixels = mgrid_grads.shape[1]
sidelen = int(np.sqrt(num_pixels))
mgrid = dataio.get_mgrid(sidelen).detach().cpu().numpy()
x = np.linspace(-1, 1, sidelen)
y = np.linspace(-1, 1, sidelen)
mgrid_grads = mgrid_grads[0, ...].detach().cpu().numpy()
mgrid_grads_mag = np.linalg.norm(mgrid_grads, axis=-1)
gt_contour = gt_contour[0, ...].detach().cpu().numpy()
gt_normals = gt_normals[0, ...].detach().cpu().numpy()
mgrid_sds = mgrid_sds[0, ...].detach().cpu().numpy()
mgrid_sds = mgrid_sds.reshape(sidelen, sidelen)
fig, (axa, axb, axc) = plt.subplots(nrows=1, ncols=3, figsize=(30, 10))
axa.cla(), axb.cla(), axc.cla()
# PLOT A: Ground truth mesh
axa.set_xlim([mgrid.min(),
mgrid.max()]), axa.set_ylim([mgrid.min(),
mgrid.max()])
axa.plot(gt_contour[..., 1], gt_contour[..., 0] * -1.)
q = axa.quiver(gt_contour[..., 1],
gt_contour[..., 0] * -1.,
gt_normals[..., 1],
gt_normals[..., 0] * -1.,
scale=25.)
axa.set_title('Ground truth level set')
# PLOT A: Predicted SDs
axb.imshow(mgrid_sds)
axb.contour(mgrid_sds, levels=[0], colors='k', linestyles='-')
axb.set_title('Predicted Signed Distance')
# PLOT B: GRADIENT DIRECTIONS
grad_subsample = 1
quiver_coords = mgrid.reshape(
sidelen, sidelen,
2)[::grad_subsample, ::grad_subsample, :].reshape(-1, 2)
quiver_mgrid_grads = mgrid_grads.reshape(
sidelen, sidelen,
2)[::grad_subsample, ::grad_subsample, :].reshape(-1, 2)
axc.set_xlim([mgrid.min(),
mgrid.max()]), axc.set_ylim([mgrid.min(),
mgrid.max()])
axc.set_xlabel('x'), axc.set_ylabel('y')
axc.set_xticks([mgrid.min(), 0., mgrid.max()
]), axc.set_yticks([mgrid.min(), 0.,
mgrid.max()])
q = axc.quiver(quiver_coords[..., 1], quiver_coords[..., 0] * -1,
quiver_mgrid_grads[..., 1], quiver_mgrid_grads[..., 0] * -1)
axc.set_title('Orientations of Gradients')
plt.show()
in_features=img_dataset_test.img_channels,
out_features=img_dataset_test.img_channels,
image_resolution=image_resolution,
partial_conv=opt.partial_conv)
model.cuda()
model.eval()
root_path = os.path.join(opt.logging_root, opt.experiment_name)
utils.cond_mkdir(root_path)
# Load checkpoint
model.load_state_dict(torch.load(opt.checkpoint_path))
# First experiment: Upsample training image
model_input = {
'coords': dataio.get_mgrid(image_resolution)[None, :].cuda(),
'img_sparse': coord_dataset_train[0][0]['img_sparse'].unsqueeze(0).cuda()
}
model_output = model(model_input)
out_img = dataio.lin2img(model_output['model_out'],
image_resolution).squeeze().permute(
1, 2, 0).detach().cpu().numpy()
out_img += 1
out_img /= 2.
out_img = np.clip(out_img, 0., 1.)
imageio.imwrite(os.path.join(root_path, 'upsampled_train.png'), out_img)
# Second experiment: sample larger range
model_input = {
