def main(
caffe_prototxt, # pylint: disable=too-many-arguments, too-many-locals, too-many-statements
caffe_model,
image_list_file,
dset_root,
output_folder,
caffe_install_path):
"""Store and visualize the pose results for a model."""
LOGGER.info("Storing pose results to folder `%s`.", output_folder)
LOGGER.info("Using caffe from `%s`.", caffe_install_path)
sys.path.insert(0, path.join(caffe_install_path))
import caffe # pylint: disable=import-error
caffe.set_mode_gpu()
with open(image_list_file, 'r') as inf:
image_list = inf.readlines()
image_list = [
path.join(dset_root, line[1:]) for line in image_list
if line.startswith('/')
]
n_landmarks = int(image_list_file[image_list_file.find('_') +
1:image_list_file.find('_') + 3])
for imgnames in tqdm.tqdm(image_list):
imgname = imgnames.split(" ")[0].strip()
LOGGER.debug("Processing `%s`...", imgname)
im = caffe.io.load_image(imgname) # pylint: disable=invalid-name
# caffe.io loads as RGB, and in range [0., 1.].
im = (im * 255.)[:, :, ::-1].astype('uint8')
landmark_locs = estimate_pose(im, caffe_prototxt, caffe_model, caffe)
if landmark_locs.shape[1] == 91 and n_landmarks == 14:
# Extract LSP joints.
landmark_locs = landmark_locs[:, reduction_91tolsp]
np.save(path.join(output_folder,
path.basename(imgname) + '.npy'), landmark_locs)
vis_im = visualize_pose(im[:, :, ::-1],
landmark_locs,
scale=1.,
dash_length=5)
sm.imsave(
path.join(output_folder,
path.basename(imgname) + '.npy.vis.png'), vis_im)
def add_dataset(
dset_fp,
dset_fromroot,
list_ids,
up3d_fp, # pylint: disable=too-many-locals, too-many-arguments, too-many-statements, too-many-branches
train_list_f,
val_list_f,
train_val_list_f,
test_list_f,
scale_f,
train_spec,
val_spec,
test_spec,
target_person_size,
landmarks,
train_crop,
test_crop,
running_idx,
only_missing=False,
with_rlswap=True,
write_gtjoints_as_lm=False,
human_annotations=False):
"""Add a dataset to the collection."""
test_ids = [int(id_[1:6]) for id_ in test_spec]
train_ids = [int(id_[1:6]) for id_ in train_spec]
val_ids = [int(id_[1:6]) for id_ in val_spec]
LOGGER.info("Split: %d train, %d val, %d test.", len(train_ids),
len(val_ids), len(test_ids))
LOGGER.info("Writing dataset...")
for im_idx in tqdm.tqdm(train_ids + val_ids + test_ids):
image = scipy.misc.imread(
path.join(up3d_fp, '%05d_image.png' % (im_idx)))
with open(path.join(up3d_fp, '%05d_fit_crop_info.txt' % (im_idx)),
'r') as inf:
cropinfo = [int(val) for val in inf.readline().strip().split()]
assert image.ndim == 3
out_exists = (path.exists(
path.join(dset_fp, '%05d_image.png' % (running_idx)))
and path.exists(
path.join(dset_fp, '%05d_ann_vis.png' %
(running_idx))))
if with_rlswap and im_idx not in test_ids:
out_exists = out_exists and (path.exists(
path.join(dset_fp, '%05d_image.png' %
(running_idx + 1))) and path.exists(
path.join(dset_fp, '%05d_ann_vis.png' %
(running_idx + 1))))
if not (only_missing and out_exists or write_gtjoints_as_lm):
if human_annotations:
landmark_pos = np.load(
path.join(up3d_fp, '%05d_joints.npy' % (im_idx)))
else:
landmark_pos = get_landmark_positions(
path.join(up3d_fp, '%05d_body.pkl' % (im_idx)),
(cropinfo[1], cropinfo[0]), landmarks)
fac_y = cropinfo[0] / float(cropinfo[3] - cropinfo[2])
fac_x = cropinfo[1] / float(cropinfo[5] - cropinfo[4])
landmark_pos[:2, :] /= np.mean([fac_x, fac_y])
landmark_pos[0, :] += cropinfo[4]
landmark_pos[1, :] += cropinfo[2]
joints = np.load(path.join(up3d_fp, '%05d_joints.npy' % (im_idx)))
joints = np.vstack((joints, np.all(joints > 0, axis=0)[None, :]))
person_size = robust_person_size(joints)
norm_factor = float(target_person_size) / person_size
joints[:2, :] *= norm_factor
if not (only_missing and out_exists or write_gtjoints_as_lm):
landmark_pos[:2, :] *= norm_factor
if write_gtjoints_as_lm:
landmark_pos = joints.copy()
image = scipy.misc.imresize(image, norm_factor, interp='bilinear')
if im_idx in test_ids:
crop = test_crop
else:
crop = train_crop
if image.shape[0] > crop or image.shape[1] > crop:
LOGGER.debug(
"Image (original %d, here %d) too large (%s)! Cropping...",
im_idx, running_idx, str(image.shape[:2]))
person_center = np.mean(joints[:2, joints[2, :] == 1], axis=1)
crop_y, crop_x = get_crop(image, person_center, crop)
image = image[crop_y[0]:crop_y[1], crop_x[0]:crop_x[1], :]
landmark_pos[0, :] -= crop_x[0]
landmark_pos[1, :] -= crop_y[0]
assert image.shape[0] == crop or image.shape[1] == crop, (
"Error cropping image (original %d, here %d)!" %
(im_idx, running_idx))
assert image.shape[0] <= crop and image.shape[
1] <= crop and image.shape[2] == 3, (
"Wrong image shape (original %d, here %d)!" %
(im_idx, running_idx))
vis_im = vs.visualize_pose(image, landmark_pos, scale=1.)
if not (only_missing and out_exists):
scipy.misc.imsave(
path.join(dset_fp, '%05d_image.png' % (running_idx)), image)
scipy.misc.imsave(
path.join(dset_fp, '%05d_ann_vis.png' % (running_idx)), vis_im)
if with_rlswap and im_idx not in test_ids:
if landmark_pos.shape[1] == 14:
landmark_pos_swapped = landmark_pos[:, rlswap_lsp]
else:
landmark_pos_swapped = landmark_pos[:, rlswap_landmarks_91]
landmark_pos_swapped[
0, :] = image.shape[1] - landmark_pos_swapped[0, :]
image_swapped = image[:, ::-1, :]
# Use core visualization for 14 joints.
vis_im_swapped = vs.visualize_pose(image_swapped,
landmark_pos_swapped,
scale=1)
if not (only_missing and out_exists):
scipy.misc.imsave(
path.join(dset_fp, '%05d_image.png' % (running_idx + 1)),
image_swapped)
scipy.misc.imsave(
path.join(dset_fp, '%05d_ann_vis.png' % (running_idx + 1)),
vis_im_swapped)
list_fs = []
list_id_ids = []
if im_idx in train_ids:
list_fs.append(train_val_list_f)
list_id_ids.append(2)
list_fs.append(train_list_f)
list_id_ids.append(0)
elif im_idx in val_ids:
list_fs.append(train_val_list_f)
list_id_ids.append(2)
list_fs.append(val_list_f)
list_id_ids.append(1)
elif im_idx in test_ids:
list_fs.append(test_list_f)
list_id_ids.append(3)
for list_f, list_id_idx in zip(list_fs, list_id_ids):
# pylint: disable=bad-continuation
list_f.write("""# %d
%s
3
%d
%d
%d
""" % (list_ids[list_id_idx],
path.join('/' + dset_fromroot, '%05d_image.png' % (running_idx)),
image.shape[0], image.shape[1], landmark_pos.shape[1]))
for landmark_idx, landmark_point in enumerate(landmark_pos.T):
list_f.write("%d %d %d\n" %
(landmark_idx + 1, int(
landmark_point[0]), int(landmark_point[1])))
list_f.flush()
list_ids[list_id_idx] += 1
scale_f.write("%05d_image.png %f\n" % (running_idx, norm_factor))
scale_f.flush()
running_idx += 1
if with_rlswap and im_idx not in test_ids:
for list_f, list_id_idx in zip(list_fs, list_id_ids):
# pylint: disable=bad-continuation
list_f.write("""# %d
%s
3
%d
%d
%d
""" % (list_ids[list_id_idx],
path.join('/' + dset_fromroot, '%05d_image.png' % (running_idx)),
image.shape[0], image.shape[1], landmark_pos.shape[1]))
for landmark_idx, landmark_point in enumerate(
landmark_pos_swapped.T):
list_f.write(
"%d %d %d\n" % (landmark_idx + 1, int(
landmark_point[0]), int(landmark_point[1])))
list_f.flush()
list_ids[list_id_idx] += 1
scale_f.write("%05d_image.png %f\n" % (running_idx, norm_factor))
scale_f.flush()
running_idx += 1
return running_idx
def add_dataset(dset_fp, dset_rel_fp, up3d_fp, # pylint: disable=too-many-locals, too-many-arguments, too-many-statements, too-many-branches
train_list_f, val_list_f, test_list_f,
train_spec, val_spec, test_spec,
target_person_size, landmarks, partspec, crop, running_idx,
only_missing=False):
"""Add a dataset to the collection."""
test_ids = [int(id_[1:6]) for id_ in test_spec]
train_ids = [int(id_[1:6]) for id_ in train_spec]
val_ids = [int(id_[1:6]) for id_ in val_spec]
ids_list = sorted(train_ids + val_ids + test_ids)
LOGGER.info("Split: %d train, %d val, %d test.",
len(train_ids), len(val_ids), len(test_ids))
LOGGER.info("Writing dataset...")
for im_idx in tqdm.tqdm(ids_list):
image = scipy.misc.imread(path.join(up3d_fp, '%05d_image.png' % (im_idx)))
with open(path.join(up3d_fp, '%05d_fit_crop_info.txt' % (im_idx)), 'r') as inf:
cropinfo = [int(val) for val in inf.readline().strip().split()]
assert image.ndim == 3
out_exists = (path.exists(path.join(dset_fp, '%05d_image.png' % (running_idx))) and
path.exists(path.join(dset_fp, '%05d_ann.png' % (running_idx))) and
path.exists(path.join(dset_fp, '%05d_ann_vis.png' % (running_idx))) and
path.exists(path.join(dset_fp, '%05d_render.png' % (running_idx))) and
path.exists(path.join(dset_fp, '%05d_render_light.png' % (running_idx))))
if not (only_missing and out_exists):
rendering = uncrop(render_body_impl(path.join(up3d_fp, '%05d_body.pkl' % (im_idx)),
resolution=(cropinfo[1],
cropinfo[0]),
quiet=True,
use_light=False)[0],
image.shape[:2],
cropinfo)
rendering_l = uncrop(render_body_impl(path.join(up3d_fp, '%05d_body.pkl' % (im_idx)),
resolution=(cropinfo[1],
cropinfo[0]),
quiet=True,
use_light=True)[0],
image.shape[:2],
cropinfo)
joints = np.load(path.join(up3d_fp, '%05d_joints.npy' % (im_idx)))
joints = np.vstack((joints, np.all(joints > 0, axis=0)[None, :]))
person_size = robust_person_size(joints)
norm_factor = float(target_person_size) / person_size
landmark_pos = get_landmark_positions(path.join(up3d_fp, '%05d_body.pkl' % (im_idx)),
(cropinfo[1], cropinfo[0]),
landmarks)
fac_y = cropinfo[0] / float(cropinfo[3] - cropinfo[2])
fac_x = cropinfo[1] / float(cropinfo[5] - cropinfo[4])
landmark_pos[:2, :] /= np.mean([fac_x, fac_y])
landmark_pos[0, :] += cropinfo[4]
landmark_pos[1, :] += cropinfo[2]
landmark_pos[:2, :] *= norm_factor
if not (only_missing and out_exists):
image = scipy.misc.imresize(image, norm_factor, interp='bilinear')
rendering = scipy.misc.imresize(rendering, norm_factor, interp='nearest')
rendering_l = scipy.misc.imresize(rendering_l, norm_factor, interp='bilinear')
if image.shape[0] > crop or image.shape[1] > crop:
LOGGER.debug("Image (original %d, here %d) too large (%s)! Cropping...",
im_idx, running_idx, str(image.shape[:2]))
person_center = np.mean(joints[:2, joints[2, :] == 1], axis=1) * norm_factor
crop_y, crop_x = get_crop(image, person_center, crop)
image = image[crop_y[0]:crop_y[1],
crop_x[0]:crop_x[1], :]
rendering = rendering[crop_y[0]:crop_y[1],
crop_x[0]:crop_x[1], :]
rendering_l = rendering_l[crop_y[0]:crop_y[1],
crop_x[0]:crop_x[1], :]
landmark_pos[0, :] -= crop_x[0]
landmark_pos[1, :] -= crop_y[0]
assert image.shape[0] == crop or image.shape[1] == crop, (
"Error cropping image (original %d, here %d)!" % (im_idx,
running_idx))
assert image.shape[0] <= crop and image.shape[1] <= crop and image.shape[2] == 3, (
"Wrong image shape (original %d, here %d)!" % (im_idx, running_idx))
class_groups = six_region_groups if partspec == '6' else None
annotation = regions_to_classes(rendering, class_groups, warn_id=str(im_idx))
if partspec == '1':
annotation = (annotation > 0).astype('uint8')
assert np.max(annotation) <= int(partspec), (
"Wrong annotation value (original %d, here %d): %s!" % (
im_idx, running_idx, str(np.unique(annotation))))
if running_idx == 0:
assert np.max(annotation) == int(partspec), (
"Probably an error in the number of parts!")
pose_vis_im = vs.visualize_pose(cv2.cvtColor(annotation*8, cv2.COLOR_GRAY2RGB),
landmark_pos,
scale=1.)
scipy.misc.imsave(path.join(dset_fp, '%05d_image.png' % (running_idx)), image)
scipy.misc.imsave(path.join(dset_fp, '%05d_ann.png' % (running_idx)), annotation)
scipy.misc.imsave(path.join(dset_fp, '%05d_seg_ann_vis.png' % (running_idx)),
apply_colormap(annotation, vmax=int(partspec)))
# scipy.misc.imsave(path.join(dset_fp, '%05d_render.png' % (running_idx)), rendering)
scipy.misc.imsave(path.join(dset_fp, '%05d_render_light.png' % (running_idx)), rendering_l) # pylint: disable=line-too-long
scipy.misc.imsave(path.join(dset_fp, '%05d_pose_ann_vis.png' % (running_idx)), pose_vis_im)
landmark_pos = np.concatenate((landmark_pos, joints[2][None, :]))
np.save(str(path.join(dset_fp, '%05d_joints.npy' % (running_idx))), landmark_pos, allow_pickle=False)
if im_idx in train_ids:
list_f = train_list_f
elif im_idx in val_ids:
list_f = val_list_f
elif im_idx in test_ids:
list_f = test_list_f
list_f.write("/%s/%05d_image.png /%s/%05d_ann.png %f\n" % (
dset_rel_fp, running_idx, dset_rel_fp, running_idx, norm_factor))
list_f.flush()
running_idx += 1
return running_idx
def predict_pose_from(
image_name, # pylint: disable=too-many-arguments, too-many-branches
out_name=None,
scales='1.',
visualize=True,
write_pmaps=False,
folder_image_suffix='.png',
use_cpu=False,
every_nth=1):
"""
Load an image file, predict the pose and write it out.
`IMAGE_NAME` may be an image or a directory, for which all images with
`folder_image_suffix` will be processed.
"""
scales = [float(val) for val in scales.split(',')]
if _os.path.isdir(image_name):
folder_name = image_name[:]
_LOGGER.info(
"Specified image name is a folder. Processing all images "
"with suffix %s.", folder_image_suffix)
images = sorted(
_glob.glob(_os.path.join(folder_name, '*' + folder_image_suffix)))
images = [im_fp for im_fp in images if not im_fp.endswith('vis.png')]
images = images[::every_nth]
process_folder = True
else:
images = [image_name]
process_folder = False
if use_cpu:
_caffe.set_mode_cpu()
else:
_caffe.set_mode_gpu()
out_name_provided = out_name
if process_folder and out_name is not None and not _os.path.exists(
out_name):
_os.mkdir(out_name)
for image_name in images:
if out_name_provided is None:
out_name = image_name + '_pose.npz'
pmap_out_name = image_name + '_pose_probabilities.npz'
elif process_folder:
out_name = _os.path.join(
out_name_provided,
_os.path.basename(image_name) + '_pose.npz')
_LOGGER.info(
"Predicting the pose on `%s` (saving to `%s`) in best of "
"scales %s.", image_name, out_name, scales)
image = _scipy.misc.imread(image_name)
if image.ndim == 2:
_LOGGER.warn(
"The image is grayscale! This may deteriorate performance!")
image = _np.dstack((image, image, image))
else:
image = image[:, :, ::-1]
pose, pmaps = estimate_pose(image, scales)
_np.savez_compressed(out_name, pose=pose)
if write_pmaps:
_np.savez_compressed(pmap_out_name, pose_probabilities=pmaps)
if visualize:
visim = visualize_pose(image[:, :, ::-1],
pose,
dash_length=5,
opacity=0.8,
scale=1.)
vis_name = out_name + '_vis.png'
_scipy.misc.imsave(vis_name, visim)
if write_pmaps:
VIS_FAC = 255.
vismaps = _np.max(pmaps, axis=2)
visim = _np.tile(
(vismaps * VIS_FAC).astype('uint8')[:, :, None], (1, 1, 3))
if _CV2_AVAILABLE:
visim = _cv2.applyColorMap(
visim, _cv2.COLORMAP_AUTUMN)[:, :, ::-1]
visim = (visim.astype('float32') * 0.6 +
image.astype('float32') * 0.4).astype('uint8')
else:
visim = vismaps * 255
vis_name = pmap_out_name + '_vis.png'
_scipy.misc.imsave(vis_name, visim)
def add_dataset(dset_fp,
landmarks,
partspec,
resolution_wh=256,
num_zfill_in_name=5,
start=0,
num=0,
num_augment_per_sample=0,
global_pose_augment=False,
global_pose_max_angle=2 * np.pi,
flip_init_glob_pose=True,
shape_augment=False,
shape_range=[-2, 2],
fat_height_range=[-4, 4],
pose_augment=False,
pose_axis_deviation_scale=0.1,
pose_angle_deviation_range=[-np.pi / 6, np.pi / 6]):
"""Add a dataset to the collection."""
ids_list = [
str(f[:num_zfill_in_name]) for f in sorted(os.listdir(dset_fp))
if f.endswith("_body.pkl") and '-' not in f
]
to_render_ids = ids_list[start:start + num]
LOGGER.info("Writing dataset. Shape augment: {}, "
"Glob pose augment: {}, "
"Pose augment: {}".format(str(shape_augment),
str(global_pose_augment),
str(pose_augment)))
for im_idx in to_render_ids:
print('Index', im_idx)
smpl_path = path.join(dset_fp, '{}_body.pkl'.format(im_idx))
with open(smpl_path, 'rb') as f:
smpl_data = pickle.load(f)
pose = smpl_data['pose']
if 'betas' in smpl_data.keys():
betas = smpl_data['betas']
elif 'shape' in smpl_data.keys():
betas = smpl_data['shape']
rt = np.array([0.0, 0.0, 0.0])
f = np.array(5000.0)
trans = np.array([0.0, 0.0, 0.0])
t = np.array([0.0, 0.0, 40.0])
# ------- First render original, un-augmented data (if doing UP3D augmentation) -------
camera = {
'rt': rt,
'f': f,
'trans': trans,
't': t,
'betas': betas,
'pose': pose
}
resolution = (resolution_wh, resolution_wh)
factor = 1.0
renderings = render(MODEL_NEUTRAL, (np.asarray(resolution) * 1. /
factor).astype('int'),
camera,
1,
False,
use_light=False)
renderings = [
scipy.misc.imresize(renderim, (resolution[1], resolution[0]),
interp='nearest') for renderim in renderings
]
rendering = renderings[
0] # Only rendering one rotated view - single element in list
landmark_pos = get_landmark_positions(betas, pose, trans, resolution,
landmarks, rt, t, f)
vis = check_landmark_visibility(landmark_pos, resolution_wh)
class_groups = six_region_groups if partspec == '6' else None
annotation = regions_to_classes(rendering,
class_groups,
warn_id=im_idx)
if partspec == '1':
annotation = (annotation > 0).astype('uint8')
# assert np.max(annotation) <= int(partspec), (
# "Wrong annotation value (%s): %s!" % (im_idx, str(np.unique(annotation))))
# if int(im_idx) == 0:
# assert np.max(annotation) == int(partspec), ("Probably an error in the number of parts!")
pose_vis_im = vs.visualize_pose(cv2.cvtColor(annotation * 8,
cv2.COLOR_GRAY2RGB),
landmark_pos,
scale=1.)
scipy.misc.imsave(path.join(dset_fp, '{}_ann.png'.format(im_idx)),
annotation)
scipy.misc.imsave(
path.join(dset_fp, '{}_seg_ann_vis.png'.format(im_idx)),
apply_colormap(annotation, vmax=int(partspec)))
scipy.misc.imsave(
path.join(dset_fp, '{}_pose_ann_vis.png'.format(im_idx)),
pose_vis_im)
landmark_pos_with_vis = np.concatenate([landmark_pos, vis[None, :]],
axis=0)
np.save(str(path.join(dset_fp, '{}_joints.npy'.format(im_idx))),
landmark_pos_with_vis,
allow_pickle=False)
# --------------- Render augmented data (if doing UP3D augmentation) ---------------
# UP3D Augmentation by random sampling
if num_augment_per_sample > 0:
if global_pose_augment:
assert 'global_pose' in dset_fp, "Dataset path is probably wrong!"
# Random sampling of global rotations
new_r_globals = uniform_rotation_sample_global_pose(
global_pose_max_angle=global_pose_max_angle,
flip_init_global_pose=flip_init_glob_pose,
num=num_augment_per_sample - 1)
# First global rotation augmentation is set to be a backwards facing one
# since this is the second most common global pose modality after front-facing
R_global_backface = np.matmul(
cv2.Rodrigues(np.array([0, np.pi, 0]))[0],
cv2.Rodrigues(np.array([np.pi, 0, 0]))[0])
r_global_backface = np.squeeze(
cv2.Rodrigues(R_global_backface)[0])
r_global_backface += 0.1 * np.random.randn(
3) # add some random noise
new_r_globals.insert(0, r_global_backface)
if shape_augment:
assert 'shape' in dset_fp, "Dataset path is probably wrong!"
# Random sampling of shape deviations from original
betas_delta = uniform_sample_smpl_shape_deviation(
range=shape_range,
fat_and_height_range=fat_height_range,
num=num_augment_per_sample)
if pose_augment:
assert 'pose' in dset_fp, "Dataset path is probably wrong!"
# Random sampling of axis and angle deviations from original pose
new_poses = uniform_axis_angle_sample_pose_deviation(
pose[3:], pose_axis_deviation_scale,
pose_angle_deviation_range, num)
for aug_idx in range(num_augment_per_sample):
print('Aug', aug_idx)
aug_pose = np.copy(pose)
aug_betas = np.copy(betas)
if global_pose_augment:
aug_pose[:3] = new_r_globals[aug_idx]
if shape_augment:
aug_betas = aug_betas + betas_delta[aug_idx]
if pose_augment:
aug_pose[3:] = new_poses[aug_idx]
aug_camera = {
'rt': rt,
'f': f,
'trans': trans,
't': t,
'betas': aug_betas,
'pose': aug_pose
}
resolution = (resolution_wh, resolution_wh)
factor = 1.0
aug_renderings = render(MODEL_NEUTRAL, (np.asarray(resolution) *
1. / factor).astype('int'),
aug_camera,
1,
False,
use_light=False)
aug_renderings = [
scipy.misc.imresize(renderim, (resolution[1], resolution[0]),
interp='nearest')
for renderim in aug_renderings
]
aug_rendering = aug_renderings[
0] # Rendering 1 rotated view - 1 element in list
aug_landmark_pos = get_landmark_positions(aug_betas, aug_pose,
trans, resolution,
landmarks, rt, t, f)
aug_vis = check_landmark_visibility(aug_landmark_pos,
resolution_wh)
class_groups = six_region_groups if partspec == '6' else None
aug_annotation = regions_to_classes(aug_rendering,
class_groups,
warn_id=im_idx)
if partspec == '1':
aug_annotation = (aug_annotation > 0).astype('uint8')
# assert np.max(annotation) <= int(partspec), (
# "Wrong annotation value (%s): %s!" % (im_idx, str(np.unique(annotation))))
# if int(im_idx) == 0:
# assert np.max(annotation) == int(partspec), ("Probably an error in the number of parts!")
aug_pose_vis_im = vs.visualize_pose(cv2.cvtColor(
aug_annotation * 8, cv2.COLOR_GRAY2RGB),
aug_landmark_pos,
scale=1.)
scipy.misc.imsave(
path.join(dset_fp, '{}-{}_ann.png'.format(im_idx, aug_idx)),
aug_annotation)
scipy.misc.imsave(
path.join(dset_fp,
'{}-{}_seg_ann_vis.png'.format(im_idx, aug_idx)),
apply_colormap(aug_annotation, vmax=int(partspec)))
scipy.misc.imsave(
path.join(dset_fp,
'{}-{}_pose_ann_vis.png'.format(im_idx, aug_idx)),
aug_pose_vis_im)
aug_landmark_pos_with_vis = np.concatenate(
[aug_landmark_pos, aug_vis[None, :]], axis=0)
np.save(str(
path.join(dset_fp, '{}-{}_joints.npy'.format(im_idx,
aug_idx))),
aug_landmark_pos_with_vis,
allow_pickle=False)
aug_smpl_save_path = path.join(
dset_fp, '{}-{}_body.pkl'.format(im_idx, aug_idx))
with open(aug_smpl_save_path, 'wb') as aug_f:
pickle.dump({
'betas': aug_betas,
'pose': aug_pose
},
aug_f,
protocol=2)