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 laplace_mse(model_output, gt):
# compute laplacian on the model
laplace = diff_operators.laplace(model_output['model_out'],
model_output['model_in'])
# compare them with the ground truth
laplace_loss = torch.mean((laplace - gt['laplace'])**2)
return {'laplace_loss': laplace_loss}
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)
