def visualize_images(X, X_lm_hm, landmarks=None, show_recon=True, show_landmarks=True, show_heatmaps=False, draw_wireframe=False, smoothing_level=2, heatmap_opacity=0.8, f=1): if show_recon: disp_X = vis.to_disp_images(X, denorm=True) else: disp_X = vis.to_disp_images(torch.zeros_like(X), denorm=False) heatmap_opacity = 1 if X_lm_hm is not None: if smoothing_level > 0: X_lm_hm = smooth_heatmaps(X_lm_hm) if smoothing_level > 1: X_lm_hm = smooth_heatmaps(X_lm_hm) if show_heatmaps: pred_heatmaps = to_single_channel_heatmap(to_numpy(X_lm_hm)) pred_heatmaps = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_CUBIC) for im in pred_heatmaps ] disp_X = [ vis.overlay_heatmap(disp_X[i], pred_heatmaps[i], heatmap_opacity) for i in range(len(pred_heatmaps)) ] if show_landmarks and landmarks is not None: pred_color = (0, 255, 255) disp_X = vis.add_landmarks_to_images(disp_X, landmarks, color=pred_color, draw_wireframe=draw_wireframe) return disp_X
result['landmarks'], self.landmark_sigma, self.landmark_ids) return result if __name__ == '__main__': from utils.nn import Batch, to_numpy import utils.common from landmarks import lmconfig as lmcfg utils.common.init_random(3) lmcfg.config_landmarks('wflw') ds = WFLW(train=True, deterministic=True, use_cache=True, daug=0) # ds.filter_labels({'pose':0, 'blur':0, 'occlusion':1}) dl = td.DataLoader(ds, batch_size=1, shuffle=False, num_workers=0) cfg.WITH_LANDMARK_LOSS = False for data in dl: batch = Batch(data, gpu=False) images = vis._to_disp_images(batch.images, denorm=True) # lms = lmutils.convert_landmarks(to_numpy(batch.landmarks), lmutils.LM98_TO_LM68) lms = batch.landmarks images = vis.add_landmarks_to_images(images, lms, draw_wireframe=False, color=(0, 255, 0), radius=3) vis.vis_square(images, nCols=1, fx=1., fy=1., normalize=False)
if __name__ == "__main__": model = "./demo" net = fabrec.load_net(model, num_landmarks=17) net.eval() im_dir = "./images" img0 = "5.jpg" with torch.no_grad(): img = load_image(im_dir, img0, channels=1, size=(256, 256)) img /= 256 img = ToTensor()(img) img = cfg.RHPE_NORMALIZER(img) img = torch.unsqueeze(img, 0) print(img, img.shape) X_recon, lms_in_crop, X_lm_hm = net.detect_landmarks(img) outputs = add_landmarks_to_images(img, lms_in_crop, skeleton=HandDataset.SKELETON, denorm=True, draw_wireframe=True, color=(0, 255, 255)) X_recon = X_recon[0, :, :, :] X_recon = to_disp_image(X_recon, True) cv2.imwrite("images/outputs.jpg", outputs[0]) cv2.imwrite("images/reconstruct.jpg", X_recon)
def visualize_batch_CVPR(images, landmarks, X_recon, X_lm_hm, lm_preds, lm_heatmaps=None, ds=None, wait=0, horizontal=False, f=1.0, radius=2): gt_color = (0, 255, 0) pred_color = (0, 255, 255) nimgs = min(10, len(images)) images = nn.atleast4d(images)[:nimgs] if landmarks is None: print('num landmarks', lmcfg.NUM_LANDMARKS) lm_gt = np.zeros((nimgs, lmcfg.NUM_LANDMARKS, 2)) else: lm_gt = nn.atleast3d(to_numpy(landmarks))[:nimgs] # 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 = lm_preds[:nimgs] * f lm_gt *= f rows = [] 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() ] rows.append(vis.make_grid(disp_X_recon, nCols=nimgs)) # recon_images = vis._to_disp_images(X_recon[:nimgs], denorm=True) disp_X_recon_pred = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in recon_images.copy() ] disp_X_recon_pred = vis.add_landmarks_to_images(disp_X_recon_pred, lm_preds, color=pred_color, radius=radius) rows.append(vis.make_grid(disp_X_recon_pred, nCols=nimgs)) disp_X_recon_gt = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in recon_images.copy() ] disp_X_recon_gt = vis.add_landmarks_to_images(disp_X_recon_gt, lm_gt, color=gt_color, radius=radius) rows.append(vis.make_grid(disp_X_recon_gt, nCols=nimgs)) # overlay landmarks on images disp_X_recon_hm = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in recon_images.copy() ] disp_X_recon_hm = [ vis.overlay_heatmap(disp_X_recon_hm[i], pred_heatmaps[i]) for i in range(len(pred_heatmaps)) ] rows.append(vis.make_grid(disp_X_recon_hm, nCols=nimgs)) # input images with prediction (and ground truth) disp_images_pred = vis._to_disp_images(images[:nimgs], denorm=True) disp_images_pred = [ cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST) for im in disp_images_pred ] # disp_images_pred = vis.add_landmarks_to_images(disp_images_pred, lm_gt, color=gt_color, radius=radius) disp_images_pred = vis.add_landmarks_to_images(disp_images_pred, lm_preds, color=pred_color, radius=radius) rows.append(vis.make_grid(disp_images_pred, nCols=nimgs)) if horizontal: assert (nimgs == 1) disp_rows = vis.make_grid(rows, nCols=2) else: disp_rows = vis.make_grid(rows, nCols=1) 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_batch(images, landmarks, X_recon, X_lm_hm, lm_preds_max, lm_heatmaps=None, images_mod=None, lm_preds_cnn=None, ds=None, wait=0, ssim_maps=None, landmarks_to_draw=lmcfg.ALL_LANDMARKS, ocular_norm='outer', horizontal=False, f=1.0, overlay_heatmaps_input=False, overlay_heatmaps_recon=False, clean=False): gt_color = (0, 255, 0) pred_color = (0, 0, 255) nimgs = min(10, len(images)) images = nn.atleast4d(images)[:nimgs] nme_per_lm = None if landmarks is None: # print('num landmarks', lmcfg.NUM_LANDMARKS) lm_gt = np.zeros((nimgs, lmcfg.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) 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 images_mod is not None: disp_images = vis._to_disp_images(images_mod[: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=ocular_norm, landmarks_to_eval=lmcfg.LANDMARKS_NO_OUTLINE) lm_errs_max_outline = calc_landmark_nme_per_img( lm_gt, lm_preds_max, ocular_norm=ocular_norm, landmarks_to_eval=lmcfg.LANDMARKS_ONLY_OUTLINE) lm_errs_max_all = calc_landmark_nme_per_img( lm_gt, lm_preds_max, ocular_norm=ocular_norm, landmarks_to_eval=lmcfg.ALL_LANDMARKS) 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_confs(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_confs(disp_X_recon, lmcfg.LANDMARKS_NO_OUTLINE, 'bm-2') # disp_X_recon = add_confs(disp_X_recon, lmcfg.LANDMARKS_ONLY_OUTLINE, 'bm-1') # disp_X_recon = add_confs(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( ds_utils.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)
return np.concatenate(sizes) if __name__ == '__main__': import torch from utils import vis from utils.nn import Batch from datasets import ds_utils from datasets import affectnet torch.manual_seed(0) torch.cuda.manual_seed_all(0) train = True datasets = [ affectnet.AffectNet(train=train, max_samples=1000), # vggface2.VggFace2(train=train, max_samples=1000), ] multi_ds = MultiFaceDataset(datasets, train=True, max_samples=5000) print(multi_ds) dl = td.DataLoader(multi_ds, batch_size=40, shuffle=False, num_workers=0) for data in dl: batch = Batch(data, gpu=False) inputs = batch.images.clone() ds_utils.denormalize(inputs) imgs = vis.add_landmarks_to_images(inputs.numpy(), batch.landmarks.numpy()) # imgs = vis.add_pose_to_images(inputs.numpy(), batch.poses.numpy()) # imgs = vis.add_emotion_to_images(imgs, batch.emotions.numpy()) vis.vis_square(imgs, nCols=20, fx=0.6, fy=0.6, normalize=False)
def visualize_batch(batch, X_recon, X_lm_hm, lm_preds_max, lm_preds_cnn=None, ds=None, wait=0, ssim_maps=None, landmarks_to_draw=lmcfg.LANDMARKS_TO_EVALUATE, ocular_norm='pupil', horizontal=False, f=1.0): nimgs = min(10, len(batch)) gt_color = (0, 1, 0) 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 ] if batch.lm_heatmaps is not None: gt_heatmaps = to_single_channel_heatmap( to_numpy(batch.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) # scale landmarks 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 = to_numpy(batch.landmarks[:nimgs]) * f if lm_gt.shape[1] == 98: lm_gt = convert_landmarks(lm_gt, LM98_TO_LM68) input_images = vis._to_disp_images(batch.images[:nimgs], denorm=True) if batch.images_mod is not None: disp_images = vis._to_disp_images(batch.images_mod[:nimgs], denorm=True) else: disp_images = vis._to_disp_images(batch.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() ] # draw landmarks to input images if pred_heatmaps is not None: disp_images = [ vis.overlay_heatmap(disp_images[i], pred_heatmaps[i]) for i in range(len(pred_heatmaps)) ] nme_per_lm = calc_landmark_nme(lm_gt, lm_preds_max, ocular_norm=ocular_norm) lm_ssim_errs = calc_landmark_ssim_error(batch.images[:nimgs], X_recon[:nimgs], batch.landmarks[:nimgs]) # # Show input images # disp_images = vis.add_landmarks_to_images( disp_images, lm_gt[:nimgs], color=gt_color, draw_dots=True, draw_wireframe=False, 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=(0.0, 0.0, 1.0), draw_dots=True, draw_wireframe=False, gt_landmarks=lm_gt, draw_gt_offsets=True, landmarks_to_draw=landmarks_to_draw) # if lm_preds_cnn is not None: # disp_images = vis.add_landmarks_to_images(disp_images, lm_preds_cnn, color=(1, 1, 0), # gt_landmarks=lm_gt, draw_gt_offsets=False, # draw_wireframe=True, landmarks_to_draw=landmarks_to_draw) rows = [vis.make_grid(disp_images, nCols=nimgs, normalize=False)] # # Show reconstructions # X_recon_errs = 255.0 * torch.abs(batch.images - X_recon).reshape( len(batch.images), -1).mean(dim=1) 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))] lm_errs_max = calc_landmark_nme_per_img( lm_gt, lm_preds_max, ocular_norm=ocular_norm, landmarks_to_eval=lmcfg.LANDMARKS_NO_OUTLINE) lm_errs_max_outline = calc_landmark_nme_per_img( lm_gt, lm_preds_max, ocular_norm=ocular_norm, landmarks_to_eval=lmcfg.LANDMARKS_ONLY_OUTLINE) lm_errs_max_all = calc_landmark_nme_per_img( lm_gt, lm_preds_max, ocular_norm=ocular_norm, landmarks_to_eval=lmcfg.ALL_LANDMARKS) 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_preds_max[:nimgs], color=(0, 0, 1), landmarks_to_draw=landmarks_to_draw, draw_wireframe=False, lm_errs=nme_per_lm, # lm_confs=lm_confs, lm_confs=1 - 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_gt, color=gt_color, draw_wireframe=False, landmarks_to_draw=landmarks_to_draw) # landmarks from CNN prediction if lm_preds_cnn is not None: nme_per_lm = calc_landmark_nme(lm_gt, lm_preds_cnn, ocular_norm=ocular_norm) disp_X_recon = vis.add_landmarks_to_images( disp_X_recon, lm_preds_cnn, color=(1, 1, 0), landmarks_to_draw=lmcfg.ALL_LANDMARKS, draw_wireframe=False, lm_errs=nme_per_lm, gt_landmarks=lm_gt, draw_gt_offsets=True, draw_dots=True, offset_line_color=(1, 1, 1)) lm_errs_cnn = calc_landmark_nme_per_img( lm_gt, lm_preds_cnn, ocular_norm=ocular_norm, landmarks_to_eval=landmarks_to_draw) lm_errs_cnn_outline = calc_landmark_nme_per_img( lm_gt, lm_preds_cnn, ocular_norm=ocular_norm, landmarks_to_eval=lmcfg.LANDMARKS_ONLY_OUTLINE) lm_errs_cnn_all = calc_landmark_nme_per_img( lm_gt, lm_preds_cnn, ocular_norm=ocular_norm, landmarks_to_eval=lmcfg.ALL_LANDMARKS) disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_errs_cnn, loc='tr', format_string='{:>5.2f}', vmax=15) disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_errs_cnn_outline, loc='tr+1', format_string='{:>5.2f}', vmax=15) disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_errs_cnn_all, loc='tr+2', format_string='{:>5.2f}', vmax=15) # 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) 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: 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.append(vis.make_grid(disp_X_recon, nCols=nimgs)) if ssim_maps is not None: disp_ssim_maps = to_numpy( ds_utils.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)) X_gen_lm_hm = None X_gen_vis = None show_random_faces = False if show_random_faces: with torch.no_grad(): z_random = self.enc_rand(nimgs, self.saae.z_dim).cuda() outputs = self.saae.P(z_random) X_gen_vis = outputs[:, :3] if outputs.shape[1] > 3: X_gen_lm_hm = outputs[:, 3:] disp_X_gen = to_numpy( ds_utils.denormalized(X_gen_vis)[:nimgs].permute(0, 2, 3, 1)) if X_gen_lm_hm is not None: if lmcfg.LANDMARK_TARGET == 'colored': gen_heatmaps = [to_image(X_gen_lm_hm[i]) for i in range(nimgs)] elif lmcfg.LANDMARK_TARGET == 'multi_channel': X_gen_lm_hm = X_gen_lm_hm.max(dim=1)[0] gen_heatmaps = [to_image(X_gen_lm_hm[i]) for i in range(nimgs)] else: gen_heatmaps = [ to_image(X_gen_lm_hm[i, 0]) for i in range(nimgs) ] disp_X_gen = [ vis.overlay_heatmap(disp_X_gen[i], gen_heatmaps[i]) for i in range(len(pred_heatmaps)) ] # inputs = torch.cat([X_gen_vis, X_gen_lm_hm.detach()], dim=1) inputs = X_gen_lm_hm.detach() # disabled for multi_channel LM targets # lm_gen_preds = self.saae.lm_coords(inputs).reshape(len(inputs), -1, 2) # disp_X_gen = vis.add_landmarks_to_images(disp_X_gen, lm_gen_preds[:nimgs], color=(0,1,1)) disp_gen_heatmaps = [ vis.color_map(hm, vmin=0, vmax=1.0) for hm in gen_heatmaps ] img_gen_heatmaps = cv2.resize(vis.make_grid(disp_gen_heatmaps, nCols=nimgs, padval=0), None, fx=1.0, fy=1.0) cv2.imshow('generated landmarks', cv2.cvtColor(img_gen_heatmaps, cv2.COLOR_RGB2BGR)) rows.append(vis.make_grid(disp_X_gen, nCols=nimgs)) # self.saae.D.train(train_state_D) # self.saae.Q.train(train_state_Q) # self.saae.P.train(train_state_P) if horizontal: assert (nimgs == 1) disp_rows = vis.make_grid(rows, nCols=2) else: disp_rows = vis.make_grid(rows, nCols=1) 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)
for lm in landmarks: lm_x, lm_y = lm[0], lm[1] cv2.circle(img, (int(lm_x), int(lm_y)), 3, (0, 0, 255), -1) cv2.imshow('landmarks', cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) cv2.waitKey(0) if __name__ == '__main__': from utils import vis import torch torch.manual_seed(0) torch.cuda.manual_seed_all(0) ds = W300(train=True, deterministic=True, use_cache=True, test_split='challenging', daug=0, align_face_orientation=False, return_modified_images=False) dl = td.DataLoader(ds, batch_size=50, shuffle=False, num_workers=0) print(ds) cfg.WITH_LANDMARK_LOSS = False for data in dl: batch = Batch(data, gpu=False) inputs = batch.images.clone() imgs = vis._to_disp_images(inputs, denorm=True) imgs = vis.add_landmarks_to_images(imgs, batch.landmarks, radius=3, color=(0,255,0)) # imgs = vis.add_landmarks_to_images(imgs, data['landmarks_of'].numpy(), color=(1,0,0)) vis.vis_square(imgs, nCols=5, fx=1, fy=1, normalize=False)
def draw_results(X_resized, X_recon, levels_z=None, landmarks=None, landmarks_pred=None, cs_errs=None, ncols=15, fx=0.5, fy=0.5, additional_status_text=''): clean_images = True if clean_images: landmarks = None nimgs = len(X_resized) ncols = min(ncols, nimgs) img_size = X_recon.shape[-1] disp_X = vis.to_disp_images(X_resized, denorm=True) disp_X_recon = vis.to_disp_images(X_recon, denorm=True) # reconstruction error in pixels l1_dists = 255.0 * to_numpy( (X_resized - X_recon).abs().reshape(len(disp_X), -1).mean(dim=1)) # SSIM errors ssim = np.zeros(nimgs) for i in range(nimgs): ssim[i] = compare_ssim(disp_X[i], disp_X_recon[i], data_range=1.0, multichannel=True) landmarks = to_numpy(landmarks) cs_errs = to_numpy(cs_errs) text_size = img_size / 256 text_thickness = 2 # # Visualise resized input images and reconstructed images # if landmarks is not None: disp_X = vis.add_landmarks_to_images( disp_X, landmarks, draw_wireframe=False, landmarks_to_draw=lmcfg.LANDMARKS_19) disp_X_recon = vis.add_landmarks_to_images( disp_X_recon, landmarks, draw_wireframe=False, landmarks_to_draw=lmcfg.LANDMARKS_19) if landmarks_pred is not None: disp_X = vis.add_landmarks_to_images(disp_X, landmarks_pred, color=(1, 0, 0)) disp_X_recon = vis.add_landmarks_to_images(disp_X_recon, landmarks_pred, color=(1, 0, 0)) if not clean_images: disp_X_recon = vis.add_error_to_images(disp_X_recon, l1_dists, format_string='{:.1f}', size=text_size, thickness=text_thickness) 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, size=text_size, thickness=text_thickness) if cs_errs is not None: disp_X_recon = vis.add_error_to_images(disp_X_recon, cs_errs, loc='bl-2', format_string='{:>4.2f}', vmax=0.0, vmin=0.4, size=text_size, thickness=text_thickness) # landmark errors lm_errs = np.zeros(1) if landmarks is not None: try: from landmarks import lmutils lm_errs = lmutils.calc_landmark_nme_per_img( landmarks, landmarks_pred) disp_X_recon = vis.add_error_to_images(disp_X_recon, lm_errs, loc='br', format_string='{:>5.2f}', vmax=15, size=img_size / 256, thickness=2) except: pass img_input = vis.make_grid(disp_X, nCols=ncols, normalize=False) img_recon = vis.make_grid(disp_X_recon, nCols=ncols, normalize=False) img_input = cv2.resize(img_input, None, fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC) img_recon = cv2.resize(img_recon, None, fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC) img_stack = [img_input, img_recon] # # Visualise hidden layers # VIS_HIDDEN = False if VIS_HIDDEN: img_z = vis.draw_z_vecs(levels_z, size=(img_size, 30), ncols=ncols) img_z = cv2.resize(img_z, dsize=(img_input.shape[1], img_z.shape[0]), interpolation=cv2.INTER_NEAREST) img_stack.append(img_z) cs_errs_mean = np.mean(cs_errs) if cs_errs is not None else np.nan status_bar_text = ("l1 recon err: {:.2f}px " "ssim: {:.3f}({:.3f}) " "lms err: {:2} {}").format(l1_dists.mean(), cs_errs_mean, 1 - ssim.mean(), lm_errs.mean(), additional_status_text) img_status_bar = vis.draw_status_bar(status_bar_text, status_bar_width=img_input.shape[1], status_bar_height=30, dtype=img_input.dtype) img_stack.append(img_status_bar) return np.vstack(img_stack)
result['image_mod'] = crop_occ if self.return_landmark_heatmaps: result['lm_heatmaps'] = create_landmark_heatmaps(result['landmarks'], self.landmark_sigma, self.landmark_ids) return result if __name__ == '__main__': from utils.nn import Batch import utils.common utils.common.init_random() lmcfg.config_landmarks('aflw') ds = AFLW(train=True, deterministic=True, use_cache=True) dl = td.DataLoader(ds, batch_size=1, shuffle=False, num_workers=0) cfg.WITH_LANDMARK_LOSS = False for data in dl: batch = Batch(data, gpu=False) inputs = batch.images.clone() ds_utils.denormalize(inputs) imgs = vis.add_landmarks_to_images(inputs.numpy(), batch.landmarks.numpy(), radius=3, color=(0,255,0)) print(batch.fnames) vis.vis_square(imgs, nCols=1, fx=1.0, fy=1.0, normalize=False)