def add_confidences(disp_X_recon, lmids, loc): means = lm_confs[:, lmids].mean(axis=1) colors = vis.color_map(to_numpy(1 - means), cmap=plt.cm.jet, vmin=0.0, vmax=0.4) return vis.add_error_to_images(disp_X_recon, means, loc=loc, format_string='{:>4.2f}', colors=colors)
def generate_images(self, z): train_state_D = self.saae.D.training train_state_P = self.saae.P.training self.saae.D.eval() self.saae.P.eval() loc_err_gan = 'tr' with torch.no_grad(): X_gen_vis = self.saae.P(z)[:, :3] err_gan_gen = self.saae.D(X_gen_vis) imgs = vis.to_disp_images(X_gen_vis, denorm=True) self.saae.D.train(train_state_D) self.saae.P.train(train_state_P) return vis.add_error_to_images(imgs, errors=1 - err_gan_gen, loc=loc_err_gan, format_string='{:.2f}', vmax=1.0)
def visualize_batch(self, batch, X_recon, ssim_maps, nimgs=8, ds=None, wait=0): nimgs = min(nimgs, len(batch)) train_state_D = self.saae.D.training train_state_Q = self.saae.Q.training train_state_P = self.saae.P.training self.saae.D.eval() self.saae.Q.eval() self.saae.P.eval() loc_err_gan = 'tr' text_size_errors = 0.65 input_images = vis.to_disp_images(batch.images[:nimgs], denorm=True) target_images = batch.target_images if batch.target_images is not None else batch.images disp_images = vis.to_disp_images(target_images[:nimgs], denorm=True) # draw GAN score if self.args.with_gan: with torch.no_grad(): err_gan_inputs = self.saae.D(batch.images[:nimgs]) disp_images = vis.add_error_to_images(disp_images, errors=1 - err_gan_inputs, loc=loc_err_gan, format_string='{:>5.2f}', vmax=1.0) # disp_images = vis.add_landmarks_to_images(disp_images, batch.landmarks[:nimgs], color=(0,1,0), radius=1, # draw_wireframe=False) rows = [vis.make_grid(disp_images, nCols=nimgs, normalize=False)] recon_images = vis.to_disp_images(X_recon[:nimgs], denorm=True) disp_X_recon = recon_images.copy() print_stats = True if print_stats: # lm_ssim_errs = None # if batch.landmarks is not None: # lm_recon_errs = lmutils.calc_landmark_recon_error(batch.images[:nimgs], X_recon[:nimgs], batch.landmarks[:nimgs], reduction='none') # disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_recon_errs, size=text_size_errors, loc='bm', # format_string='({:>3.1f})', vmin=0, vmax=10) # lm_ssim_errs = lmutils.calc_landmark_ssim_error(batch.images[:nimgs], X_recon[:nimgs], batch.landmarks[:nimgs]) # disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_ssim_errs.mean(axis=1), size=text_size_errors, loc='bm-1', # format_string='({:>3.2f})', vmin=0.2, vmax=0.8) X_recon_errs = 255.0 * torch.abs(batch.images - X_recon).reshape( len(batch.images), -1).mean(dim=1) # disp_X_recon = vis.add_landmarks_to_images(disp_X_recon, batch.landmarks[:nimgs], radius=1, color=None, # lm_errs=lm_ssim_errs, draw_wireframe=False) disp_X_recon = vis.add_error_to_images(disp_X_recon[:nimgs], errors=X_recon_errs, size=text_size_errors, format_string='{:>4.1f}') if self.args.with_gan: with torch.no_grad(): err_gan = self.saae.D(X_recon[:nimgs]) disp_X_recon = vis.add_error_to_images( disp_X_recon, errors=1 - err_gan, loc=loc_err_gan, format_string='{:>5.2f}', vmax=1.0) ssim = np.zeros(nimgs) for i in range(nimgs): data_range = 255.0 if input_images[0].dtype == np.uint8 else 1.0 ssim[i] = compare_ssim(input_images[i], recon_images[i], data_range=data_range, multichannel=True) disp_X_recon = vis.add_error_to_images(disp_X_recon, 1 - ssim, loc='bl-1', size=text_size_errors, format_string='{:>4.2f}', vmin=0.2, vmax=0.8) if ssim_maps is not None: disp_X_recon = vis.add_error_to_images( disp_X_recon, ssim_maps.reshape(len(ssim_maps), -1).mean(axis=1), size=text_size_errors, loc='bl-2', format_string='{:>4.2f}', vmin=0.0, vmax=0.4) rows.append(vis.make_grid(disp_X_recon, nCols=nimgs)) if ssim_maps is not None: disp_ssim_maps = to_numpy( nn.denormalized(ssim_maps)[:nimgs].transpose(0, 2, 3, 1)) for i in range(len(disp_ssim_maps)): disp_ssim_maps[i] = vis.color_map( disp_ssim_maps[i].mean(axis=2), vmin=0.0, vmax=2.0) grid_ssim_maps = vis.make_grid(disp_ssim_maps, nCols=nimgs) cv2.imshow('ssim errors', cv2.cvtColor(grid_ssim_maps, cv2.COLOR_RGB2BGR)) self.saae.D.train(train_state_D) self.saae.Q.train(train_state_Q) self.saae.P.train(train_state_P) f = 1 disp_rows = vis.make_grid(rows, nCols=1, normalize=False, fx=f, fy=f) wnd_title = 'recon errors ' if ds is not None: wnd_title += ds.__class__.__name__ cv2.imshow(wnd_title, cv2.cvtColor(disp_rows, cv2.COLOR_RGB2BGR)) cv2.waitKey(wait)
def visualize_vessels(images, X_recon, vessel_hm, pred_vessel_hm=None, ds=None, wait=0, horizontal=False, f=1.0, overlay_heatmaps_input=True, overlay_heatmaps_recon=True, scores=None, nimgs=5): nimgs = min(nimgs, len(images)) images = images[:nimgs] rows = [] input_images = vis.to_disp_images(images[:nimgs], denorm=True) disp_images = vis.to_disp_images(images[:nimgs], denorm=True) disp_images = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in disp_images ] rows.append(vis.make_grid(disp_images, nCols=nimgs, normalize=False)) recon_images = vis.to_disp_images(X_recon[:nimgs], denorm=True) disp_X_recon = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in recon_images.copy() ] if vessel_hm is not None and overlay_heatmaps_input: vessel_hm = to_numpy(vessel_hm[:nimgs]) disp_images = [ vis.overlay_heatmap(disp_images[i], vessel_hm[i, 0], 0.5) for i in range(len(vessel_hm)) ] rows.append(vis.make_grid(disp_images, nCols=nimgs, normalize=False)) if pred_vessel_hm is not None and overlay_heatmaps_recon: pred_vessel_hm = to_numpy(pred_vessel_hm[:nimgs]) disp_X_recon_overlay = [ vis.overlay_heatmap(disp_X_recon[i], pred_vessel_hm[i, 0], 1.0) for i in range(len(pred_vessel_hm)) ] if scores is not None: disp_X_recon_overlay = vis.add_error_to_images( disp_X_recon_overlay, scores, loc='tr', format_string='{:.3f}') rows.append(vis.make_grid(disp_X_recon_overlay, nCols=nimgs)) rows.append(vis.make_grid(disp_X_recon, nCols=nimgs)) if horizontal: assert (nimgs == 1) disp_rows = vis.make_grid(rows, nCols=4) else: disp_rows = vis.make_grid(rows, nCols=1) wnd_title = 'Predicted vessels ' if ds is not None: wnd_title += ds.__class__.__name__ cv2.imshow(wnd_title, cv2.cvtColor(disp_rows, cv2.COLOR_RGB2BGR)) cv2.waitKey(wait)
def visualize_batch(images, landmarks, X_recon, X_lm_hm, lm_preds_max, lm_heatmaps=None, target_images=None, lm_preds_cnn=None, ds=None, wait=0, ssim_maps=None, landmarks_to_draw=None, ocular_norm='outer', horizontal=False, f=1.0, overlay_heatmaps_input=False, overlay_heatmaps_recon=False, clean=False, landmarks_only_outline=range(17), landmarks_no_outline=range(17, 68)): gt_color = (0, 255, 0) pred_color = (0, 0, 255) image_size = images.shape[3] assert image_size in [128, 256] nimgs = min(10, len(images)) images = nn.atleast4d(images)[:nimgs] num_landmarks = lm_preds_max.shape[1] if landmarks_to_draw is None: landmarks_to_draw = range(num_landmarks) nme_per_lm = None if landmarks is None: # print('num landmarks', lmcfg.NUM_LANDMARKS) lm_gt = np.zeros((nimgs, num_landmarks, 2)) else: lm_gt = nn.atleast3d(to_numpy(landmarks))[:nimgs] nme_per_lm = calc_landmark_nme(lm_gt, lm_preds_max[:nimgs], ocular_norm=ocular_norm, image_size=image_size) lm_ssim_errs = 1 - calc_landmark_ssim_score(images, X_recon[:nimgs], lm_gt) lm_confs = None # show landmark heatmaps pred_heatmaps = None if X_lm_hm is not None: pred_heatmaps = to_single_channel_heatmap(to_numpy(X_lm_hm[:nimgs])) pred_heatmaps = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in pred_heatmaps ] gt_heatmaps = None if lm_heatmaps is not None: gt_heatmaps = to_single_channel_heatmap( to_numpy(lm_heatmaps[:nimgs])) gt_heatmaps = np.array([ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in gt_heatmaps ]) show_landmark_heatmaps(pred_heatmaps, gt_heatmaps, nimgs, f=1) lm_confs = to_numpy(X_lm_hm).reshape(X_lm_hm.shape[0], X_lm_hm.shape[1], -1).max(axis=2) # resize images for display and scale landmarks accordingly lm_preds_max = lm_preds_max[:nimgs] * f if lm_preds_cnn is not None: lm_preds_cnn = lm_preds_cnn[:nimgs] * f lm_gt *= f input_images = vis.to_disp_images(images[:nimgs], denorm=True) if target_images is not None: disp_images = vis.to_disp_images(target_images[:nimgs], denorm=True) else: disp_images = vis.to_disp_images(images[:nimgs], denorm=True) disp_images = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in disp_images ] recon_images = vis.to_disp_images(X_recon[:nimgs], denorm=True) disp_X_recon = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in recon_images.copy() ] # overlay landmarks on input images if pred_heatmaps is not None and overlay_heatmaps_input: disp_images = [ vis.overlay_heatmap(disp_images[i], pred_heatmaps[i]) for i in range(len(pred_heatmaps)) ] if pred_heatmaps is not None and overlay_heatmaps_recon: disp_X_recon = [ vis.overlay_heatmap(disp_X_recon[i], pred_heatmaps[i]) for i in range(len(pred_heatmaps)) ] # # Show input images # disp_images = vis.add_landmarks_to_images(disp_images, lm_gt[:nimgs], color=gt_color) disp_images = vis.add_landmarks_to_images(disp_images, lm_preds_max[:nimgs], lm_errs=nme_per_lm, color=pred_color, draw_wireframe=False, gt_landmarks=lm_gt, draw_gt_offsets=True) # disp_images = vis.add_landmarks_to_images(disp_images, lm_gt[:nimgs], color=(1,1,1), radius=1, # draw_dots=True, draw_wireframe=True, landmarks_to_draw=landmarks_to_draw) # disp_images = vis.add_landmarks_to_images(disp_images, lm_preds_max[:nimgs], lm_errs=nme_per_lm, # color=(1.0, 0.0, 0.0), # draw_dots=True, draw_wireframe=True, radius=1, # gt_landmarks=lm_gt, draw_gt_offsets=False, # landmarks_to_draw=landmarks_to_draw) # # Show reconstructions # X_recon_errs = 255.0 * torch.abs(images - X_recon[:nimgs]).reshape( len(images), -1).mean(dim=1) if not clean: disp_X_recon = vis.add_error_to_images(disp_X_recon[:nimgs], errors=X_recon_errs, format_string='{:>4.1f}') # modes of heatmaps # disp_X_recon = [overlay_heatmap(disp_X_recon[i], pred_heatmaps[i]) for i in range(len(pred_heatmaps))] if not clean: lm_errs_max = calc_landmark_nme_per_img(lm_gt, lm_preds_max, ocular_norm, landmarks_no_outline, image_size=image_size) lm_errs_max_outline = calc_landmark_nme_per_img(lm_gt, lm_preds_max, ocular_norm, landmarks_only_outline, image_size=image_size) lm_errs_max_all = calc_landmark_nme_per_img( lm_gt, lm_preds_max, ocular_norm, list(landmarks_only_outline) + list(landmarks_no_outline), image_size=image_size) disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_errs_max, loc='br-2', format_string='{:>5.2f}', vmax=15) disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_errs_max_outline, loc='br-1', format_string='{:>5.2f}', vmax=15) disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_errs_max_all, loc='br', format_string='{:>5.2f}', vmax=15) disp_X_recon = vis.add_landmarks_to_images(disp_X_recon, lm_gt, color=gt_color, draw_wireframe=True) # disp_X_recon = vis.add_landmarks_to_images(disp_X_recon, lm_preds_max[:nimgs], # color=pred_color, draw_wireframe=False, # lm_errs=nme_per_lm, lm_confs=lm_confs, # lm_rec_errs=lm_ssim_errs, gt_landmarks=lm_gt, # draw_gt_offsets=True, draw_dots=True) disp_X_recon = vis.add_landmarks_to_images(disp_X_recon, lm_preds_max[:nimgs], color=pred_color, draw_wireframe=True, gt_landmarks=lm_gt, draw_gt_offsets=True, lm_errs=nme_per_lm, draw_dots=True, radius=2) def add_confidences(disp_X_recon, lmids, loc): means = lm_confs[:, lmids].mean(axis=1) colors = vis.color_map(to_numpy(1 - means), cmap=plt.cm.jet, vmin=0.0, vmax=0.4) return vis.add_error_to_images(disp_X_recon, means, loc=loc, format_string='{:>4.2f}', colors=colors) # disp_X_recon = add_confidences(disp_X_recon, lmcfg.LANDMARKS_NO_OUTLINE, 'bm-2') # disp_X_recon = add_confidences(disp_X_recon, lmcfg.LANDMARKS_ONLY_OUTLINE, 'bm-1') # disp_X_recon = add_confidences(disp_X_recon, lmcfg.ALL_LANDMARKS, 'bm') # print ssim errors ssim = np.zeros(nimgs) for i in range(nimgs): ssim[i] = compare_ssim(input_images[i], recon_images[i], data_range=1.0, multichannel=True) if not clean: disp_X_recon = vis.add_error_to_images(disp_X_recon, 1 - ssim, loc='bl-1', format_string='{:>4.2f}', vmax=0.8, vmin=0.2) # print ssim torch errors if ssim_maps is not None and not clean: disp_X_recon = vis.add_error_to_images(disp_X_recon, ssim_maps.reshape( len(ssim_maps), -1).mean(axis=1), loc='bl-2', format_string='{:>4.2f}', vmin=0.0, vmax=0.4) rows = [vis.make_grid(disp_images, nCols=nimgs, normalize=False)] rows.append(vis.make_grid(disp_X_recon, nCols=nimgs)) if ssim_maps is not None: disp_ssim_maps = to_numpy( nn.denormalized(ssim_maps)[:nimgs].transpose(0, 2, 3, 1)) for i in range(len(disp_ssim_maps)): disp_ssim_maps[i] = vis.color_map(disp_ssim_maps[i].mean(axis=2), vmin=0.0, vmax=2.0) grid_ssim_maps = vis.make_grid(disp_ssim_maps, nCols=nimgs, fx=f, fy=f) cv2.imshow('ssim errors', cv2.cvtColor(grid_ssim_maps, cv2.COLOR_RGB2BGR)) if horizontal: assert (nimgs == 1) disp_rows = vis.make_grid(rows, nCols=2) else: disp_rows = vis.make_grid(rows, nCols=1) wnd_title = 'Predicted Landmarks ' if ds is not None: wnd_title += ds.__class__.__name__ cv2.imshow(wnd_title, cv2.cvtColor(disp_rows, cv2.COLOR_RGB2BGR)) cv2.waitKey(wait)