| import cv2
|
| import numpy as np
|
| from PIL import Image
|
| import matplotlib.pyplot as plt
|
|
|
|
|
|
|
| def log_input_image(x, opts):
|
| return tensor2im(x)
|
|
|
|
|
| def tensor2im(var):
|
| var = var.cpu().detach().transpose(0, 2).transpose(0, 1).numpy()
|
| var = ((var + 1) / 2)
|
| var[var < 0] = 0
|
| var[var > 1] = 1
|
| var = var * 255
|
| return Image.fromarray(var.astype('uint8'))
|
|
|
|
|
|
|
| def get_colors():
|
|
|
| colors = [[0, 0, 0], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255],
|
| [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204],
|
| [255, 51, 153], [0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]]
|
| return colors
|
|
|
|
|
| def vis_faces(log_hooks):
|
| display_count = len(log_hooks)
|
| fig = plt.figure(figsize=(8, 4 * display_count))
|
| gs = fig.add_gridspec(display_count, 3)
|
| for i in range(display_count):
|
| hooks_dict = log_hooks[i]
|
| fig.add_subplot(gs[i, 0])
|
| vis_faces_no_id(hooks_dict, fig, gs, i)
|
| plt.tight_layout()
|
| return fig
|
|
|
|
|
| def vis_faces_no_id(hooks_dict, fig, gs, i):
|
| plt.imshow(hooks_dict['input_face'], cmap="gray")
|
| plt.title('Input')
|
| fig.add_subplot(gs[i, 1])
|
| plt.imshow(hooks_dict['y_hat'])
|
| plt.title('Output (TriPlaneNet)')
|
| fig.add_subplot(gs[i, 2])
|
| plt.imshow(hooks_dict['y_hat_psp'])
|
| plt.title('Output (pSp)')
|
|
|
