|
"""This script is to generate skin attention mask for Deep3DFaceRecon_pytorch |
|
""" |
|
|
|
import math |
|
import numpy as np |
|
import os |
|
import cv2 |
|
|
|
class GMM: |
|
def __init__(self, dim, num, w, mu, cov, cov_det, cov_inv): |
|
self.dim = dim |
|
self.num = num |
|
self.w = w |
|
self.mu= mu |
|
self.cov = cov |
|
self.cov_det = cov_det |
|
self.cov_inv = cov_inv |
|
|
|
self.factor = [0]*num |
|
for i in range(self.num): |
|
self.factor[i] = (2*math.pi)**(self.dim/2) * self.cov_det[i]**0.5 |
|
|
|
def likelihood(self, data): |
|
assert(data.shape[1] == self.dim) |
|
N = data.shape[0] |
|
lh = np.zeros(N) |
|
|
|
for i in range(self.num): |
|
data_ = data - self.mu[i] |
|
|
|
tmp = np.matmul(data_,self.cov_inv[i]) * data_ |
|
tmp = np.sum(tmp,axis=1) |
|
power = -0.5 * tmp |
|
|
|
p = np.array([math.exp(power[j]) for j in range(N)]) |
|
p = p/self.factor[i] |
|
lh += p*self.w[i] |
|
|
|
return lh |
|
|
|
|
|
def _rgb2ycbcr(rgb): |
|
m = np.array([[65.481, 128.553, 24.966], |
|
[-37.797, -74.203, 112], |
|
[112, -93.786, -18.214]]) |
|
shape = rgb.shape |
|
rgb = rgb.reshape((shape[0] * shape[1], 3)) |
|
ycbcr = np.dot(rgb, m.transpose() / 255.) |
|
ycbcr[:, 0] += 16. |
|
ycbcr[:, 1:] += 128. |
|
return ycbcr.reshape(shape) |
|
|
|
|
|
def _bgr2ycbcr(bgr): |
|
rgb = bgr[..., ::-1] |
|
return _rgb2ycbcr(rgb) |
|
|
|
|
|
gmm_skin_w = [0.24063933, 0.16365987, 0.26034665, 0.33535415] |
|
gmm_skin_mu = [np.array([113.71862, 103.39613, 164.08226]), |
|
np.array([150.19858, 105.18467, 155.51428]), |
|
np.array([183.92976, 107.62468, 152.71820]), |
|
np.array([114.90524, 113.59782, 151.38217])] |
|
gmm_skin_cov_det = [5692842.5, 5851930.5, 2329131., 1585971.] |
|
gmm_skin_cov_inv = [np.array([[0.0019472069, 0.0020450759, -0.00060243998],[0.0020450759, 0.017700525, 0.0051420014],[-0.00060243998, 0.0051420014, 0.0081308950]]), |
|
np.array([[0.0027110141, 0.0011036990, 0.0023122299],[0.0011036990, 0.010707724, 0.010742856],[0.0023122299, 0.010742856, 0.017481629]]), |
|
np.array([[0.0048026871, 0.00022935172, 0.0077668377],[0.00022935172, 0.011729696, 0.0081661865],[0.0077668377, 0.0081661865, 0.025374353]]), |
|
np.array([[0.0011989699, 0.0022453172, -0.0010748957],[0.0022453172, 0.047758564, 0.020332102],[-0.0010748957, 0.020332102, 0.024502251]])] |
|
|
|
gmm_skin = GMM(3, 4, gmm_skin_w, gmm_skin_mu, [], gmm_skin_cov_det, gmm_skin_cov_inv) |
|
|
|
gmm_nonskin_w = [0.12791070, 0.31130761, 0.34245777, 0.21832393] |
|
gmm_nonskin_mu = [np.array([99.200851, 112.07533, 140.20602]), |
|
np.array([110.91392, 125.52969, 130.19237]), |
|
np.array([129.75864, 129.96107, 126.96808]), |
|
np.array([112.29587, 128.85121, 129.05431])] |
|
gmm_nonskin_cov_det = [458703648., 6466488., 90611376., 133097.63] |
|
gmm_nonskin_cov_inv = [np.array([[0.00085371657, 0.00071197288, 0.00023958916],[0.00071197288, 0.0025935620, 0.00076557708],[0.00023958916, 0.00076557708, 0.0015042332]]), |
|
np.array([[0.00024650150, 0.00045542428, 0.00015019422],[0.00045542428, 0.026412144, 0.018419769],[0.00015019422, 0.018419769, 0.037497383]]), |
|
np.array([[0.00037054974, 0.00038146760, 0.00040408765],[0.00038146760, 0.0085505722, 0.0079136286],[0.00040408765, 0.0079136286, 0.010982352]]), |
|
np.array([[0.00013709733, 0.00051228428, 0.00012777430],[0.00051228428, 0.28237113, 0.10528370],[0.00012777430, 0.10528370, 0.23468947]])] |
|
|
|
gmm_nonskin = GMM(3, 4, gmm_nonskin_w, gmm_nonskin_mu, [], gmm_nonskin_cov_det, gmm_nonskin_cov_inv) |
|
|
|
prior_skin = 0.8 |
|
prior_nonskin = 1 - prior_skin |
|
|
|
|
|
|
|
def skinmask(imbgr): |
|
im = _bgr2ycbcr(imbgr) |
|
|
|
data = im.reshape((-1,3)) |
|
|
|
lh_skin = gmm_skin.likelihood(data) |
|
lh_nonskin = gmm_nonskin.likelihood(data) |
|
|
|
tmp1 = prior_skin * lh_skin |
|
tmp2 = prior_nonskin * lh_nonskin |
|
post_skin = tmp1 / (tmp1+tmp2) |
|
|
|
post_skin = post_skin.reshape((im.shape[0],im.shape[1])) |
|
|
|
post_skin = np.round(post_skin*255) |
|
post_skin = post_skin.astype(np.uint8) |
|
post_skin = np.tile(np.expand_dims(post_skin,2),[1,1,3]) |
|
|
|
return post_skin |
|
|
|
|
|
def get_skin_mask(img_path): |
|
print('generating skin masks......') |
|
names = [i for i in sorted(os.listdir( |
|
img_path)) if 'jpg' in i or 'png' in i or 'jpeg' in i or 'PNG' in i] |
|
save_path = os.path.join(img_path, 'mask') |
|
if not os.path.isdir(save_path): |
|
os.makedirs(save_path) |
|
|
|
for i in range(0, len(names)): |
|
name = names[i] |
|
print('%05d' % (i), ' ', name) |
|
full_image_name = os.path.join(img_path, name) |
|
img = cv2.imread(full_image_name).astype(np.float32) |
|
skin_img = skinmask(img) |
|
cv2.imwrite(os.path.join(save_path, name), skin_img.astype(np.uint8)) |
|
|
