Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import os | |
| import cv2 | |
| import matplotlib.pyplot as plt | |
| from scipy.ndimage import zoom | |
| def print_training_params_to_file(init_locals): | |
| """save param log file""" | |
| del init_locals['self'] | |
| with open(os.path.join(init_locals['save_log_path'], 'Training_Parameters.txt'), 'w') as f: | |
| f.write('Training Parameters:\n\n') | |
| for key, value in init_locals.items(): | |
| f.write('* %s: %s\n' % (key, value)) | |
| def heat_maps_to_landmarks(maps, image_size=256, num_landmarks=68): | |
| """find landmarks from heatmaps (arg max on each map)""" | |
| landmarks = np.zeros((num_landmarks,2)).astype('float32') | |
| for m_ind in range(num_landmarks): | |
| landmarks[m_ind, :] = np.unravel_index(maps[:, :, m_ind].argmax(), (image_size, image_size)) | |
| return landmarks | |
| def heat_maps_to_landmarks_alloc_once(maps, landmarks, image_size=256, num_landmarks=68): | |
| """find landmarks from heatmaps (arg max on each map) with pre-allocation""" | |
| for m_ind in range(num_landmarks): | |
| landmarks[m_ind, :] = np.unravel_index(maps[:, :, m_ind].argmax(), (image_size, image_size)) | |
| def batch_heat_maps_to_landmarks_alloc_once(batch_maps, batch_landmarks, batch_size, image_size=256, num_landmarks=68): | |
| """find landmarks from heatmaps (arg max on each map) - for multiple images""" | |
| for i in range(batch_size): | |
| heat_maps_to_landmarks_alloc_once( | |
| maps=batch_maps[i, :, :, :], landmarks=batch_landmarks[i, :, :], image_size=image_size, | |
| num_landmarks=num_landmarks) | |
| def normalize_map(map_in): | |
| map_min = map_in.min() | |
| return (map_in - map_min) / (map_in.max() - map_min) | |
| def map_to_rgb(map_gray): | |
| cmap = plt.get_cmap('jet') | |
| rgba_map_image = cmap(map_gray) | |
| map_rgb = np.delete(rgba_map_image, 3, 2) * 255 | |
| return map_rgb | |
| def create_img_with_landmarks(image, landmarks, image_size=256, num_landmarks=68, scale=255, circle_size=2): | |
| """add landmarks to a face image""" | |
| image = image.reshape(image_size, image_size, -1) | |
| if scale is 0: | |
| image = 127.5 * (image + 1) | |
| elif scale is 1: | |
| image *= 255 | |
| landmarks = landmarks.reshape(num_landmarks, 2) | |
| landmarks = np.clip(landmarks, 0, image_size-1) | |
| for (y, x) in landmarks.astype('int'): | |
| cv2.circle(image, (x, y), circle_size, (255, 0, 0), -1) | |
| return image | |
| def heat_maps_to_image(maps, landmarks=None, image_size=256, num_landmarks=68): | |
| """create one image from multiple heatmaps""" | |
| if landmarks is None: | |
| landmarks = heat_maps_to_landmarks(maps, image_size=image_size, num_landmarks=num_landmarks) | |
| x, y = np.mgrid[0:image_size, 0:image_size] | |
| pixel_dist = np.sqrt( | |
| np.square(np.expand_dims(x, 2) - landmarks[:, 0]) + np.square(np.expand_dims(y, 2) - landmarks[:, 1])) | |
| nn_landmark = np.argmin(pixel_dist, 2) | |
| map_image = maps[x, y, nn_landmark] | |
| map_image = (map_image-map_image.min())/(map_image.max()-map_image.min()) # normalize for visualization | |
| return map_image | |
| def merge_images_landmarks_maps_gt(images, maps, maps_gt, landmarks=None, image_size=256, num_landmarks=68, | |
| num_samples=9, scale=255, circle_size=2, fast=False): | |
| """create image for log - containing input face images, predicted heatmaps and GT heatmaps (if exists)""" | |
| images = images[:num_samples] | |
| if maps.shape[1] is not image_size: | |
| images = zoom(images, (1, 0.25, 0.25, 1)) | |
| image_size /= 4 | |
| image_size=int(image_size) | |
| if maps_gt is not None: | |
| if maps_gt.shape[1] is not image_size: | |
| maps_gt = zoom(maps_gt, (1, 0.25, 0.25, 1)) | |
| cmap = plt.get_cmap('jet') | |
| row = int(np.sqrt(num_samples)) | |
| if maps_gt is None: | |
| merged = np.zeros([row * image_size, row * image_size * 2, 3]) | |
| else: | |
| merged = np.zeros([row * image_size, row * image_size * 3, 3]) | |
| for idx, img in enumerate(images): | |
| i = idx // row | |
| j = idx % row | |
| if landmarks is None: | |
| img_landmarks = heat_maps_to_landmarks(maps[idx, :, :, :], image_size=image_size, | |
| num_landmarks=num_landmarks) | |
| else: | |
| img_landmarks = landmarks[idx] | |
| if fast: | |
| map_image = np.amax(maps[idx, :, :, :], 2) | |
| map_image = (map_image - map_image.min()) / (map_image.max() - map_image.min()) | |
| else: | |
| map_image = heat_maps_to_image(maps[idx, :, :, :], img_landmarks, image_size=image_size, | |
| num_landmarks=num_landmarks) | |
| rgba_map_image = cmap(map_image) | |
| map_image = np.delete(rgba_map_image, 3, 2) * 255 | |
| img = create_img_with_landmarks(img, img_landmarks, image_size, num_landmarks, scale=scale, | |
| circle_size=circle_size) | |
| if maps_gt is not None: | |
| if fast: | |
| map_gt_image = np.amax(maps_gt[idx, :, :, :], 2) | |
| map_gt_image = (map_gt_image - map_gt_image.min()) / (map_gt_image.max() - map_gt_image.min()) | |
| else: | |
| map_gt_image = heat_maps_to_image(maps_gt[idx, :, :, :], image_size=image_size, | |
| num_landmarks=num_landmarks) | |
| rgba_map_gt_image = cmap(map_gt_image) | |
| map_gt_image = np.delete(rgba_map_gt_image, 3, 2) * 255 | |
| merged[i * image_size:(i + 1) * image_size, (j * 3) * image_size:(j * 3 + 1) * image_size, :] = img | |
| merged[i * image_size:(i + 1) * image_size, (j * 3 + 1) * image_size:(j * 3 + 2) * image_size, | |
| :] = map_image | |
| merged[i * image_size:(i + 1) * image_size, (j * 3 + 2) * image_size:(j * 3 + 3) * image_size, | |
| :] = map_gt_image | |
| else: | |
| merged[i * image_size:(i + 1) * image_size, (j * 2) * image_size:(j * 2 + 1) * image_size, :] = img | |
| merged[i * image_size:(i + 1) * image_size, (j * 2 + 1) * image_size:(j * 2 + 2) * image_size,:] = map_image | |
| return merged | |
| def map_comapre_channels(images, maps1, maps2, image_size=64, num_landmarks=68, scale=255): | |
| """create image for log - present one face image, along with all its heatmaps (one for each landmark)""" | |
| map1 = maps1[0] | |
| if maps2 is not None: | |
| map2 = maps2[0] | |
| image = images[0] | |
| if image.shape[0] is not image_size: | |
| image = zoom(image, (0.25, 0.25, 1)) | |
| if scale is 1: | |
| image *= 255 | |
| elif scale is 0: | |
| image = 127.5 * (image + 1) | |
| row = np.ceil(np.sqrt(num_landmarks)).astype(np.int64) | |
| if maps2 is not None: | |
| merged = np.zeros([row * image_size, row * image_size * 2, 3]) | |
| else: | |
| merged = np.zeros([row * image_size, row * image_size, 3]) | |
| for idx in range(num_landmarks): | |
| i = idx // row | |
| j = idx % row | |
| channel_map = map_to_rgb(normalize_map(map1[:, :, idx])) | |
| if maps2 is not None: | |
| channel_map2 = map_to_rgb(normalize_map(map2[:, :, idx])) | |
| merged[i * image_size:(i + 1) * image_size, (j * 2) * image_size:(j * 2 + 1) * image_size, :] =\ | |
| channel_map | |
| merged[i * image_size:(i + 1) * image_size, (j * 2 + 1) * image_size:(j * 2 + 2) * image_size, :] =\ | |
| channel_map2 | |
| else: | |
| merged[i * image_size:(i + 1) * image_size, j * image_size:(j + 1) * image_size, :] = channel_map | |
| i = (idx + 1) // row | |
| j = (idx + 1) % row | |
| if maps2 is not None: | |
| merged[i * image_size:(i + 1) * image_size, (j * 2) * image_size:(j * 2 + 1) * image_size, :] = image | |
| else: | |
| merged[i * image_size:(i + 1) * image_size, j * image_size:(j + 1) * image_size, :] = image | |
| return merged | |