Tensorflow implement of Point Distribution Model && coord2channel function in CVPR 2018 paper "Natural and effective obfuscation by head inpainting"

pdm.py

#################### 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