def write_image_summary(image_resolution, model, model_input, gt,
model_output, writer, total_steps, prefix='train_'):
gt_img = dataio.lin2img(gt['img'], image_resolution)
pred_img = dataio.lin2img(model_output['model_out'], image_resolution)
img_gradient = diff_operators.gradient(model_output['model_out'], model_output['model_in'])
img_laplace = diff_operators.laplace(model_output['model_out'], model_output['model_in'])
output_vs_gt = torch.cat((gt_img, pred_img), dim=-1)
writer.add_image(prefix + 'gt_vs_pred', make_grid(output_vs_gt, scale_each=False, normalize=True),
global_step=total_steps)
pred_img = dataio.rescale_img((pred_img+1)/2, mode='clamp').permute(0,2,3,1).squeeze(0).detach().cpu().numpy()
pred_grad = dataio.grads2img(dataio.lin2img(img_gradient)).permute(1,2,0).squeeze().detach().cpu().numpy()
pred_lapl = cv2.cvtColor(cv2.applyColorMap(dataio.to_uint8(dataio.rescale_img(
dataio.lin2img(img_laplace), perc=2).permute(0,2,3,1).squeeze(0).detach().cpu().numpy()), cmapy.cmap('RdBu')), cv2.COLOR_BGR2RGB)
gt_img = dataio.rescale_img((gt_img+1) / 2, mode='clamp').permute(0, 2, 3, 1).squeeze(0).detach().cpu().numpy()
gt_grad = dataio.grads2img(dataio.lin2img(gt['gradients'])).permute(1, 2, 0).squeeze().detach().cpu().numpy()
gt_lapl = cv2.cvtColor(cv2.applyColorMap(dataio.to_uint8(dataio.rescale_img(
dataio.lin2img(gt['laplace']), perc=2).permute(0, 2, 3, 1).squeeze(0).detach().cpu().numpy()), cmapy.cmap('RdBu')), cv2.COLOR_BGR2RGB)
writer.add_image(prefix + 'pred_img', torch.from_numpy(pred_img).permute(2, 0, 1), global_step=total_steps)
writer.add_image(prefix + 'pred_grad', torch.from_numpy(pred_grad).permute(2, 0, 1), global_step=total_steps)
writer.add_image(prefix + 'pred_lapl', torch.from_numpy(pred_lapl).permute(2,0,1), global_step=total_steps)
writer.add_image(prefix + 'gt_img', torch.from_numpy(gt_img).permute(2,0,1), global_step=total_steps)
writer.add_image(prefix + 'gt_grad', torch.from_numpy(gt_grad).permute(2, 0, 1), global_step=total_steps)
writer.add_image(prefix + 'gt_lapl', torch.from_numpy(gt_lapl).permute(2, 0, 1), global_step=total_steps)
write_psnr(dataio.lin2img(model_output['model_out'], image_resolution),
dataio.lin2img(gt['img'], image_resolution), writer, total_steps, prefix+'img_')
def write_laplace_summary(model, model_input, gt, model_output, writer, total_steps, prefix='train_'):
# Plot comparison images
gt_img = dataio.lin2img(gt['img'])
pred_img = dataio.lin2img(model_output['model_out'])
output_vs_gt = torch.cat((dataio.rescale_img(gt_img), dataio.rescale_img(pred_img,perc=1e-2)), dim=-1)
writer.add_image(prefix + 'comp_gt_vs_pred', make_grid(output_vs_gt, scale_each=False, normalize=True),
global_step=total_steps)
# Plot comparisons laplacian (this is what has been fitted)
gt_laplace = dataio.lin2img(gt['laplace'])
pred_laplace = diff_operators.laplace(model_output['model_out'], model_output['model_in'])
pred_laplace = dataio.lin2img(pred_laplace)
output_vs_gt_laplace = torch.cat((gt_laplace, pred_laplace), dim=-1)
writer.add_image(prefix + 'comp_gt_vs_pred_laplace', make_grid(output_vs_gt_laplace, scale_each=False, normalize=True),
global_step=total_steps)
# Plot image gradient
img_gradient = diff_operators.gradient(model_output['model_out'], model_output['model_in'])
grads_img = dataio.grads2img(dataio.lin2img(img_gradient))
writer.add_image(prefix + 'pred_grad', make_grid(grads_img, scale_each=False, normalize=True),
global_step=total_steps)
# Plot gt image
writer.add_image(prefix + 'gt_img', make_grid(gt_img, scale_each=False, normalize=True),
global_step=total_steps)
# Plot gt laplacian
# writer.add_image(prefix + 'gt_laplace', make_grid(gt_laplace, scale_each=False, normalize=True),
# global_step=total_steps)
gt_laplace_img = dataio.to_uint8(dataio.to_numpy(dataio.rescale_img(gt_laplace, 'scale', 1)))
gt_laplace_img = cv2.applyColorMap(gt_laplace_img.squeeze(), cmapy.cmap('RdBu'))
gt_laplace_img = cv2.cvtColor(gt_laplace_img, cv2.COLOR_BGR2RGB)
writer.add_image(prefix + 'gt_lapl', torch.from_numpy(gt_laplace_img).permute(2, 0, 1), global_step=total_steps)
# Plot pred image
writer.add_image(prefix + 'pred_img', make_grid(pred_img, scale_each=False, normalize=True),
global_step=total_steps)
# Plot pred gradient
pred_gradients = diff_operators.gradient(model_output['model_out'], model_output['model_in'])
pred_grads_img = dataio.grads2img(dataio.lin2img(pred_gradients))
writer.add_image(prefix + 'pred_grad', make_grid(pred_grads_img, scale_each=False, normalize=True),
global_step=total_steps)
# Plot pred laplacian
# writer.add_image(prefix + 'pred_lapl', make_grid(pred_laplace, scale_each=False, normalize=True),
# global_step=total_steps)
pred_laplace_img = dataio.to_uint8(dataio.to_numpy(dataio.rescale_img(pred_laplace,'scale',1)))
pred_laplace_img = cv2.applyColorMap(pred_laplace_img.squeeze(),cmapy.cmap('RdBu'))
pred_laplace_img = cv2.cvtColor(pred_laplace_img, cv2.COLOR_BGR2RGB)
writer.add_image(prefix + 'pred_lapl', torch.from_numpy(pred_laplace_img).permute(2,0,1), global_step=total_steps)
min_max_summary(prefix + 'coords', model_input['coords'], writer, total_steps)
min_max_summary(prefix + 'gt_laplace', gt_laplace, writer, total_steps)
min_max_summary(prefix + 'pred_laplace', pred_laplace, writer, total_steps)
min_max_summary(prefix + 'pred_img', pred_img, writer, total_steps)
min_max_summary(prefix + 'gt_img', gt_img, writer, total_steps)
def write_gradcomp_summary(model, model_input, gt, model_output, writer, total_steps, prefix='train_'):
# Plot gt gradients (this is what has been fitted)
gt_gradients = gt['gradients']
gt_grads_img = dataio.grads2img(dataio.lin2img(gt_gradients))
pred_gradients = diff_operators.gradient(model_output['model_out'], model_output['model_in'])
pred_grads_img = dataio.grads2img(dataio.lin2img(pred_gradients))
output_vs_gt_gradients = torch.cat((gt_grads_img, pred_grads_img), dim=-1)
writer.add_image(prefix + 'comp_gt_vs_pred_gradients', make_grid(output_vs_gt_gradients, scale_each=False, normalize=True),
global_step=total_steps)
# Plot gt
gt_grads1 = gt['grads1']
gt_grads1_img = dataio.grads2img(dataio.lin2img(gt_grads1))
gt_grads2 = gt['grads2']
gt_grads2_img = dataio.grads2img(dataio.lin2img(gt_grads2))
writer.add_image(prefix + 'gt_grads1', make_grid(gt_grads1_img, scale_each=False, normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'gt_grads2', make_grid(gt_grads2_img, scale_each=False, normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'gt_gradcomp', make_grid(gt_grads_img, scale_each=False, normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'pred_gradcomp', make_grid(pred_grads_img, scale_each=False, normalize=True),
global_step=total_steps)
# Plot gt image
gt_img1 = dataio.lin2img(gt['img1'])
gt_img2 = dataio.lin2img(gt['img2'])
writer.add_image(prefix + 'gt_img1', make_grid(gt_img1, scale_each=False, normalize=True),
global_step=total_steps)
writer.add_image(prefix + 'gt_img2', make_grid(gt_img2, scale_each=False, normalize=True),
global_step=total_steps)
# Plot pred compo image
pred_img = dataio.rescale_img(dataio.lin2img(model_output['model_out']))
writer.add_image(prefix + 'pred_comp_img', make_grid(pred_img, scale_each=False, normalize=True),
global_step=total_steps)
min_max_summary(prefix + 'coords', model_input['coords'], writer, total_steps)
min_max_summary(prefix + 'gt_laplace', gt_gradients, writer, total_steps)
min_max_summary(prefix + 'pred_img', pred_img, writer, total_steps)