def get_seg_map(self,
landmarks,
no_guassian=False,
size=256,
original_angles=None):
landmarks = landmarks[:, :, :2].cpu().numpy().astype(np.float)
all_heatmap = []
all_orig_heatmap = []
if original_angles is None:
original_angles = torch.zeros(landmarks.shape[0])
# key_points = []
for i in range(landmarks.shape[0]):
heatmap = curve.points_to_heatmap_68points(landmarks[i], 13, size,
self.opt.heatmap_size)
heatmap2 = curve.combine_map(heatmap, no_guassian=no_guassian)
if self.opt.isTrain:
if np.random.randint(2):
heatmap = np.zeros_like(heatmap)
else:
if torch.abs(original_angles[i]) < 0.255:
heatmap = np.zeros_like(heatmap)
all_heatmap.append(heatmap2)
all_orig_heatmap.append(heatmap)
# key_points.append(self.landmark_68_to_5(landmarks[i]))
all_heatmap = np.stack(all_heatmap, axis=0)
all_orig_heatmap = np.stack(all_orig_heatmap, axis=0)
all_heatmap = torch.from_numpy(all_heatmap.astype(np.float32)).cuda()
all_orig_heatmap = torch.from_numpy(all_orig_heatmap.astype(
np.float32)).cuda()
all_orig_heatmap = all_orig_heatmap.permute(0, 3, 1, 2)
all_orig_heatmap[all_orig_heatmap > 0] = 2.0
return all_heatmap, all_orig_heatmap
def get_seg_map(self, vertices, no_guassian=False, size=256):
landmarks = self.torch_get_68_points(vertices)
landmarks = landmarks[:, :, :2].cpu().numpy().astype(np.float)
all_heatmap = []
all_orig_heatmap = []
for i in range(landmarks.shape[0]):
heatmap = curve.points_to_heatmap_68points(landmarks[i], 13, size, self.opt.heatmap_size)
heatmap2 = curve.combine_map(heatmap, no_guassian=no_guassian)
all_heatmap.append(heatmap2)
all_orig_heatmap.append(heatmap)
all_heatmap = np.stack(all_heatmap, axis=0)
all_orig_heatmap = np.stack(all_orig_heatmap, axis=0)
all_heatmap = torch.from_numpy(all_heatmap.astype(np.float32)).cuda()
all_orig_heatmap = torch.from_numpy(all_orig_heatmap.astype(np.float32)).cuda()
all_orig_heatmap = all_orig_heatmap.permute(0, 3, 1, 2)
all_orig_heatmap[all_orig_heatmap > 0] = 1.0
return all_heatmap, all_orig_heatmap