Spaces:
Running
Running
File size: 6,156 Bytes
d015578 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import numpy as np
from scipy import interpolate
import cv2
class AddBoundary(object):
def __init__(self, num_landmarks=68, map_size=64, sigma=1, min_dpi=64):
self.num_landmarks = num_landmarks
self.sigma = sigma
if isinstance(map_size, (tuple, list)):
self.width = map_size[0]
self.height = map_size[1]
else:
self.width = map_size
self.height = map_size
if max(map_size) > min_dpi:
self.dpi = max(map_size)
else:
self.dpi = min_dpi
self.fig_size =[self.height/self.dpi, self.width/self.dpi]
def __call__(self, sample):
landmarks = sample['landmarks_float']
mask_lnd = sample['mask_ldm_float']
boundaries = self.get_dataset_boundaries(landmarks, mask_lnd)
functions = {}
for key, points in boundaries.items():
if len(points) != 0:
temp = points[0]
new_points = points[0:1, :]
for point in points[1:]:
if point[0] == temp[0] and point[1] == temp[1]:
continue
else:
new_points = np.concatenate((new_points, np.expand_dims(point, 0)), axis=0)
temp = point
points = new_points
if points.shape[0] == 1:
points = np.concatenate((points, points+0.001), axis=0)
k = min(4, points.shape[0])
functions[key] = interpolate.splprep([points[:, 0], points[:, 1]], k=k-1,s=0)
boundary_maps = np.zeros((len(boundaries), self.height, self.width))
for i_map, key in enumerate(functions.keys()):
boundary_map = np.zeros((self.height, self.width))
xnew = np.arange(0, 1, 1/self.dpi)
out = interpolate.splev(xnew, functions[key][0], der=0)
out = np.round(out).astype(int).transpose()
out = out[out[:, 0] < self.height]
out = out[out[:, 1] < self.width]
boundary_map[out[:,1], out[:,0]]= 255
# Smooth
sigma = self.sigma
temp = 255 - boundary_map.astype(np.uint8)
temp = cv2.distanceTransform(temp, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
temp = temp.astype(np.float32)
temp = np.where(temp < 3*sigma, np.exp(-(temp*temp)/(2*sigma*sigma)), 0 )
boundary_maps[i_map] = temp
sample['boundary'] = boundary_maps
return sample
def get_dataset_boundaries(self, landmarks, mask_lnd):
boundaries = {}
if self.num_landmarks == 68:
cheek = landmarks[0:17]
boundaries['cheek'] = cheek[mask_lnd[0:17] > 0]
left_eyebrow = landmarks[17:22]
boundaries['left_eyebrow'] = left_eyebrow[mask_lnd[17:22] > 0]
right_eyebrow = landmarks[22:27]
boundaries['right_eyebrow'] = right_eyebrow[mask_lnd[22:27] > 0]
nose = landmarks[27:31]
boundaries['nose'] = nose[mask_lnd[27:31] > 0]
nose_bot = landmarks[31:36]
boundaries['nose_bot'] = nose_bot[mask_lnd[31:36] > 0]
uper_left_eyelid = landmarks[36:40]
boundaries['upper_left_eyelid'] = uper_left_eyelid[mask_lnd[36:40] > 0]
lower_left_eyelid = np.array([landmarks[i] for i in [36, 41, 40, 39]])
lower_left_eyelid_mask = np.array([mask_lnd[i] for i in [36, 41, 40, 39]])
boundaries['lower_left_eyelid'] = lower_left_eyelid[lower_left_eyelid_mask > 0]
upper_right_eyelid = landmarks[42:46]
boundaries['upper_right_eyelid'] = upper_right_eyelid[mask_lnd[42:46] > 0]
lower_right_eyelid = np.array([landmarks[i] for i in [42, 47, 46, 45]])
lower_right_eyelid_mask = np.array([mask_lnd[i] for i in [42, 47, 46, 45]])
boundaries['lower_right_eyelid'] = lower_right_eyelid[lower_right_eyelid_mask > 0]
upper_outer_lip = landmarks[48:55]
boundaries['upper_outer_lip'] = upper_outer_lip[mask_lnd[48:55] > 0]
lower_outer_lip = np.array([landmarks[i] for i in [48, 59, 58, 57, 56, 55, 54]])
lower_outer_lip_mask = np.array([mask_lnd[i] for i in [48, 59, 58, 57, 56, 55, 54]])
boundaries['lower_outer_lip'] = lower_outer_lip[lower_outer_lip_mask > 0]
upper_inner_lip = np.array([landmarks[i] for i in [60, 61, 62, 63, 64]])
upper_inner_lip_mask = np.array([mask_lnd[i] for i in [60, 61, 62, 63, 64]])
boundaries['upper_inner_lip'] = upper_inner_lip[upper_inner_lip_mask > 0]
lower_inner_lip = np.array([landmarks[i] for i in [60, 67, 66, 65, 64]])
lower_inner_lip_mask = np.array([mask_lnd[i] for i in [60, 67, 66, 65, 64]])
boundaries['lower_inner_lip'] = lower_inner_lip[lower_inner_lip_mask > 0]
elif self.num_landmarks == 98:
boundaries['cheek'] = landmarks[0:33]
boundaries['upper_left_eyebrow'] = landmarks[33:38]
boundaries['lower_left_eyebrow'] = np.array([landmarks[i] for i in [33, 41, 40, 39, 38]])
boundaries['upper_right_eyebrow'] = landmarks[42:47]
boundaries['lower_right_eyebrow'] = landmarks[46:51]
boundaries['nose'] = landmarks[51:55]
boundaries['nose_bot'] = landmarks[55:60]
boundaries['upper_left_eyelid'] = landmarks[60:65]
boundaries['lower_left_eyelid'] = np.array([landmarks[i] for i in [60, 67, 66, 65, 64]])
boundaries['upper_right_eyelid'] = landmarks[68:73]
boundaries['lower_right_eyelid'] = np.array([landmarks[i] for i in [68, 75, 74, 73, 72]])
boundaries['upper_outer_lip'] = landmarks[76:83]
boundaries['lower_outer_lip'] = np.array([landmarks[i] for i in [76, 87, 86, 85, 84, 83, 82]])
boundaries['upper_inner_lip'] = np.array([landmarks[i] for i in [88, 89, 90, 91, 92]])
boundaries['lower_inner_lip'] = np.array([landmarks[i] for i in [88, 95, 94, 93, 92]])
return boundaries
|