mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2024-12-24 23:11:12 -08:00
281 lines
10 KiB
Python
281 lines
10 KiB
Python
import os
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from numpy import linalg as npla
|
|
|
|
from facelib import FaceType, LandmarksProcessor
|
|
from core.leras import nn
|
|
|
|
"""
|
|
ported from https://github.com/1adrianb/face-alignment
|
|
"""
|
|
class FANExtractor(object):
|
|
def __init__ (self, landmarks_3D=False, place_model_on_cpu=False):
|
|
|
|
model_path = Path(__file__).parent / ( "2DFAN.npy" if not landmarks_3D else "3DFAN.npy")
|
|
if not model_path.exists():
|
|
raise Exception("Unable to load FANExtractor model")
|
|
|
|
nn.initialize(data_format="NHWC")
|
|
tf = nn.tf
|
|
|
|
class ConvBlock(nn.ModelBase):
|
|
def on_build(self, in_planes, out_planes):
|
|
self.in_planes = in_planes
|
|
self.out_planes = out_planes
|
|
|
|
self.bn1 = nn.BatchNorm2D(in_planes)
|
|
self.conv1 = nn.Conv2D (in_planes, out_planes//2, kernel_size=3, strides=1, padding='SAME', use_bias=False )
|
|
|
|
self.bn2 = nn.BatchNorm2D(out_planes//2)
|
|
self.conv2 = nn.Conv2D (out_planes//2, out_planes//4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
|
|
|
|
self.bn3 = nn.BatchNorm2D(out_planes//4)
|
|
self.conv3 = nn.Conv2D (out_planes//4, out_planes//4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
|
|
|
|
if self.in_planes != self.out_planes:
|
|
self.down_bn1 = nn.BatchNorm2D(in_planes)
|
|
self.down_conv1 = nn.Conv2D (in_planes, out_planes, kernel_size=1, strides=1, padding='VALID', use_bias=False )
|
|
else:
|
|
self.down_bn1 = None
|
|
self.down_conv1 = None
|
|
|
|
def forward(self, input):
|
|
x = input
|
|
x = self.bn1(x)
|
|
x = tf.nn.relu(x)
|
|
x = out1 = self.conv1(x)
|
|
|
|
x = self.bn2(x)
|
|
x = tf.nn.relu(x)
|
|
x = out2 = self.conv2(x)
|
|
|
|
x = self.bn3(x)
|
|
x = tf.nn.relu(x)
|
|
x = out3 = self.conv3(x)
|
|
|
|
x = tf.concat ([out1, out2, out3], axis=-1)
|
|
|
|
if self.in_planes != self.out_planes:
|
|
downsample = self.down_bn1(input)
|
|
downsample = tf.nn.relu (downsample)
|
|
downsample = self.down_conv1 (downsample)
|
|
x = x + downsample
|
|
else:
|
|
x = x + input
|
|
|
|
return x
|
|
|
|
class HourGlass (nn.ModelBase):
|
|
def on_build(self, in_planes, depth):
|
|
self.b1 = ConvBlock (in_planes, 256)
|
|
self.b2 = ConvBlock (in_planes, 256)
|
|
|
|
if depth > 1:
|
|
self.b2_plus = HourGlass(256, depth-1)
|
|
else:
|
|
self.b2_plus = ConvBlock(256, 256)
|
|
|
|
self.b3 = ConvBlock(256, 256)
|
|
|
|
def forward(self, input):
|
|
up1 = self.b1(input)
|
|
|
|
low1 = tf.nn.avg_pool(input, [1,2,2,1], [1,2,2,1], 'VALID')
|
|
low1 = self.b2 (low1)
|
|
|
|
low2 = self.b2_plus(low1)
|
|
low3 = self.b3(low2)
|
|
|
|
up2 = nn.upsample2d(low3)
|
|
|
|
return up1+up2
|
|
|
|
class FAN (nn.ModelBase):
|
|
def __init__(self):
|
|
super().__init__(name='FAN')
|
|
|
|
def on_build(self):
|
|
self.conv1 = nn.Conv2D (3, 64, kernel_size=7, strides=2, padding='SAME')
|
|
self.bn1 = nn.BatchNorm2D(64)
|
|
|
|
self.conv2 = ConvBlock(64, 128)
|
|
self.conv3 = ConvBlock(128, 128)
|
|
self.conv4 = ConvBlock(128, 256)
|
|
|
|
self.m = []
|
|
self.top_m = []
|
|
self.conv_last = []
|
|
self.bn_end = []
|
|
self.l = []
|
|
self.bl = []
|
|
self.al = []
|
|
for i in range(4):
|
|
self.m += [ HourGlass(256, 4) ]
|
|
self.top_m += [ ConvBlock(256, 256) ]
|
|
|
|
self.conv_last += [ nn.Conv2D (256, 256, kernel_size=1, strides=1, padding='VALID') ]
|
|
self.bn_end += [ nn.BatchNorm2D(256) ]
|
|
|
|
self.l += [ nn.Conv2D (256, 68, kernel_size=1, strides=1, padding='VALID') ]
|
|
|
|
if i < 4-1:
|
|
self.bl += [ nn.Conv2D (256, 256, kernel_size=1, strides=1, padding='VALID') ]
|
|
self.al += [ nn.Conv2D (68, 256, kernel_size=1, strides=1, padding='VALID') ]
|
|
|
|
def forward(self, inp) :
|
|
x, = inp
|
|
x = self.conv1(x)
|
|
x = self.bn1(x)
|
|
x = tf.nn.relu(x)
|
|
|
|
x = self.conv2(x)
|
|
x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], 'VALID')
|
|
x = self.conv3(x)
|
|
x = self.conv4(x)
|
|
|
|
outputs = []
|
|
previous = x
|
|
for i in range(4):
|
|
ll = self.m[i] (previous)
|
|
ll = self.top_m[i] (ll)
|
|
ll = self.conv_last[i] (ll)
|
|
ll = self.bn_end[i] (ll)
|
|
ll = tf.nn.relu(ll)
|
|
tmp_out = self.l[i](ll)
|
|
outputs.append(tmp_out)
|
|
if i < 4 - 1:
|
|
ll = self.bl[i](ll)
|
|
previous = previous + ll + self.al[i](tmp_out)
|
|
x = outputs[-1]
|
|
x = tf.transpose(x, (0,3,1,2) )
|
|
return x
|
|
|
|
e = None
|
|
if place_model_on_cpu:
|
|
e = tf.device("/CPU:0")
|
|
|
|
if e is not None: e.__enter__()
|
|
self.model = FAN()
|
|
self.model.load_weights(str(model_path))
|
|
if e is not None: e.__exit__(None,None,None)
|
|
|
|
self.model.build_for_run ([ ( tf.float32, (None,256,256,3) ) ])
|
|
|
|
def extract (self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
|
|
if len(rects) == 0:
|
|
return []
|
|
|
|
if is_bgr:
|
|
input_image = input_image[:,:,::-1]
|
|
is_bgr = False
|
|
|
|
(h, w, ch) = input_image.shape
|
|
|
|
landmarks = []
|
|
for (left, top, right, bottom) in rects:
|
|
scale = (right - left + bottom - top) / 195.0
|
|
|
|
center = np.array( [ (left + right) / 2.0, (top + bottom) / 2.0] )
|
|
centers = [ center ]
|
|
|
|
if multi_sample:
|
|
centers += [ center + [-1,-1],
|
|
center + [1,-1],
|
|
center + [1,1],
|
|
center + [-1,1],
|
|
]
|
|
|
|
images = []
|
|
ptss = []
|
|
|
|
try:
|
|
for c in centers:
|
|
images += [ self.crop(input_image, c, scale) ]
|
|
|
|
images = np.stack (images)
|
|
images = images.astype(np.float32) / 255.0
|
|
|
|
predicted = []
|
|
for i in range( len(images) ):
|
|
predicted += [ self.model.run ( [ images[i][None,...] ] )[0] ]
|
|
|
|
predicted = np.stack(predicted)
|
|
|
|
for i, pred in enumerate(predicted):
|
|
ptss += [ self.get_pts_from_predict ( pred, centers[i], scale) ]
|
|
pts_img = np.mean ( np.array(ptss), 0 )
|
|
|
|
landmarks.append (pts_img)
|
|
except:
|
|
landmarks.append (None)
|
|
|
|
if second_pass_extractor is not None:
|
|
for i, lmrks in enumerate(landmarks):
|
|
try:
|
|
if lmrks is not None:
|
|
image_to_face_mat = LandmarksProcessor.get_transform_mat (lmrks, 256, FaceType.FULL)
|
|
face_image = cv2.warpAffine(input_image, image_to_face_mat, (256, 256), cv2.INTER_CUBIC )
|
|
|
|
rects2 = second_pass_extractor.extract(face_image, is_bgr=is_bgr)
|
|
if len(rects2) == 1: #dont do second pass if faces != 1 detected in cropped image
|
|
lmrks2 = self.extract (face_image, [ rects2[0] ], is_bgr=is_bgr, multi_sample=True)[0]
|
|
landmarks[i] = LandmarksProcessor.transform_points (lmrks2, image_to_face_mat, True)
|
|
except:
|
|
pass
|
|
|
|
return landmarks
|
|
|
|
def transform(self, point, center, scale, resolution):
|
|
pt = np.array ( [point[0], point[1], 1.0] )
|
|
h = 200.0 * scale
|
|
m = np.eye(3)
|
|
m[0,0] = resolution / h
|
|
m[1,1] = resolution / h
|
|
m[0,2] = resolution * ( -center[0] / h + 0.5 )
|
|
m[1,2] = resolution * ( -center[1] / h + 0.5 )
|
|
m = np.linalg.inv(m)
|
|
return np.matmul (m, pt)[0:2]
|
|
|
|
def crop(self, image, center, scale, resolution=256.0):
|
|
ul = self.transform([1, 1], center, scale, resolution).astype( np.int )
|
|
br = self.transform([resolution, resolution], center, scale, resolution).astype( np.int )
|
|
|
|
if image.ndim > 2:
|
|
newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32)
|
|
newImg = np.zeros(newDim, dtype=np.uint8)
|
|
else:
|
|
newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int)
|
|
newImg = np.zeros(newDim, dtype=np.uint8)
|
|
ht = image.shape[0]
|
|
wd = image.shape[1]
|
|
newX = np.array([max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)
|
|
newY = np.array([max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)
|
|
oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
|
|
oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
|
|
newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1] ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]
|
|
|
|
newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR)
|
|
return newImg
|
|
|
|
def get_pts_from_predict(self, a, center, scale):
|
|
a_ch, a_h, a_w = a.shape
|
|
|
|
b = a.reshape ( (a_ch, a_h*a_w) )
|
|
c = b.argmax(1).reshape ( (a_ch, 1) ).repeat(2, axis=1).astype(np.float)
|
|
c[:,0] %= a_w
|
|
c[:,1] = np.apply_along_axis ( lambda x: np.floor(x / a_w), 0, c[:,1] )
|
|
|
|
for i in range(a_ch):
|
|
pX, pY = int(c[i,0]), int(c[i,1])
|
|
if pX > 0 and pX < 63 and pY > 0 and pY < 63:
|
|
diff = np.array ( [a[i,pY,pX+1]-a[i,pY,pX-1], a[i,pY+1,pX]-a[i,pY-1,pX]] )
|
|
c[i] += np.sign(diff)*0.25
|
|
|
|
c += 0.5
|
|
|
|
return np.array( [ self.transform (c[i], center, scale, a_w) for i in range(a_ch) ] )
|