def calc_landmark_ssim_score(X,
X_recon,
lms,
wnd_size=int(cfg.INPUT_SCALE_FACTOR * 16)):
input_images = vis._to_disp_images(X, denorm=True)
recon_images = vis._to_disp_images(X_recon, denorm=True)
data_range = 255.0 if input_images[0].dtype == np.uint8 else 1.0
nimgs = len(input_images)
nlms = len(lms[0])
scores = np.zeros((nimgs, nlms), dtype=np.float32)
for i in range(nimgs):
S = compare_ssim(input_images[i],
recon_images[i],
data_range=data_range,
multichannel=True,
full=True)[1]
S = S.mean(axis=2)
for lid in range(nlms):
x = int(lms[i, lid, 0])
y = int(lms[i, lid, 1])
t = max(0, y - wnd_size // 2)
b = min(S.shape[0] - 1, y + wnd_size // 2)
l = max(0, x - wnd_size // 2)
r = min(S.shape[1] - 1, x + wnd_size // 2)
wnd = S[t:b, l:r]
scores[i, lid] = wnd.mean()
return np.nan_to_num(scores)
def show_pairs(images, features, pairs):
dists = np.sqrt(np.sum((features[0] - features[1])**2, axis=1))
ds_utils.denormalize(images[0])
ds_utils.denormalize(images[1])
images[1] = vis.add_error_to_images(images[1],
dists,
size=2.0,
thickness=2,
vmin=0,
vmax=1)
images[1] = vis.add_id_to_images(images[1],
pairs.numpy(),
size=1.2,
thickness=2,
color=(1, 0, 1))
thresh = 0.4
corrects = (dists < thresh) == pairs.cpu().numpy()
colors = [(0, 1, 0) if c else (1, 0, 0) for c in corrects]
images[1] = vis.add_cirle_to_images(images[1], colors)
images[0] = vis._to_disp_images(images[0])
img_rows = [
vis.make_grid(imgs,
fx=0.75,
fy=0.75,
nCols=len(dists),
normalize=False) for imgs in images
]
vis.vis_square(img_rows, nCols=1, normalize=False)
def calc_landmark_ncc(X, X_recon, lms):
input_images = vis._to_disp_images(X, denorm=True)
recon_images = vis._to_disp_images(X_recon, denorm=True)
nimgs = len(input_images)
nlms = len(lms[0])
wnd_size = int(cfg.INPUT_SCALE_FACTOR * 15)
nccs = np.zeros((nimgs, nlms), dtype=np.float32)
img_shape = input_images[0].shape
for i in range(nimgs):
for lid in range(nlms):
x = int(lms[i, lid, 0])
y = int(lms[i, lid, 1])
t = max(0, y - wnd_size // 2)
b = min(img_shape[0] - 1, y + wnd_size // 2)
l = max(0, x - wnd_size // 2)
r = min(img_shape[1] - 1, x + wnd_size // 2)
wnd1 = input_images[i][t:b, l:r]
wnd2 = recon_images[i][t:b, l:r]
ncc = ((wnd1 - wnd1.mean()) *
(wnd2 - wnd2.mean())).mean() / (wnd1.std() * wnd2.std())
nccs[i, lid] = ncc
return np.clip(np.nan_to_num(nccs), a_min=-1, a_max=1)
def calc_landmark_ssim_error(X, X_recon, lms):
input_images = vis._to_disp_images(X, denorm=True)
recon_images = vis._to_disp_images(X_recon, denorm=True)
nimgs = len(input_images)
nlms = len(lms[0])
wnd_size = int(cfg.INPUT_SCALE_FACTOR * 11)
errs = np.zeros((nimgs, nlms), dtype=np.float32)
for i in range(nimgs):
S = compare_ssim(input_images[i],
recon_images[i],
data_range=1.0,
multichannel=True,
full=True)[1]
S = S.mean(axis=2)
for lid in range(nlms):
x = int(lms[i, lid, 0])
y = int(lms[i, lid, 1])
t = max(0, y - wnd_size // 2)
b = min(S.shape[0] - 1, y + wnd_size // 2)
l = max(0, x - wnd_size // 2)
r = min(S.shape[1] - 1, x + wnd_size // 2)
wnd = S[t:b, l:r]
errs[i, lid] = 1 - wnd.mean()
return errs
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)
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)
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)