python 三维重建_三维重建(5):HMR源码阅读1 TF-SMPL部分阅读

def batch_skew(vec, batch_size=None): # 生成反对称矩阵

"""

vec is N x 3, batch_size is int 单位旋转向量(由旋转角度构成)

returns N x 3 x 3. Skew_sym version of each matrix.

"""

with tf.name_scope("batch_skew", [vec]):

if batch_size is None:

batch_size = vec.shape.as_list()[0]

col_inds = tf.constant([1, 2, 3, 5, 6, 7]) # 列向量索引

indices = tf.reshape(

tf.reshape(tf.range(0, batch_size) * 9, [-1, 1]) + col_inds, # shape=(batchsize，6)每一行都是列索引

[-1, 1]) #shape=(batchsize*6,1)

updates = tf.reshape(

tf.stack(

[

-vec[:, 2], vec[:, 1], vec[:, 2], -vec[:, 0], -vec[:, 1],

vec[:, 0]

], # batchsize*6的矩阵

axis=1), [-1]) # shape=[batchsize*6]

out_shape = [batch_size * 9]

res = tf.scatter_nd(indices, updates, out_shape) #

res = tf.reshape(res, [batch_size, 3, 3])

return res

def batch_rodrigues(theta, name=None):

"""

Theta is N x 3

"""

with tf.name_scope(name, "batch_rodrigues", [theta]):

batch_size = theta.shape.as_list()[0]

# angle = tf.norm(theta, axis=1)

# r = tf.expand_dims(tf.div(theta, tf.expand_dims(angle + 1e-8, -1)), -1)

# angle = tf.expand_dims(tf.norm(theta, axis=1) + 1e-8, -1)

angle = tf.expand_dims(tf.norm(theta + 1e-8, axis=1), -1) # tf.norm求模(范数)，因为theta是轴角，所以求模就是总旋转角度

r = tf.expand_dims(tf.div(theta, angle), -1) # 转化成单位向量，然后套公式

angle = tf.expand_dims(angle, -1)

cos = tf.cos(angle)

sin = tf.sin(angle)

outer = tf.matmul(r, r, transpose_b=True, name="outer") # 计算外积

eyes = tf.tile(tf.expand_dims(tf.eye(3), 0), [batch_size, 1, 1]) # 单位矩阵

R = cos * eyes + (1 - cos) * outer + sin * batch_skew( # Rodrigues旋转公式

r, batch_size=batch_size)

return R # 返回旋转矩阵

def batch_lrotmin(theta, name=None):

""" NOTE: not used bc I want to reuse R and this is simple.

Output of this is used to compute joint-to-pose blend shape mapping.

Equation 9 in SMPL paper.

Args:

pose: `Tensor`, N x 72 vector holding the axis-angle rep of K joints.

This includes the global rotation so K=24

Returns

diff_vec : `Tensor`: N x 207 rotation matrix of 23=(K-1) joints with identity subtracted.,

"""

with tf.name_scope(name, "batch_lrotmin", [theta]):

with tf.name_scope("ignore_global"):

theta = theta[:, 3:]

# N*23 x 3 x 3

Rs = batch_rodrigues(tf.reshape(theta, [-1, 3]))

lrotmin = tf.reshape(Rs - tf.eye(3), [-1, 207])

return lrotmin

def batch_global_rigid_transformation(Rs, Js, parent, rotate_base=False):

"""

Computes absolute joint locations given pose.

rotate_base: if True, rotates the global rotation by 90 deg in x axis.

if False, this is the original SMPL coordinate.

Args:

Rs: N x 24 x 3 x 3 rotation vector of K joints

Js: N x 24 x 3, joint locations before posing

parent: 24 holding the parent id for each index

Returns

new_J : `Tensor`: N x 24 x 3 location of absolute joints 全局坐标

A : `Tensor`: N x 24 4 x 4 relative joint transformations for LBS. 相对变换

"""

with tf.name_scope("batch_forward_kinematics", [Rs, Js]):

N = Rs.shape[0].value # batchsize

if rotate_base:

print('Flipping the SMPL coordinate frame!!!!')

rot_x = tf.constant(

[[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=Rs.dtype)

rot_x = tf.reshape(tf.tile(rot_x, [N, 1]), [N, 3, 3])

root_rotation = tf.matmul(Rs[:, 0, :, :], rot_x)

else:

root_rotation = Rs[:, 0, :, :] # 根节点的旋转向量

# Now Js is N x 24 x 3 x 1

Js = tf.expand_dims(Js, -1)

def make_A(R, t, name=None):

# Rs is N x 3 x 3, ts is N x 3 x 1

with tf.name_scope(name, "Make_A", [R, t]):

R_homo = tf.pad(R, [[0, 0], [0, 1], [0, 0]])

t_homo = tf.concat([t, tf.ones([N, 1, 1])], 1)

return tf.concat([R_homo, t_homo], 2)

A0 = make_A(root_rotation, Js[:, 0])

results = [A0]

for i in range(1, parent.shape[0]):

j_here = Js[:, i] - Js[:, parent[i]]

A_here = make_A(Rs[:, i], j_here)

res_here = tf.matmul(

results[parent[i]], A_here, name="propA%d" % i)

results.append(res_here)

# 10 x 24 x 4 x 4

results = tf.stack(results, axis=1)

new_J = results[:, :, :3, 3]

# --- Compute relative A: Skinning is based on

# how much the bone moved (not the final location of the bone)

# but (final_bone - init_bone)

# ---

Js_w0 = tf.concat([Js, tf.zeros([N, 24, 1, 1])], 2)

init_bone = tf.matmul(results, Js_w0)

# Append empty 4 x 3:

init_bone = tf.pad(init_bone, [[0, 0], [0, 0], [0, 0], [3, 0]])

A = results - init_bone

return new_J, A

3. batchsize_smpl.py

"""

Tensorflow SMPL implementation as batch.

Specify joint types:

'coco': Returns COCO+ 19 joints

'lsp': Returns H3.6M-LSP 14 joints

Note: To get original smpl joints, use self.J_transformed

"""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import numpy as np

import cPickle as pickle

import tensorflow as tf

from .batch_lbs import batch_rodrigues, batch_global_rigid_transformation

# There are chumpy variables so convert them to numpy.

def undo_chumpy(x):

return x if isinstance(x, np.ndarray) else x.r

class SMPL(object):

def __init__(self, pkl_path, joint_type='cocoplus', dtype=tf.float32):

"""

pkl_path is the path to a SMPL model

"""

# -- Load SMPL params --

with open(pkl_path, 'r') as f:

dd = pickle.load(f)

# Mean template vertices

self.v_template = tf.Variable(

undo_chumpy(dd['v_template']),

name='v_template',

dtype=dtype,

trainable=False)

# Size of mesh [Number of vertices, 3]

self.size = [self.v_template.shape[0].value, 3]

self.num_betas = dd['shapedirs'].shape[-1]

# Shape blend shape basis: 6980 x 3 x 10

# reshaped to 6980*30 x 10, transposed to 10x6980*3

shapedir = np.reshape(

undo_chumpy(dd['shapedirs']), [-1, self.num_betas]).T

self.shapedirs = tf.Variable(

shapedir, name='shapedirs', dtype=dtype, trainable=False)

# Regressor for joint locations given shape - 6890 x 24

self.J_regressor = tf.Variable(

dd['J_regressor'].T.todense(),

name="J_regressor",

dtype=dtype,

trainable=False)

# Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*30 x 207

num_pose_basis = dd['posedirs'].shape[-1]

# 207 x 20670

posedirs = np.reshape(

undo_chumpy(dd['posedirs']), [-1, num_pose_basis]).T

self.posedirs = tf.Variable(

posedirs, name='posedirs', dtype=dtype, trainable=False)

# indices of parents for each joints

self.parents = dd['kintree_table'][0].astype(np.int32)

# LBS weights

self.weights = tf.Variable(

undo_chumpy(dd['weights']),

name='lbs_weights',

dtype=dtype,

trainable=False)

# This returns 19 keypoints: 6890 x 19

self.joint_regressor = tf.Variable(

dd['cocoplus_regressor'].T.todense(),

name="cocoplus_regressor",

dtype=dtype,

trainable=False)

if joint_type == 'lsp': # 14 LSP joints!

self.joint_regressor = self.joint_regressor[:, :14]

if joint_type not in ['cocoplus', 'lsp']:

print('BAD!! Unknown joint type: %s, it must be either "cocoplus" or "lsp"' % joint_type)

import ipdb

ipdb.set_trace()

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 6980 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 # 形状参数 * 从模型shapedirs里得到参数 + 顶点模型 = 形状 公式8

v_shaped = tf.reshape(

tf.matmul(beta, self.shapedirs, name='shape_bs'),

[-1, self.size[0], self.size[1]]) + self.v_template # 公式 10

# 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

Rs = tf.reshape( # 将theta写成24*3的shape，然后进行Rodrigues变换成旋转矩阵

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), # 公式9

[-1, self.size[0], self.size[1]]) + v_shaped # 公式 6

#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