def __process_affine(self, image, target, theta, nopoints, aux_info=None):
image, target, theta = image.clone(), target.copy(), theta.clone()
(C, H, W), (height, width) = image.size(), self.shape
if nopoints: # do not have label
norm_trans_points = torch.zeros((3, self.NUM_PTS))
heatmaps = torch.zeros(
(self.NUM_PTS + 1, height // self.downsample,
width // self.downsample))
mask = torch.ones((self.NUM_PTS + 1, 1, 1), dtype=torch.uint8)
else:
norm_trans_points = apply_affine2point(target.get_points(), theta,
(H, W))
norm_trans_points = apply_boundary(norm_trans_points)
real_trans_points = norm_trans_points.clone()
real_trans_points[:2, :] = denormalize_points(
self.shape, real_trans_points[:2, :])
heatmaps, mask = generate_label_map(real_trans_points.numpy(),
height // self.downsample,
width // self.downsample,
self.sigma, self.downsample,
nopoints,
self.heatmap_type) # H*W*C
heatmaps = torch.from_numpy(heatmaps.transpose(
(2, 0, 1))).type(torch.FloatTensor)
mask = torch.from_numpy(mask.transpose(
(2, 0, 1))).type(torch.ByteTensor)
affineImage = affine2image(image, theta, self.shape)
return affineImage, heatmaps, mask, norm_trans_points, theta
def __process_affine(self, image, target, theta, nopoints, aux_info=None):
image, target, theta = image.clone(), target.copy(), theta.clone()
(C, H, W), (height, width) = image.size(), self.shape
if nopoints: # do not have label
norm_trans_points = torch.zeros((3, self.NUM_PTS))
heatmaps = torch.zeros((self.NUM_PTS + 1, height // self.downsample, width // self.downsample))
mask = torch.ones((self.NUM_PTS + 1, 1, 1), dtype=torch.uint8)
transpose_theta = identity2affine(False)
else:
norm_trans_points = apply_affine2point(target.get_points(), theta, (H, W))
norm_trans_points = apply_boundary(norm_trans_points)
real_trans_points = norm_trans_points.clone()
real_trans_points[:2, :] = denormalize_points(self.shape, real_trans_points[:2, :])
heatmaps, mask = generate_label_map(real_trans_points.numpy(), height // self.downsample,
width // self.downsample, self.sigma, self.downsample, nopoints,
self.heatmap_type) # H*W*C
heatmaps = torch.from_numpy(heatmaps.transpose((2, 0, 1))).type(torch.FloatTensor)
mask = torch.from_numpy(mask.transpose((2, 0, 1))).type(torch.ByteTensor)
if self.mean_face is None:
# warnings.warn('In LandmarkDataset use identity2affine for transpose_theta because self.mean_face is None.')
transpose_theta = identity2affine(False)
else:
if torch.sum(norm_trans_points[2, :] == 1) < 3:
warnings.warn(
'In LandmarkDataset after transformation, no visiable point, using identity instead. Aux: {:}'.format(
aux_info))
transpose_theta = identity2affine(False)
else:
transpose_theta = solve2theta(norm_trans_points, self.mean_face.clone())
affineImage = affine2image(image, theta, self.shape)
if self.cutout is not None: affineImage = self.cutout(affineImage)
return affineImage, heatmaps, mask, norm_trans_points, theta, transpose_theta
def _process_(self, frames, target, view_info, index, skipopt):
# transform the image and points
frames, target, theta = self.transform(frames, target)
(C, H, W), (height, width) = frames[0].size(), self.shape
# obtain the visiable indicator vector
if target.is_none(): nopoints = True
else: nopoints = False
affineFrames, forward_flow, backward_flow, heatmaps, mask, norm_trans_points, THETA = self.__process_affine(
frames, target, theta, nopoints, skipopt)
torch_index = torch.IntTensor([index])
torch_nopoints = torch.ByteTensor([nopoints])
torch_shape = torch.IntTensor([H, W])
torch_is_3D = torch.ByteTensor([view_info is not None])
if view_info is None:
MultiViewKRT = torch.zeros((self.max_views, 3, 4))
MultiViewTensor = torch.zeros(
(self.max_views, 1 if self.use_gray else 3, self.shape[0],
self.shape[1]))
MultiViewShapes = torch.zeros((self.max_views, 2),
dtype=torch.float32)
MultiViewThetas = torch.zeros((self.max_views, 3, 3))
else:
MultiViewIs, MultiViewLs, MultiViewKRT = [], [], []
for Mimage, Mlabel, MKRT in view_info:
MultiViewIs.append(Mimage)
MultiViewLs.append(Mlabel)
MultiViewKRT.append(MKRT)
MultiViewIs, MultiViewLs, MultiViewTheta = self.transform(
MultiViewIs, MultiViewLs)
MultiViewTensor = []
for ximage, xtheta in zip(MultiViewIs, MultiViewTheta):
xTensor = affine2image(ximage, xtheta, self.shape)
MultiViewTensor.append(xTensor)
MultiViewKRT = torch.stack(MultiViewKRT)
MultiViewTensor = torch.stack(MultiViewTensor)
MultiViewShapes = torch.FloatTensor([
(image.size(-1), image.size(-2)) for image in MultiViewIs
]) # shape : W, H
MultiViewThetas = torch.stack(MultiViewTheta)
return affineFrames, forward_flow, backward_flow, heatmaps, mask, norm_trans_points, THETA, MultiViewTensor, MultiViewThetas, MultiViewShapes, MultiViewKRT, torch_is_3D, torch_index, torch_nopoints, torch_shape
def pro_debug_save(save_dir, name, image, heatmap, normpoint, meantheta,
predmap, recover):
name, ext = name.split('.')
save_dir.mkdir(parents=True, exist_ok=True)
C, H, W = image.size()
oriimage = recover(image)
oriimage.save(str(save_dir / '{:}-ori.{:}'.format(name, ext)))
if C == 1: color = (255, )
else: color = (102, 255, 102)
ptsimage = draw_image_by_points(oriimage, normpoint, 2, color, False,
False, True)
ptsimage.save(str(save_dir / '{:}-pts.{:}'.format(name, ext)))
meanI = affine2image(image, meantheta, (H, W))
meanimg = recover(meanI)
meanimg.save(str(save_dir / '{:}-tomean.{:}'.format(name, ext)))
_save_heatmap(oriimage, heatmap, save_dir, name, ext, 'GT')
_save_heatmap(oriimage, predmap, save_dir, name, ext, 'PD')
def __process_affine(self,
frames,
target,
theta,
nopoints,
skipopt,
aux_info=None):
frames, target, theta = [frame.clone() for frame in frames
], target.copy(), theta.clone()
(C, H, W), (height, width) = frames[0].size(), self.shape
if nopoints: # do not have label
norm_trans_points = torch.zeros((3, self.NUM_PTS))
heatmaps = torch.zeros(
(self.NUM_PTS + 1, height // self.downsample,
width // self.downsample))
mask = torch.ones((self.NUM_PTS + 1, 1, 1), dtype=torch.uint8)
else:
norm_trans_points = apply_affine2point(target.get_points(), theta,
(H, W))
norm_trans_points = apply_boundary(norm_trans_points)
real_trans_points = norm_trans_points.clone()
real_trans_points[:2, :] = denormalize_points(
self.shape, real_trans_points[:2, :])
heatmaps, mask = generate_label_map(real_trans_points.numpy(),
height // self.downsample,
width // self.downsample,
self.sigma, self.downsample,
nopoints,
self.heatmap_type) # H*W*C
heatmaps = torch.from_numpy(heatmaps.transpose(
(2, 0, 1))).type(torch.FloatTensor)
mask = torch.from_numpy(mask.transpose(
(2, 0, 1))).type(torch.ByteTensor)
affineFrames = [
affine2image(frame, theta, self.shape) for frame in frames
]
if not skipopt:
Gframes = [self.tensor2img(frame) for frame in affineFrames]
forward_flow, backward_flow = [], []
for idx in range(len(Gframes)):
if idx > 0:
forward_flow.append(
self.optflow.calc(Gframes[idx - 1], Gframes[idx],
None))
if idx + 1 < len(Gframes):
#backward_flow.append( self.optflow.calc(Gframes[idx], Gframes[idx+1], None) ) ## HDXY
backward_flow.append(
self.optflow.calc(Gframes[idx + 1], Gframes[idx],
None))
forward_flow = torch.stack(
[torch.from_numpy(x) for x in forward_flow])
backward_flow = torch.stack(
[torch.from_numpy(x) for x in backward_flow])
else:
forward_flow, backward_flow = torch.zeros(
(len(affineFrames) - 1, height, width, 2)), torch.zeros(
(len(affineFrames) - 1, height, width, 2))
# affineFrames #frames x #channel x #height x #width
# forward_flow (#frames-1) x #height x #width x 2
# backward_flow (#frames-1) x #height x #width x 2
return torch.stack(
affineFrames
), forward_flow, backward_flow, heatmaps, mask, norm_trans_points, theta
def __process_affine(self,
frames,
target,
theta,
nopoints,
skip_opt,
aux_info=None):
frames, target, theta = [frame.clone() for frame in frames
], target.copy(), theta.clone()
(C, H, W), (height, width) = frames[0].size(), self.shape
if nopoints: # do not have label
norm_trans_points = torch.zeros((3, self.NUM_PTS))
heatmaps = torch.zeros(
(self.NUM_PTS + 1, height // self.downsample,
width // self.downsample))
mask = torch.ones((self.NUM_PTS + 1, 1, 1), dtype=torch.uint8)
transpose_theta = identity2affine(False)
else:
norm_trans_points = apply_affine2point(target.get_points(), theta,
(H, W))
norm_trans_points = apply_boundary(norm_trans_points)
real_trans_points = norm_trans_points.clone()
real_trans_points[:2, :] = denormalize_points(
self.shape, real_trans_points[:2, :])
heatmaps, mask = generate_label_map(real_trans_points.numpy(),
height // self.downsample,
width // self.downsample,
self.sigma, self.downsample,
nopoints,
self.heatmap_type) # H*W*C
heatmaps = torch.from_numpy(heatmaps.transpose(
(2, 0, 1))).type(torch.FloatTensor)
mask = torch.from_numpy(mask.transpose(
(2, 0, 1))).type(torch.ByteTensor)
if torch.sum(norm_trans_points[2, :] ==
1) < 3 or self.mean_face is None:
warnings.warn(
'In GeneralDatasetV2 after transformation, no visiable point, using identity instead. Aux: {:}'
.format(aux_info))
transpose_theta = identity2affine(False)
else:
transpose_theta = solve2theta(norm_trans_points,
self.mean_face.clone())
affineFrames = [
affine2image(frame, theta, self.shape) for frame in frames
]
if not skip_opt:
Gframes = [self.tensor2img(frame) for frame in affineFrames]
forward_flow, backward_flow = [], []
for idx in range(len(Gframes)):
if idx > 0:
forward_flow.append(
self.optflow.calc(Gframes[idx - 1], Gframes[idx],
None))
if idx + 1 < len(Gframes):
#backward_flow.append( self.optflow.calc(Gframes[idx], Gframes[idx+1], None) )
backward_flow.append(
self.optflow.calc(Gframes[idx + 1], Gframes[idx],
None))
forward_flow = torch.stack(
[torch.from_numpy(x) for x in forward_flow])
backward_flow = torch.stack(
[torch.from_numpy(x) for x in backward_flow])
else:
forward_flow, backward_flow = torch.zeros(
(len(affineFrames) - 1, height, width, 2)), torch.zeros(
(len(affineFrames) - 1, height, width, 2))
# affineFrames #frames x #channel x #height x #width
# forward_flow (#frames-1) x #height x #width x 2
# backward_flow (#frames-1) x #height x #width x 2
return torch.stack(
affineFrames
), forward_flow, backward_flow, heatmaps, mask, norm_trans_points, theta, transpose_theta