Spaces:
Running
on
A10G
Running
on
A10G
File size: 5,333 Bytes
a22eb82 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
"""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 # feature dimension
self.num = num # number of Gaussian components
self.w = w # weights of Gaussian components (a list of scalars)
self.mu= mu # mean of Gaussian components (a list of 1xdim vectors)
self.cov = cov # covariance matrix of Gaussian components (a list of dimxdim matrices)
self.cov_det = cov_det # pre-computed determinet of covariance matrices (a list of scalars)
self.cov_inv = cov_inv # pre-computed inverse covariance matrices (a list of dimxdim matrices)
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
# calculate skin attention mask
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) # posterior probability
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]) # reshape to H*W*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))
|