mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-05 12:36:42 -07:00
66 lines
2.9 KiB
Python
66 lines
2.9 KiB
Python
import numpy as np
|
|
import os
|
|
import cv2
|
|
|
|
from pathlib import Path
|
|
|
|
from .mtcnn import *
|
|
|
|
class MTCExtractor(object):
|
|
def __init__(self, keras, tf, tf_session):
|
|
self.scale_to = 1920
|
|
self.keras = keras
|
|
self.tf = tf
|
|
self.tf_session = tf_session
|
|
|
|
self.min_face_size = self.scale_to * 0.042
|
|
self.thresh1 = 0.7
|
|
self.thresh2 = 0.85
|
|
self.thresh3 = 0.6
|
|
self.scale_factor = 0.95
|
|
|
|
'''
|
|
self.min_face_size = self.scale_to * 0.042
|
|
self.thresh1 = 7
|
|
self.thresh2 = 85
|
|
self.thresh3 = 6
|
|
self.scale_factor = 0.95
|
|
'''
|
|
|
|
def __enter__(self):
|
|
with self.tf.variable_scope('pnet2'):
|
|
data = self.tf.placeholder(self.tf.float32, (None,None,None,3), 'input')
|
|
pnet2 = PNet(self.tf, {'data':data})
|
|
pnet2.load(str(Path(__file__).parent/'det1.npy'), self.tf_session)
|
|
with self.tf.variable_scope('rnet2'):
|
|
data = self.tf.placeholder(self.tf.float32, (None,24,24,3), 'input')
|
|
rnet2 = RNet(self.tf, {'data':data})
|
|
rnet2.load(str(Path(__file__).parent/'det2.npy'), self.tf_session)
|
|
with self.tf.variable_scope('onet2'):
|
|
data = self.tf.placeholder(self.tf.float32, (None,48,48,3), 'input')
|
|
onet2 = ONet(self.tf, {'data':data})
|
|
onet2.load(str(Path(__file__).parent/'det3.npy'), self.tf_session)
|
|
|
|
self.pnet_fun = self.keras.backend.function([pnet2.layers['data']],[pnet2.layers['conv4-2'], pnet2.layers['prob1']])
|
|
self.rnet_fun = self.keras.backend.function([rnet2.layers['data']],[rnet2.layers['conv5-2'], rnet2.layers['prob1']])
|
|
self.onet_fun = self.keras.backend.function([onet2.layers['data']],[onet2.layers['conv6-2'], onet2.layers['conv6-3'], onet2.layers['prob1']])
|
|
|
|
faces, pnts = detect_face ( np.zeros ( (self.scale_to, self.scale_to, 3)), self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor )
|
|
return self
|
|
|
|
def __exit__(self, exc_type=None, exc_value=None, traceback=None):
|
|
return False #pass exception between __enter__ and __exit__ to outter level
|
|
|
|
def extract_from_bgr (self, input_image):
|
|
input_image = input_image[:,:,::-1].copy()
|
|
(h, w, ch) = input_image.shape
|
|
|
|
|
|
input_scale = self.scale_to / (w if w > h else h)
|
|
input_image = cv2.resize (input_image, ( int(w*input_scale), int(h*input_scale) ), interpolation=cv2.INTER_LINEAR)
|
|
|
|
detected_faces, pnts = detect_face ( input_image, self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor )
|
|
detected_faces = [ ( int(face[0]/input_scale), int(face[1]/input_scale), int(face[2]/input_scale), int(face[3]/input_scale)) for face in detected_faces ]
|
|
|
|
return detected_faces
|
|
|