github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-05 20:42:11 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactoring

Browse source
This commit is contained in:
Colombo 2020-03-09 13:08:32 +04:00
parent 45abcff3d1
commit a030ff6951
2 changed files with 25 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

14
facelib/TernausNet.py
View file

@ -19,8 +19,8 @@ TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentat
class TernausNet(object):
VERSION = 1
def __init__ (self, name, resolution, face_type_str=None, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False):
nn.initialize(data_format="NHWC")
def __init__ (self, name, resolution, face_type_str=None, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False, data_format="NHWC"):
nn.initialize(data_format=data_format)
tf = nn.tf
class Ternaus(nn.ModelBase):
@ -87,23 +87,23 @@ class TernausNet(object):
x = self.conv_center(x)
x = tf.nn.relu(self.conv1_up(x))
x = tf.concat( [x,x4], -1)
x = tf.concat( [x,x4], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv1(x))
x = tf.nn.relu(self.conv2_up(x))
x = tf.concat( [x,x3], -1)
x = tf.concat( [x,x3], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv2(x))
x = tf.nn.relu(self.conv3_up(x))
x = tf.concat( [x,x2], -1)
x = tf.concat( [x,x2], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv3(x))
x = tf.nn.relu(self.conv4_up(x))
x = tf.concat( [x,x1], -1)
x = tf.concat( [x,x1], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv4(x))
x = tf.nn.relu(self.conv5_up(x))
x = tf.concat( [x,x0], -1)
x = tf.concat( [x,x0], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv5(x))
logits = self.out_conv(x)

27
samplelib/SampleProcessor.py
View file

@ -10,11 +10,12 @@ from facelib import FaceType, LandmarksProcessor
class SampleProcessor(object):
class SampleType(IntEnum):
NONE = 0
FACE_IMAGE = 1
FACE_MASK = 2
LANDMARKS_ARRAY = 3
PITCH_YAW_ROLL = 4
PITCH_YAW_ROLL_SIGMOID = 5
IMAGE = 1
FACE_IMAGE = 2
FACE_MASK = 3
LANDMARKS_ARRAY = 4
PITCH_YAW_ROLL = 5
PITCH_YAW_ROLL_SIGMOID = 6
class ChannelType(IntEnum):
NONE = 0
@ -92,11 +93,12 @@ class SampleProcessor(object):
ct_mode = opts.get('ct_mode', None)
data_format = opts.get('data_format', 'NHWC')
if sample_type == SPST.FACE_MASK:
if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE:
border_replicate = False
elif sample_type == SPST.FACE_IMAGE:
border_replicate = True
border_replicate = opts.get('border_replicate', border_replicate)
borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT
@ -230,9 +232,16 @@ class SampleProcessor(object):
out_sample = np.clip (out_sample * 2.0 - 1.0, -1.0, 1.0)
if data_format == "NCHW":
out_sample = np.transpose(out_sample, (2,0,1) )
#else:
# img = imagelib.warp_by_params (params, img, warp, transform, can_flip=True, border_replicate=True)
# img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
elif sample_type == SPST.IMAGE:
img = sample_bgr
img = imagelib.warp_by_params (params, img, warp, transform, can_flip=True, border_replicate=True)
img = cv2.resize( img, (resolution, resolution), cv2.INTER_CUBIC )
out_sample = img
if data_format == "NCHW":
out_sample = np.transpose(out_sample, (2,0,1) )
elif sample_type == SPST.LANDMARKS_ARRAY:
l = sample_landmarks
l = np.concatenate ( [ np.expand_dims(l[:,0] / w,-1), np.expand_dims(l[:,1] / h,-1) ], -1 )