def read_data(self, filename_queue, has_3d=False):
with tf.name_scope(None, 'read_data', [filename_queue]):
reader = tf.TFRecordReader()
_, example_serialized = reader.read(filename_queue)
if has_3d:
image, image_size, label, center, fname, pose, shape,\
gt3d, has_smpl3d = data_utils.parse_example_proto(
example_serialized, has_3d=has_3d)
# Need to send pose bc image can get flipped.
image, label, pose, gt3d = self.image_preprocessing(image,
image_size,
label,
center,
pose=pose,
gt3d=gt3d)
# Convert pose to rotation.
# Do not ignore the global!!
rotations = batch_rodrigues(tf.reshape(pose, [-1, 3]))
gt3d_flat = tf.reshape(gt3d, [-1])
# Label 3d is:
# [rotations, shape-beta, 3Djoints]
# [216=24*3*3, 10, 42=14*3]
label3d = tf.concat(
[tf.reshape(rotations, [-1]), shape, gt3d_flat], 0)
# label should be K x 3
label = tf.transpose(label)
if has_3d:
return image, label, label3d, has_smpl3d
def get_smpl_loader_temporal(self):
"""
Loads real delta poses from mocap data.
Returns:
delta_pose_batch (BxTx216).
"""
data_dirs = [
join(self.dataset_dir, 'mocap_neutrMosh_temporal_pose',
'neutrSMPL_{}_*.tfrecord'.format(dataset))
for dataset in self.mocap_datasets
]
files = []
for data_dir in data_dirs:
files += glob(data_dir)
if len(files) == 0:
print('Couldnt find any files!!')
import ipdb
ipdb.set_trace()
with tf.name_scope('input_smpl_loader_temporal'):
filename_queue = tf.train.string_input_producer(files,
shuffle=True,
capacity=128)
mosh_batch_size = self.batch_size * self.T
min_after_dequeue = 32
num_threads = 2
capacity = min_after_dequeue + (num_threads + 4) * mosh_batch_size
pose, deltas = data_utils.read_smpl_data_temporal(filename_queue)
pose = batch_rodrigues(tf.reshape(pose, (-1, 3)))
delta_pose = batch_rodrigues(tf.reshape(deltas, (-1, 3)))
pose_batch, delta_pose_batch = tf.train.batch(
[pose, delta_pose],
batch_size=mosh_batch_size,
num_threads=num_threads,
capacity=capacity,
name='input_smpl_batch')
return (tf.reshape(pose_batch, (self.batch_size, self.T, 216)),
tf.reshape(delta_pose_batch,
(self.batch_size, self.T, 216)))
def __init__(self, config, poses_aa, shapes, joints, kps, batch_size=None):
super(OmegasGt, self).__init__(config, batch_size=batch_size)
self.length = poses_aa.shape[1]
self.poses_aa = poses_aa
poses_rot = batch_rodrigues(tf.reshape(poses_aa, (-1, 3)))
self.poses_rot = tf.reshape(poses_rot, (self.batch_size, -1, 24, 3, 3))
self.shapes = shapes
self.joints = joints
self.kps = kps
self.deltas_rot = compute_deltas_batched(
poses_prev=self.poses_rot[:, :-1],
poses_curr=self.poses_rot[:, 1:])
def compute_smpl(self):
"""
Batch computation of vertices, joints, rotation matrices, and keypoints.
Due to the overhead added to computation graph, call this once.
"""
if self.smpl_computed:
print('SMPL should only be computed once!')
B = self.batch_size
T = self.length
verts, joints, poses_rot, transforms = self.smpl(
beta=tf.reshape(self.shapes, (B * T, 10)),
theta=tf.reshape(self.poses_aa, (B * T, 24, 3)),
get_skin=True)
# print(verts)
self.joints = tf.reshape(joints, (B, T, self.config.num_kps, 3))
self.poses_rot = tf.reshape(poses_rot, (B, T, 24, 3, 3))
# Make sure joints are B*T x num_kps x 3.
if self.use_optcam and self.is_training:
print('Using optimal camera!!')
# Just drop the z here ([1, 0, 0])
kps = joints[:, :, :2]
else:
kps = batch_orth_proj_idrot(joints,
tf.reshape(self.cams, (B * T, 3)))
self.kps = tf.reshape(kps, (B, T, self.config.num_kps, 2))
if self.deltas_aa.shape[1] != 0:
deltas_rot = batch_rodrigues(tf.reshape(self.deltas_aa, (-1, 3)))
self.deltas_rot = tf.reshape(deltas_rot,
(self.batch_size, -1, 24, 3, 3))
self.all_verts = tf.reshape(verts,
(B, T, 6890, 3))[:self.vis_max_batch]
if self.vis_t_indices is None:
self.verts = self.all_verts
else:
self.verts = Omegas.gather(values=self.all_verts,
indices=self.vis_t_indices)
# george modify
self.all_trans = tf.reshape(transforms,
(B, T, 24, 3))[:self.vis_max_batch]
if self.vis_t_indices is None:
self.transform = self.all_trans
else:
self.transform = Omegas.gather(values=self.all_trans,
indices=self.vis_t_indices)
self.smpl_computed = True
def setup_discriminator(self, fake_rotations, fake_shapes):
# Compute the rotation matrices of "rea" pose.
# These guys are in 24 x 3.
real_rotations = batch_rodrigues(tf.reshape(self.pose_loader, [-1, 3]))
real_rotations = tf.reshape(real_rotations, [-1, 24, 9])
# Ignoring global rotation. N x 23*9
# The # of real rotation is B*num_stage so it's balanced.
real_rotations = real_rotations[:, 1:, :]
all_fake_rotations = tf.reshape(tf.concat(
fake_rotations, 0), [self.batch_size * self.num_stage, -1, 9])
comb_rotations = tf.concat([real_rotations, all_fake_rotations],
0,
name="combined_pose")
comb_rotations = tf.expand_dims(comb_rotations, 2)
all_fake_shapes = tf.concat(fake_shapes, 0)
comb_shapes = tf.concat([self.shape_loader, all_fake_shapes],
0,
name="combined_shape")
disc_input = {
'weight_decay': self.d_wd,
'shapes': comb_shapes,
'poses': comb_rotations
}
self.d_out, self.D_var = Discriminator_separable_rotations(
**disc_input)
self.d_out_real, self.d_out_fake = tf.split(self.d_out, 2)
# Compute losses:
with tf.name_scope("comp_d_loss"):
self.d_loss_real = tf.reduce_mean(
tf.reduce_sum((self.d_out_real - 1)**2, axis=1))
self.d_loss_fake = tf.reduce_mean(
tf.reduce_sum((self.d_out_fake)**2, axis=1))
# Encoder loss
self.e_loss_disc = tf.reduce_mean(
tf.reduce_sum((self.d_out_fake - 1)**2, axis=1))
def __call__(self, beta, theta, get_skin=False, name=None):
"""
Obtain SMPL with shape (beta) & pose (theta) inputs.
Theta includes the global rotation.
Args:
beta: N x 10
theta: N x 72 (with 3-D axis-angle rep)
Updates:
self.J_transformed: N x 24 x 3 joint location after shaping
& posing with beta and theta
Returns:
- joints: N x 19 or 14 x 3 joint locations depending on joint_type
If get_skin is True, also returns
- Verts: N x 6890 x 3
"""
with tf.name_scope(name, "smpl_main", [beta, theta]):
num_batch = beta.shape[0].value
# 1. Add shape blend shapes
# (N x 10) x (10 x 6890*3) = N x 6890 x 3
v_shaped = tf.reshape(
tf.matmul(beta, self.shapedirs, name='shape_bs'),
[-1, self.size[0], self.size[1]]) + self.v_template
# 2. Infer shape-dependent joint locations.
Jx = tf.matmul(v_shaped[:, :, 0], self.J_regressor)
Jy = tf.matmul(v_shaped[:, :, 1], self.J_regressor)
Jz = tf.matmul(v_shaped[:, :, 2], self.J_regressor)
J = tf.stack([Jx, Jy, Jz], axis=2)
# 3. Add pose blend shapes
# N x 24 x 3 x 3
# only do rodrigues if theta has axis angle representation
Rs = tf.reshape(batch_rodrigues(tf.reshape(theta, [-1, 3])),
[-1, 24, 3, 3])
with tf.name_scope("lrotmin"):
# Ignore global rotation.
pose_feature = tf.reshape(Rs[:, 1:, :, :] - tf.eye(3),
[-1, 207])
# (N x 207) x (207, 20670) -> N x 6890 x 3
v_posed = tf.reshape(tf.matmul(pose_feature, self.posedirs),
[-1, self.size[0], self.size[1]]) + v_shaped
#4. Get the global joint location
self.J_transformed, A = batch_global_rigid_transformation(
Rs, J, self.parents)
# 5. Do skinning:
# W is N x 6890 x 24
W = tf.reshape(tf.tile(self.weights, [num_batch, 1]),
[num_batch, -1, 24])
# (N x 6890 x 24) x (N x 24 x 16)
T = tf.reshape(tf.matmul(W, tf.reshape(A, [num_batch, 24, 16])),
[num_batch, -1, 4, 4])
v_posed_homo = tf.concat(
[v_posed, tf.ones([num_batch, v_posed.shape[1], 1])], 2)
v_homo = tf.matmul(T, tf.expand_dims(v_posed_homo, -1))
verts = v_homo[:, :, :3, 0]
# Get cocoplus or lsp joints:
joint_x = tf.matmul(verts[:, :, 0], self.joint_regressor)
joint_y = tf.matmul(verts[:, :, 1], self.joint_regressor)
joint_z = tf.matmul(verts[:, :, 2], self.joint_regressor)
joints = tf.stack([joint_x, joint_y, joint_z], axis=2)
if get_skin:
return verts, joints, Rs
else:
return joints
def get_smpl_loader(self):
"""
Loads real pose and shape parameters from mocap data.
Returns:
pose_batch (BxTx216).
shape_batch (BxTx10).
"""
data_dirs = [
join(self.dataset_dir, 'mocap_neutrMosh',
'neutrSMPL_{}_*.tfrecord'.format(dataset))
for dataset in self.mocap_datasets
]
files = []
for data_dir in data_dirs:
files += glob(data_dir)
if len(files) == 0:
print('Couldnt find any files!!')
import ipdb
ipdb.set_trace()
with tf.name_scope('input_smpl_loader'):
filename_queue = tf.train.string_input_producer(files,
shuffle=True,
capacity=128)
mosh_batch_size = self.batch_size * self.T
if self.predict_delta:
new_t = len(self.delta_t_values) * self.T
batch_size_delta = self.batch_size * new_t
else:
batch_size_delta = 0
if self.do_hallucinate:
mosh_batch_size *= 2
if self.do_hallucinate_preds:
batch_size_delta *= 2
mosh_batch_size += batch_size_delta
if self.config.use_hmr_only:
# /2 bc doesnt have the movie-strip branch.
mosh_batch_size = int(mosh_batch_size / 2)
if self.config.use_hmr_only and not self.do_hallucinate:
mosh_batch_size = self.T * self.batch_size
min_after_dequeue = 32
num_threads = 2
capacity = min_after_dequeue + (num_threads + 4) * mosh_batch_size
pose, shape = data_utils.read_smpl_data(filename_queue)
pose = batch_rodrigues(tf.reshape(pose, (-1, 3)))
pose_batch, shape_batch = tf.train.batch(
[pose, shape],
batch_size=mosh_batch_size,
num_threads=num_threads,
capacity=capacity,
name='input_smpl_batch')
pose_batch = tf.reshape(pose_batch, (mosh_batch_size, 216))
shape_batch = tf.reshape(shape_batch, (mosh_batch_size, 10))
return pose_batch, shape_batch