Exemple #1
0
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
Exemple #2
0
                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)
Exemple #4
0
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)
Exemple #5
0
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)
Exemple #6
0
        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)
Exemple #7
0
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)
Exemple #8
0
        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)
Exemple #10
0
            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)