Skip to content

Instantly share code, notes, and snippets.

@charliememory charliememory/pdm.py
Last active Mar 5, 2019

Embed
What would you like to do?
Tensorflow implement of Point Distribution Model && coord2channel function in CVPR 2018 paper "Natural and effective obfuscation by head inpainting"
#################### Point Distribution Model #####################
## ref to OpenFace PDM model https://github.com/TadasBaltrusaitis/OpenFace/tree/79d7116dae7f6b5335016dcb0e9ea64e1f931287/model_training/pdm_generation
def tf_Euler2Rot(euler):
batch_size = euler.get_shape().as_list()[0]
rx, ry, rz = tf.split(euler, 3, axis=1)
Rx_0 = tf.cast(tf.tile(tf.expand_dims([1.0, 0.0, 0.0],0), [batch_size,1]), tf.float64)
Rx_1 = tf.concat([tf.zeros_like(rx, tf.float64), tf.cos(rx), -tf.sin(rx)], axis=1)
Rx_2 = tf.concat([tf.zeros_like(rx, tf.float64), tf.sin(rx), tf.cos(rx)], axis=1)
Rx = tf.concat([tf.expand_dims(Rx_0,1), tf.expand_dims(Rx_1,1), tf.expand_dims(Rx_2,1)], axis=1)
Ry_0 = tf.concat([tf.cos(ry), tf.zeros_like(ry, tf.float64), tf.sin(ry)], axis=1)
Ry_1 = tf.cast(tf.tile(tf.expand_dims([0.0, 1.0, 0.0],0), [batch_size,1]), tf.float64)
Ry_2 = tf.concat([-tf.sin(ry), tf.zeros_like(ry, tf.float64), tf.cos(ry)], axis=1)
Ry = tf.concat([tf.expand_dims(Ry_0,1), tf.expand_dims(Ry_1,1), tf.expand_dims(Ry_2,1)], axis=1)
Rz_0 = tf.concat([tf.cos(rz), -tf.sin(rz), tf.zeros_like(rz, tf.float64)], axis=1)
Rz_1 = tf.concat([tf.sin(rz), tf.cos(rz), tf.zeros_like(rz, tf.float64)], axis=1)
Rz_2 = tf.cast(tf.tile(tf.expand_dims([0.0, 0.0, 1.0],0), [batch_size,1]), tf.float64)
Rz = tf.concat([tf.expand_dims(Rz_0,1), tf.expand_dims(Rz_1,1), tf.expand_dims(Rz_2,1)], axis=1)
Rot = tf.matmul(tf.matmul(Rx, Ry), Rz)
return Rot
import scipy.io
def tf_pdm2landCoord(params, global_params, select_eigen_num, use_weighted, img_size=256):
params = tf.cast(params, tf.float64)
global_params = tf.cast(global_params, tf.float64)
batch_size = params.get_shape().as_list()[0]
pdm = scipy.io.loadmat('/path_to_OpenFace/model_training/pdm_generation/Wild_data_pdm/pdm_68_aligned_wild.mat') # 41 eigen vecs
# pdm = scipy.io.loadmat('/path_to_OpenFace/model_training/pdm_generation/Wild_data_pdm/pdm_68_aligned_pipa_0.9999.mat') # 41 eigen vecs
SelectEigenNum = select_eigen_num
assert(SelectEigenNum<=41)
MeanVals = pdm['M'].T ## 1x204
EigenVals = pdm['E'][0:SelectEigenNum,:].T ## 1x30
EigenVectors = pdm['V'][:,0:SelectEigenNum].T ## 30x204
EigenVectors_weighted = pdm['V_weighted'][:,0:SelectEigenNum].T ## 30x204
# ShapeOffset = tf.matmul(params[:,0:SelectEigenNum], EigenVectors)
if use_weighted:
ShapeOffset = tf.matmul(params[:,0:SelectEigenNum], EigenVectors_weighted)
else:
ShapeOffset = tf.matmul(params[:,0:SelectEigenNum], EigenVectors)
dim = ShapeOffset.get_shape().as_list()[1]
shape3D = tf.convert_to_tensor(MeanVals) + ShapeOffset
shape3D = tf.reshape(shape3D, [batch_size, 3, dim/3])
shape3D = tf.transpose(shape3D,[0,2,1])
R = tf_Euler2Rot(tf.slice(global_params, [0,1], [-1,3]))
T = tf.concat([tf.slice(global_params, [0,4], [-1,2]), tf.zeros([batch_size,1], tf.float64)], axis=-1)
a = tf.expand_dims(tf.slice(global_params, [0,0], [-1,1]), -1)
a = tf.sigmoid(a) ## scale parameter should > 0
shape2D = tf.matmul(a * R, tf.transpose(shape3D,[0,2,1])) + tf.expand_dims(T, -1)
shape2D = tf.transpose(shape2D,[0,2,1])
shape2D = tf.slice(shape2D, [0,0,0], [-1,-1,2])
return tf.cast(shape2D, tf.float32)
#################### transfer coordinates to heatmap #####################
def tf_coord2channel_simple_rcv(RCV, keypoint_num=18, is_normalized=True ,img_H=128, img_W=64):
## MASK_RC00: if the keypoint is not detected, the RC is [0,0],
## we will set the [0,0] of all channels to -1
print('######coord2channel_simple#####')
batch_size = RCV.get_shape().as_list()[0]
RCV = tf.reshape(RCV, [batch_size, keypoint_num, 3])
R = tf.slice(RCV, [0,0,0], [-1,-1,1])
C = tf.slice(RCV, [0,0,1], [-1,-1,1])
V = tf.slice(RCV, [0,0,2], [-1,-1,1])
if is_normalized:
R = (R + 1)/2.0*img_H ## reverse norm to 256,256
R = tf.maximum(tf.zeros_like(R), R) ## keep the coords in image
R = tf.minimum(tf.ones_like(R)*img_H-1, R) ## keep the coords in image
C = (C + 1)/2.0*img_W ## reverse norm to 256,256
C = tf.maximum(tf.zeros_like(C), C) ## keep the coords in image
C = tf.minimum(tf.ones_like(C)*img_W-1, C) ## keep the coords in image
coords = tf.concat([R,C], axis=-1)
## Note: reshape starts from the last axis
coords = tf.to_int32(coords)
## coords stores x,y
R = tf.slice(coords, [0,0,0], [-1,-1,1])
R = tf.reshape(R, [-1])
C = tf.slice(coords, [0,0,1], [-1,-1,1])
C = tf.reshape(C, [-1])
batch_size = coords.get_shape().as_list()[0]
batch_idx = tf.range(0, batch_size)
batch_idx = tf.reshape(batch_idx, (batch_size, 1))
B = tf.tile(batch_idx, (1, keypoint_num))
B = tf.reshape(B, [-1])
kp_idx = tf.range(0, keypoint_num)
K = tf.tile(kp_idx, [batch_size])
indices = tf.stack([B, R, C, K], axis=1)
updates = 2*tf.ones([batch_size*keypoint_num]) ## first [0,2], then reduce to [-1,1]
shape=tf.constant([batch_size, img_H, img_W, keypoint_num])
landChannel = tf.scatter_nd(indices=indices, updates=updates, shape=shape)
V = tf.tile(V, [1,1,img_H*img_W])
V = tf.reshape(V, [batch_size, keypoint_num, img_H, img_W])
V = nchw_to_nhwc(V)
landChannel = landChannel*V
landChannel = landChannel - 1 ## first [0,2], then reduce to [-1,1]
return landChannel
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.