def align_shapes(im, reference_shape, init=True, bb_hat=None):
reference_shape = PointCloud(reference_shape)
if init:
bb = im.landmarks['bb'].lms.bounding_box()
im.landmarks['__initial'] = align_shape_with_bounding_box(
reference_shape, bb)
im = im.rescale_to_pointcloud(reference_shape, group='__initial')
lms = im.landmarks['PTS'].lms
init = im.landmarks['__initial'].lms
bb_hat = im.landmarks['bb'].lms
# im = im.resize((235,200))
pixels = grey_to_rgb(im).pixels.transpose(1, 2, 0).copy()
height, width = pixels.shape[:2]
padded_image = np.random.rand(395, 467, 3).astype(np.float32)
dy = max(int((395 - height - 1) / 2), 0)
dx = max(int((467 - width - 1) / 2), 0)
pts = lms.points
pts[:, 0] += dy
pts[:, 1] += dx
init_pts = init.points
init_pts[:, 0] += dy
init_pts[:, 1] += dx
bb_pts = bb_hat.points
bb_pts[:, 0] += dy
bb_pts[:, 1] += dx
lms = lms.from_vector(pts)
init = init.from_vector(init_pts)
bb_hat = bb_hat.from_vector(bb_pts)
padded_image[dy:(height + dy), dx:(width + dx), :] = pixels
gt = lms.points.astype(np.float32)
init = init.points.astype(np.float32)
return np.expand_dims(padded_image,
0), np.expand_dims(init, 0), np.expand_dims(
gt, 0), bb_hat.bounding_box()
else:
bb = bb_hat
# print(bb.points)
im.landmarks['a'] = align_shape_with_bounding_box(reference_shape, bb)
init = im.landmarks['a'].lms
init = init.points.astype(np.float32)
# print(PointCloud(init).bounding_box().points)
return np.expand_dims(init, 0)
def load_image_test(path, reference_shape, frame_num):
file_name = path[:-1] + "/%06d.jpg" % (frame_num)
im = mio.import_image(file_name)
im.landmarks['PTS'] = mio.import_landmark_file(path[:-1] +
"/annot/%06d.pts" %
(frame_num))
# im.landmarks['PTS'] = mio.import_landmark_file(path[:-1] + "/%06d.pts" % (frame_num))
bb_path = path[:-1] + "/bbs/%06d.pts" % (frame_num)
im.landmarks['bb'] = mio.import_landmark_file(bb_path)
im = im.crop_to_landmarks_proportion(0.3, group='bb')
reference_shape = PointCloud(reference_shape)
bb = im.landmarks['bb'].lms.bounding_box()
im.landmarks['__initial'] = align_shape_with_bounding_box(
reference_shape, bb)
im = im.rescale_to_pointcloud(reference_shape, group='__initial')
lms = im.landmarks['PTS'].lms
initial = im.landmarks['__initial'].lms
# if the image is greyscale then convert to rgb.
pixels = grey_to_rgb(im).pixels.transpose(1, 2, 0)
gt_truth = lms.points.astype(np.float32)
estimate = initial.points.astype(np.float32)
return 1, pixels.astype(np.float32).copy(), gt_truth, estimate
def load_image(path,
reference_shape,
is_training=False,
group='PTS',
mirror_image=False):
"""Load an annotated image.
In the directory of the provided image file, there
should exist a landmark file (.pts) with the same
basename as the image file.
Args:
path: a path containing an image file.
reference_shape: a numpy array [num_landmarks, 2]
is_training: whether in training mode or not.
group: landmark group containing the grounth truth landmarks.
mirror_image: flips horizontally the image's pixels and landmarks.
Returns:
pixels: a numpy array [width, height, 3].
estimate: an initial estimate a numpy array [68, 2].
gt_truth: the ground truth landmarks, a numpy array [68, 2].
"""
im = mio.import_image(path)
bb_root = im.path.parent.relative_to(im.path.parent.parent.parent)
if 'set' not in str(bb_root):
bb_root = im.path.parent.relative_to(im.path.parent.parent)
im.landmarks['bb'] = mio.import_landmark_file(
str(Path('bbs') / bb_root / (im.path.stem + '.pts')))
im = im.crop_to_landmarks_proportion(0.3, group='bb')
reference_shape = PointCloud(reference_shape)
bb = im.landmarks['bb'].lms.bounding_box()
im.landmarks['__initial'] = align_shape_with_bounding_box(
reference_shape, bb)
im = im.rescale_to_pointcloud(reference_shape, group='__initial')
if mirror_image:
im = utils.mirror_image(im)
lms = im.landmarks[group].lms
initial = im.landmarks['__initial'].lms
# if the image is greyscale then convert to rgb.
pixels = grey_to_rgb(im).pixels.transpose(1, 2, 0)
gt_truth = lms.points.astype(np.float32)
estimate = initial.points.astype(np.float32)
return pixels.astype(np.float32).copy(), gt_truth, estimate
def load_image(path, reference_shape, is_training=False, group='PTS',
mirror_image=False):
"""Load an annotated image.
In the directory of the provided image file, there
should exist a landmark file (.pts) with the same
basename as the image file.
Args:
path: a path containing an image file.
reference_shape: a numpy array [num_landmarks, 2]
is_training: whether in training mode or not.
group: landmark group containing the grounth truth landmarks.
mirror_image: flips horizontally the image's pixels and landmarks.
Returns:
pixels: a numpy array [width, height, 3].
estimate: an initial estimate a numpy array [68, 2].
gt_truth: the ground truth landmarks, a numpy array [68, 2].
"""
im = mio.import_image(path)
bb_root = im.path.parent.relative_to(im.path.parent.parent.parent)
if 'set' not in str(bb_root):
bb_root = im.path.parent.relative_to(im.path.parent.parent)
im.landmarks['bb'] = mio.import_landmark_file(str(Path('bbs') / bb_root / (
im.path.stem + '.pts')))
im = im.crop_to_landmarks_proportion(0.3, group='bb')
reference_shape = PointCloud(reference_shape)
bb = im.landmarks['bb'].lms.bounding_box()
im.landmarks['__initial'] = align_shape_with_bounding_box(reference_shape,
bb)
im = im.rescale_to_pointcloud(reference_shape, group='__initial')
if mirror_image:
im = utils.mirror_image(im)
lms = im.landmarks[group].lms
initial = im.landmarks['__initial'].lms
# if the image is greyscale then convert to rgb.
pixels = grey_to_rgb(im).pixels.transpose(1, 2, 0)
gt_truth = lms.points.astype(np.float32)
estimate = initial.points.astype(np.float32)
return pixels.astype(np.float32).copy(), gt_truth, estimate
def get_noisy_init_from_bb(reference_shape, bb, noise_percentage=.02):
"""Roughly aligns a reference shape to a bounding box.
This adds some uniform noise for translation and scale to the
aligned shape.
Args:
reference_shape: a numpy array [num_landmarks, 2]
bb: bounding box, a numpy array [4, ]
noise_percentage: noise presentation to add.
Returns:
The aligned shape, as a numpy array [num_landmarks, 2]
"""
bb = PointCloud(bb)
reference_shape = PointCloud(reference_shape)
bb = noisy_shape_from_bounding_box(
reference_shape,
bb,
noise_percentage=[noise_percentage, 0,
noise_percentage]).bounding_box()
return align_shape_with_bounding_box(reference_shape, bb).points
def get_noisy_init_from_bb(reference_shape, bb, noise_percentage=.02):
"""Roughly aligns a reference shape to a bounding box.
This adds some uniform noise for translation and scale to the
aligned shape.
Args:
reference_shape: a numpy array [num_landmarks, 2]
bb: bounding box, a numpy array [4, ]
noise_percentage: noise presentation to add.
Returns:
The aligned shape, as a numpy array [num_landmarks, 2]
"""
bb = PointCloud(bb)
reference_shape = PointCloud(reference_shape)
bb = noisy_shape_from_bounding_box(
reference_shape, bb,
noise_percentage=[noise_percentage, 0, noise_percentage]
).bounding_box()
return align_shape_with_bounding_box(reference_shape, bb).points
def _train_batch(self, image_batch, increment=False, group=None,
bounding_box_group_glob=None, verbose=False):
# Rescale images wrt the scale factor between the existing
# reference_shape and their ground truth (group) shapes
image_batch = rescale_images_to_reference_shape(
image_batch, group, self.reference_shape,
verbose=verbose)
# Create a callable that generates perturbations of the bounding boxes
# of the provided images.
generated_bb_func = generate_perturbations_from_gt(
image_batch, self.n_perturbations,
self._perturb_from_gt_bounding_box, gt_group=group,
bb_group_glob=bounding_box_group_glob, verbose=verbose)
# For each scale (low --> high)
for j in range(self.n_scales):
# Print progress if asked
if verbose:
if len(self.scales) > 1:
scale_prefix = ' - Scale {}: '.format(j)
else:
scale_prefix = ' - '
else:
scale_prefix = None
# Extract features. Features are extracted only if we are at the
# first scale or if the features of the current scale are different
# than the ones extracted at the previous scale.
if j == 0 and self.holistic_features[j] == no_op:
# Saves a lot of memory
feature_images = image_batch
elif (j == 0 or
self.holistic_features[j] != self.holistic_features[j - 1]):
# Compute features only if this is the first pass through
# the loop or the features at this scale are different from
# the features at the previous scale
feature_images = compute_features(image_batch,
self.holistic_features[j],
prefix=scale_prefix,
verbose=verbose)
# Rescale images according to scales. Note that scale_images is smart
# enough in order not to rescale the images if the current scale
# factor equals to 1.
scaled_images, scale_transforms = scale_images(
feature_images, self.scales[j], prefix=scale_prefix,
return_transforms=True, verbose=verbose)
# Extract scaled ground truth shapes for current scale
scaled_shapes = [i.landmarks[group] for i in scaled_images]
# Get shape estimations of current scale. If we are at the first
# scale, this is done by aligning the reference shape with the
# perturbed bounding boxes. If we are at the rest of the scales,
# then the current shapes are attached on the scaled_images with
# key '__sdm_current_shape_{}'.
current_shapes = []
if j == 0:
# At the first scale, the current shapes are created by aligning
# the reference shape to the perturbed bounding boxes.
msg = '{}Aligning reference shape with bounding boxes.'.format(
scale_prefix)
wrap = partial(print_progress, prefix=msg,
end_with_newline=False, verbose=verbose)
# Extract perturbations at the very bottom level
for ii in wrap(scaled_images):
c_shapes = []
for bbox in generated_bb_func(ii):
c_s = align_shape_with_bounding_box(
self.reference_shape, bbox)
c_shapes.append(c_s)
current_shapes.append(c_shapes)
else:
# At the rest of the scales, extract the current shapes that
# were attached to the images
msg = '{}Extracting shape estimations from previous ' \
'scale.'.format(scale_prefix)
wrap = partial(print_progress, prefix=msg,
end_with_newline=False, verbose=verbose)
for ii in wrap(scaled_images):
c_shapes = []
for k in list(range(self.n_perturbations)):
c_key = '__sdm_current_shape_{}'.format(k)
c_shapes.append(ii.landmarks[c_key])
current_shapes.append(c_shapes)
# Train supervised descent algorithm. This returns the shape
# estimations for the next scale.
if not increment:
current_shapes = self.algorithms[j].train(
scaled_images, scaled_shapes, current_shapes,
prefix=scale_prefix, verbose=verbose)
else:
current_shapes = self.algorithms[j].increment(
scaled_images, scaled_shapes, current_shapes,
prefix=scale_prefix, verbose=verbose)
# Scale the current shape estimations for the next level. This
# doesn't have to be done for the last scale. The only thing we need
# to do at the last scale is to remove any attached landmarks from
# the training images.
if j < (self.n_scales - 1):
if self.holistic_features[j + 1] != self.holistic_features[j]:
# Features will be extracted, thus attach current_shapes on
# the training images (image_batch)
for jj, image_shapes in enumerate(current_shapes):
for k, shape in enumerate(image_shapes):
c_key = '__sdm_current_shape_{}'.format(k)
image_batch[jj].landmarks[c_key] = \
scale_transforms[jj].apply(shape)
else:
# Features won't be extracted;. the same feature_images will
# be used for the next scale, thus attach current_shapes on
# them.
for jj, image_shapes in enumerate(current_shapes):
for k, shape in enumerate(image_shapes):
c_key = '__sdm_current_shape_{}'.format(k)
feature_images[jj].landmarks[c_key] = \
scale_transforms[jj].apply(shape)
else:
# Check if original training image (image_batch) got some current
# shape estimations attached. If yes, delete them.
if '__sdm_current_shape_0' in image_batch[0].landmarks:
for image in image_batch:
for k in list(range(self.n_perturbations)):
c_key = '__sdm_current_shape_{}'.format(k)
del image.landmarks[c_key]
def _train_batch(self,
image_batch,
increment=False,
group=None,
bounding_box_group_glob=None,
verbose=False):
# Rescale to existing reference shape
image_batch = rescale_images_to_reference_shape(image_batch,
group,
self.reference_shape,
verbose=verbose)
generated_bb_func = generate_perturbations_from_gt(
image_batch,
self.n_perturbations,
self._perturb_from_gt_bounding_box,
gt_group=group,
bb_group_glob=bounding_box_group_glob,
verbose=verbose)
# for each scale (low --> high)
current_shapes = []
for j in range(self.n_scales):
if verbose:
if len(self.scales) > 1:
scale_prefix = ' - Scale {}: '.format(j)
else:
scale_prefix = ' - '
else:
scale_prefix = None
# Handle holistic features
if j == 0 and self.holistic_features[j] == no_op:
# Saves a lot of memory
feature_images = image_batch
elif j == 0 or self.holistic_features[
j] is not self.holistic_features[j - 1]:
# Compute features only if this is the first pass through
# the loop or the features at this scale are different from
# the features at the previous scale
feature_images = compute_features(image_batch,
self.holistic_features[j],
prefix=scale_prefix,
verbose=verbose)
# handle scales
if self.scales[j] != 1:
# Scale feature images only if scale is different than 1
scaled_images = scale_images(feature_images,
self.scales[j],
prefix=scale_prefix,
verbose=verbose)
else:
scaled_images = feature_images
# Extract scaled ground truth shapes for current scale
scaled_shapes = [i.landmarks[group].lms for i in scaled_images]
if j == 0:
msg = '{}Aligning reference shape with bounding boxes.'.format(
scale_prefix)
wrap = partial(print_progress,
prefix=msg,
end_with_newline=False,
verbose=verbose)
# Extract perturbations at the very bottom level
for ii in wrap(scaled_images):
c_shapes = []
for bbox in generated_bb_func(ii):
c_s = align_shape_with_bounding_box(
self.reference_shape, bbox)
c_shapes.append(c_s)
current_shapes.append(c_shapes)
# train supervised descent algorithm
if not increment:
current_shapes = self.algorithms[j].train(scaled_images,
scaled_shapes,
current_shapes,
prefix=scale_prefix,
verbose=verbose)
else:
current_shapes = self.algorithms[j].increment(
scaled_images,
scaled_shapes,
current_shapes,
prefix=scale_prefix,
verbose=verbose)
# Scale current shapes to next resolution, don't bother
# scaling final level
if j != (self.n_scales - 1):
transform = Scale(self.scales[j + 1] / self.scales[j],
n_dims=2)
for image_shapes in current_shapes:
for k, shape in enumerate(image_shapes):
image_shapes[k] = transform.apply(shape)
def _train_batch(self, image_batch, increment=False, group=None,
bounding_box_group_glob=None, verbose=False):
# Rescale to existing reference shape
image_batch = rescale_images_to_reference_shape(
image_batch, group, self.reference_shape,
verbose=verbose)
generated_bb_func = generate_perturbations_from_gt(
image_batch, self.n_perturbations,
self._perturb_from_gt_bounding_box, gt_group=group,
bb_group_glob=bounding_box_group_glob, verbose=verbose)
# for each scale (low --> high)
current_shapes = []
for j in range(self.n_scales):
if verbose:
if len(self.scales) > 1:
scale_prefix = ' - Scale {}: '.format(j)
else:
scale_prefix = ' - '
else:
scale_prefix = None
# Handle holistic features
if j == 0 and self.holistic_features[j] == no_op:
# Saves a lot of memory
feature_images = image_batch
elif j == 0 or self.holistic_features[j] is not self.holistic_features[j - 1]:
# Compute features only if this is the first pass through
# the loop or the features at this scale are different from
# the features at the previous scale
feature_images = compute_features(image_batch,
self.holistic_features[j],
prefix=scale_prefix,
verbose=verbose)
# handle scales
if self.scales[j] != 1:
# Scale feature images only if scale is different than 1
scaled_images = scale_images(feature_images, self.scales[j],
prefix=scale_prefix,
verbose=verbose)
else:
scaled_images = feature_images
# Extract scaled ground truth shapes for current scale
scaled_shapes = [i.landmarks[group].lms for i in scaled_images]
if j == 0:
msg = '{}Aligning reference shape with bounding boxes.'.format(
scale_prefix)
wrap = partial(print_progress, prefix=msg,
end_with_newline=False, verbose=verbose)
# Extract perturbations at the very bottom level
for ii in wrap(scaled_images):
c_shapes = []
for bbox in generated_bb_func(ii):
c_s = align_shape_with_bounding_box(
self.reference_shape, bbox)
c_shapes.append(c_s)
current_shapes.append(c_shapes)
# train supervised descent algorithm
if not increment:
current_shapes = self.algorithms[j].train(
scaled_images, scaled_shapes, current_shapes,
prefix=scale_prefix, verbose=verbose)
else:
current_shapes = self.algorithms[j].increment(
scaled_images, scaled_shapes, current_shapes,
prefix=scale_prefix, verbose=verbose)
# Scale current shapes to next resolution, don't bother
# scaling final level
if j != (self.n_scales - 1):
transform = Scale(self.scales[j + 1] / self.scales[j],
n_dims=2)
for image_shapes in current_shapes:
for shape in image_shapes:
transform.apply_inplace(shape)
def prepare_images(paths, num_patches=73, verbose=True):
"""Save Train Images to TFRecord
Args:
paths: a list of strings containing the data directories.
num_patches: number of landmarks
verbose: boolean, print debugging info.
Returns:
None
"""
if len(paths) == 0:
return
# .../<Dataset>/Images/*.png -> .../<Dataset>
path_base = Path(paths[0]).parent.parent
image_paths = []
# First: get all image paths
for path in paths:
for file in Path('.').glob(path):
try:
mio.import_landmark_file(
str(Path(file.parent.parent / 'BoundingBoxes' / (file.stem + '.pts')))
)
except ValueError:
continue
image_paths.append(file)
print('Got all image paths...')
# Second: split to train, test and validate. 7:2:1
if Path(path_base / 'train_img.txt').exists():
with Path(path_base / 'train_img.txt').open('rb') as train_ifs, \
Path(path_base / 'test_img.txt').open('rb') as test_ifs, \
Path(path_base / 'val_img.txt').open('rb') as val_ifs:
train_paths = [Path(line[:-1].decode('utf-8')) for line in train_ifs.readlines()]
test_paths = [Path(line[:-1].decode('utf-8')) for line in test_ifs.readlines()]
val_paths = [Path(line[:-1].decode('utf-8')) for line in val_ifs.readlines()]
else:
random.shuffle(image_paths)
num_train = int(len(image_paths) * 0.7)
num_test = int(len(image_paths) * 0.2)
train_paths = sorted(image_paths[:num_train])
test_paths = sorted(image_paths[num_train:num_train+num_test])
val_paths = sorted(image_paths[num_train+num_test:])
with Path(path_base / 'train_img.txt').open('wb') as train_ofs, \
Path(path_base / 'test_img.txt').open('wb') as test_ofs, \
Path(path_base / 'val_img.txt').open('wb') as val_ofs:
train_ofs.writelines([str(line).encode('utf-8') + b'\n' for line in train_paths])
test_ofs.writelines([str(line).encode('utf-8') + b'\n' for line in test_paths])
val_ofs.writelines([str(line).encode('utf-8') + b'\n' for line in val_paths])
print('Found Train/Test/Validate {}/{}/{}'.format(len(train_paths), len(test_paths), len(val_paths)))
# Third: export reference shape on train
if Path(path_base / 'reference_shape.pkl').exists():
reference_shape = PointCloud(mio.import_pickle(path_base / 'reference_shape.pkl'))
else:
reference_shape = PointCloud(build_reference_shape(train_paths, num_patches))
mio.export_pickle(reference_shape.points, path_base / 'reference_shape.pkl', overwrite=True)
print('Created reference_shape.pkl')
# Fourth: image shape & pca
image_shape = [0, 0, 3] # [H, W, C]
if Path(path_base / 'pca.bin').exists() and Path(path_base / 'meta.txt').exists():
with Path(path_base / 'meta.txt').open('r') as ifs:
image_shape = [int(x) for x in ifs.read().split(' ')]
else:
with tf.io.TFRecordWriter(str(path_base / 'pca.bin')) as ofs:
counter = 0
for path in train_paths:
counter += 1
if verbose:
status = 10.0 * counter / len(train_paths)
status_str = '\rPreparing {:2.2f}%['.format(status * 10)
for i in range(int(status)):
status_str += '='
for i in range(int(status), 10):
status_str += ' '
status_str += '] {} '.format(path)
print(status_str, end='')
mp_image = mio.import_image(path)
mp_image.landmarks['bb'] = mio.import_landmark_file(
str(Path(mp_image.path.parent.parent / 'BoundingBoxes' / (mp_image.path.stem + '.pts')))
)
mp_image = mp_image.crop_to_landmarks_proportion(0.3, group='bb')
mp_image = mp_image.rescale_to_pointcloud(reference_shape, group='PTS')
mp_image = grey_to_rgb(mp_image)
assert(mp_image.pixels.shape[0] == image_shape[2])
image_shape[0] = max(mp_image.pixels.shape[1], image_shape[0])
image_shape[1] = max(mp_image.pixels.shape[2], image_shape[1])
features = tf.train.Features(
feature={
'pca/shape': tf.train.Feature(
float_list=tf.train.FloatList(value=mp_image.landmarks['PTS'].points.flatten())
),
'pca/bb': tf.train.Feature(
float_list=tf.train.FloatList(value=mp_image.landmarks['bb'].points.flatten())
),
}
)
ofs.write(tf.train.Example(features=features).SerializeToString())
if verbose:
print('')
with Path(path_base / 'meta.txt').open('w') as ofs:
for s in image_shape[:-1]:
ofs.write('{} '.format(s))
ofs.write('{}'.format(image_shape[-1]))
print('Image shape', image_shape)
# Fifth: train data
if Path(path_base / 'train.bin').exists():
pass
else:
random.shuffle(train_paths)
with tf.io.TFRecordWriter(str(path_base / 'train.bin')) as ofs:
print('Preparing train data...')
counter = 0
for path in train_paths:
counter += 1
if verbose:
status = 10.0 * counter / len(train_paths)
status_str = '\rPreparing {:2.2f}%['.format(status * 10)
for i in range(int(status)):
status_str += '='
for i in range(int(status), 10):
status_str += ' '
status_str += '] {} '.format(path)
print(status_str, end='')
mp_image = mio.import_image(path)
mp_image.landmarks['bb'] = mio.import_landmark_file(
str(Path(mp_image.path.parent.parent / 'BoundingBoxes' / (mp_image.path.stem + '.pts')))
)
mp_image = mp_image.crop_to_landmarks_proportion(0.3, group='bb')
mp_image = mp_image.rescale_to_pointcloud(reference_shape, group='PTS')
mp_image = grey_to_rgb(mp_image)
# Padding to the same size
height, width = mp_image.pixels.shape[1:] # [C, H, W]
dy = max(int((image_shape[0] - height - 1) / 2), 0)
dx = max(int((image_shape[1] - width - 1) / 2), 0)
padded_image = np.random.rand(*image_shape).astype(np.float32)
padded_image[dy:(height + dy), dx:(width + dx), :] = mp_image.pixels.transpose(1, 2, 0)
padded_landmark = mp_image.landmarks['PTS'].points
padded_landmark[:, 0] += dy
padded_landmark[:, 1] += dx
features = tf.train.Features(
feature={
'train/image': tf.train.Feature(
bytes_list=tf.train.BytesList(value=[tf.compat.as_bytes(padded_image.tostring())])
),
'train/shape': tf.train.Feature(
float_list=tf.train.FloatList(value=padded_landmark.flatten())
)
}
)
ofs.write(tf.train.Example(features=features).SerializeToString())
if verbose:
print('')
# Sixth: test data
if Path(path_base / 'test.bin').exists():
pass
else:
with tf.io.TFRecordWriter(str(path_base / 'test.bin')) as ofs:
print('Preparing test data...')
counter = 0
for path in test_paths:
counter += 1
if verbose:
status = 10.0 * counter / len(test_paths)
status_str = '\rPreparing {:2.2f}%['.format(status * 10)
for i in range(int(status)):
status_str += '='
for i in range(int(status), 10):
status_str += ' '
status_str += '] {} '.format(path)
print(status_str, end='')
mp_image = mio.import_image(path)
mp_image.landmarks['bb'] = mio.import_landmark_file(
str(Path(mp_image.path.parent.parent / 'BoundingBoxes' / (mp_image.path.stem + '.pts')))
)
mp_image = mp_image.crop_to_landmarks_proportion(0.3, group='bb')
mp_bb = mp_image.landmarks['bb'].bounding_box()
mp_image.landmarks['init'] = align_shape_with_bounding_box(reference_shape, mp_bb)
mp_image = mp_image.rescale_to_pointcloud(reference_shape, group='init')
mp_image = grey_to_rgb(mp_image)
# Padding to the same size
height, width = mp_image.pixels.shape[1:] # [C, H, W]
dy = max(int((256 - height - 1) / 2), 0) # 200*(1+0.3*2)/sqrt(2) == 226.7
dx = max(int((256 - width - 1) / 2), 0) # 200*(1+0.3*2)/sqrt(2) == 226.7
padded_image = np.random.rand(256, 256, 3).astype(np.float32)
padded_image[dy:(height + dy), dx:(width + dx), :] = mp_image.pixels.transpose(1, 2, 0)
padded_landmark = mp_image.landmarks['PTS'].points
padded_landmark[:, 0] += dy
padded_landmark[:, 1] += dx
padded_init_landmark = mp_image.landmarks['init'].points
padded_init_landmark[:, 0] += dy
padded_init_landmark[:, 1] += dx
features = tf.train.Features(
feature={
'test/image': tf.train.Feature(
bytes_list=tf.train.BytesList(
value=[tf.compat.as_bytes(padded_image.tostring())])
),
'test/shape': tf.train.Feature(
float_list=tf.train.FloatList(value=padded_landmark.flatten())
),
'test/init': tf.train.Feature(
float_list=tf.train.FloatList(value=padded_init_landmark.flatten())
)
}
)
ofs.write(tf.train.Example(features=features).SerializeToString())
if verbose:
print('')
def load_images(paths,
group1=None,
group2=None,
verbose=True,
PLOT=False,
AFLW=False,
PLOT_shape=True):
"""Loads and rescales input knn_2D to the diagonal of the reference shape.
Args:
paths: a list of strings containing the data directories.
reference_shape (meanshape): a numpy array [num_landmarks, 2]
group: landmark group containing the grounth truth landmarks.
verbose: boolean, print debugging info.
Returns:
knn_2D: a list of numpy arrays containing knn_2D.
shapes: a list of the ground truth landmarks.
reference_shape (meanshape): a numpy array [num_landmarks, 2].
shape_gen: PCAModel, a shape generator.
"""
images = []
shapes = []
bbs = []
inits = []
shape_space = []
plot_shape_x = []
plot_shape_y = []
# compute mean shape
# if AFLW:
# # reference_shape = PointCloud(mio.import_pickle(Path('/home/hliu/gmh/RL_FA/mdm_aflw/ckpt/train_aflw') / 'reference_shape.pkl'))
# reference_shape = mio.import_pickle(
# Path('/home/hliu/data2/CongcongZhu/ICME/RDN/ckpt/pred') / 'reference_shape.pkl')
# else:
reference_shape = PointCloud(build_reference_shape(paths))
for path in paths:
if verbose:
print('Importing data from {}'.format(path))
for im in mio.import_images(path, verbose=verbose, as_generator=True):
# group = group or im.landmarks[group]._group_label
# group = group or im.landmarks.keys()[0]
group1 = 'PTS'
bb_root = im.path.parent.relative_to(im.path.parent.parent.parent)
if 'set' not in str(bb_root):
bb_root = im.path.parent.relative_to(im.path.parent.parent)
if AFLW:
im.landmarks['bb'] = im.landmarks['PTS'].lms.bounding_box()
else:
im.landmarks['bb'] = mio.import_landmark_file(
str(Path('bbs') / bb_root / (im.path.stem + '.pts')))
im = im.crop_to_landmarks_proportion(0.3, group='bb')
bb = im.landmarks['bb'].lms.bounding_box()
im.landmarks['initial'] = align_shape_with_bounding_box(
reference_shape, bb)
im = im.rescale_to_pointcloud(reference_shape, group=group1)
im = grey_to_rgb(im)
images.append(im.pixels.transpose(1, 2, 0))
inits.append(im.landmarks['initial'].lms)
shapes.append(im.landmarks[group1].lms)
shape_space.append(im.landmarks[group1].lms.points)
bbs.append(im.landmarks['bb'].lms)
if PLOT_shape:
x_tmp = np.sum((im.landmarks[group1].lms.points[:, 0] -
reference_shape.points[:, 0]))
y_tmp = np.sum((im.landmarks[group1].lms.points[:, 1] -
reference_shape.points[:, 1]))
if x_tmp < 0 and y_tmp < 0:
plot_shape_x.append(x_tmp)
plot_shape_y.append(y_tmp)
shape_space = np.array(shape_space)
print('shape_space:', shape_space.shape)
train_dir = Path(FLAGS.train_dir)
# centers = utils.k_means(shape_space, 100)
# centers = np.reshape(centers, [-1, 68, 2])
# np.save(train_dir/'shape_space_origin.npy', centers)
# print('created shape_space.npy using the {} group'.format(group))
# exit(0)
if PLOT_shape:
k_nn_plot_x = []
k_nn_plot_y = []
centers = utils.k_means(shape_space, 100)
centers = np.reshape(centers, [-1, 68, 2])
for i in range(centers.shape[0]):
x_tmp = np.sum((centers[i, :, 0] - reference_shape.points[:, 0]))
y_tmp = np.sum((centers[i, :, 1] - reference_shape.points[:, 1]))
if x_tmp < 0 and y_tmp < 0:
k_nn_plot_x.append(x_tmp)
k_nn_plot_y.append(y_tmp)
# pdb.set_trace()
# plt.scatter(plot_shape_x, plot_shape_y, s=20)
# plt.scatter(k_nn_plot_x, k_nn_plot_y, s=40)
# plt.xticks(())
# plt.yticks(())
# plt.show()
# pdb.set_trace()
mio.export_pickle(reference_shape.points,
train_dir / 'reference_shape.pkl',
overwrite=True)
print('created reference_shape.pkl using the {} group'.format(group1))
pca_model = detect.create_generator(shapes, bbs)
# Pad knn_2D to max length
max_shape = np.max([im.shape for im in images], axis=0)
max_shape = [len(images)] + list(max_shape)
padded_images = np.random.rand(*max_shape).astype(np.float32)
print(padded_images.shape,
'====================================================')
if PLOT:
# plot without padding
centers = utils.k_means(shape_space, 100)
centers = np.reshape(centers, [-1, 68, 2])
plot_img = cv2.imread('a.png').transpose(2, 0, 1)
centers_tmp = np.zeros(centers.shape)
# menpo_img = mio.import_image('a.png')
menpo_img = menpo.image.Image(plot_img)
for i in range(centers.shape[0]):
menpo_img.view()
min_y = np.min(centers[i, :, 0])
min_x = np.min(centers[i, :, 1])
centers_tmp[i, :, 0] = centers[i, :, 0] - min_y + 20
centers_tmp[i, :, 1] = centers[i, :, 1] - min_x + 20
print(centers_tmp[i, :, :])
menpo_img.landmarks['center'] = PointCloud(centers_tmp[i, :, :])
menpo_img.view_landmarks(group='center',
marker_face_colour='b',
marker_size='16')
# menpo_img.landmarks['center'].view(render_legend=True)
plt.savefig('plot_shape_space/' + str(i) + '.png')
plt.close()
exit(0)
# !!!shape_space without delta, which means shape_space has already been padded!
# delta = np.zeros(shape_space.shape)
gts = []
inis = []
for i, im in enumerate(images):
height, width = im.shape[:2]
dy = max(int((max_shape[1] - height - 1) / 2), 0)
dx = max(int((max_shape[2] - width - 1) / 2), 0)
lms = shapes[i]
pts = lms.points
init = inits[i]
init_shape = init.points
pts[:, 0] += dy
pts[:, 1] += dx
init_shape[:, 0] += dy
init_shape[:, 1] += dx
shape_space[i, :, 0] += dy
shape_space[i, :, 1] += dx
# delta[i][:, 0] = dy
# delta[i][:, 1] = dx
lms = lms.from_vector(pts)
init = init.from_vector(init_shape)
padded_images[i, dy:(height + dy), dx:(width + dx)] = im
gts.append(lms)
inis.append(init)
# shape_space = np.concatenate((shape_space, delta), 2)
centers = utils.k_means(shape_space, 100)
centers = np.reshape(centers, [-1, 68, 2])
np.save(train_dir / 'shape_space.npy', centers)
print('created shape_space.npy using the {} group'.format(group1))
return padded_images, gts, reference_shape.points.astype(
'float32'), pca_model, centers, inis
