StarMirror/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2024-03-22 13:10:55 +08:00
Code Issues Projects Releases Wiki Activity

New script:

Browse Source 
5.XSeg) data_dst/src mask for XSeg trainer - fetch.bat
Copies faces containing XSeg polygons to aligned_xseg\ dir.
Useful only if you want to collect labeled faces and reuse them in other fakes.

Now you can use trained XSeg mask in the SAEHD training process.
It’s mean default ‘full_face’ mask obtained from landmarks will be replaced with the mask obtained from the trained XSeg model.
use
5.XSeg.optional) trained mask for data_dst/data_src - apply.bat
5.XSeg.optional) trained mask for data_dst/data_src - remove.bat

Normally you don’t need it. You can use it, if you want to use ‘face_style’ and ‘bg_style’ with obstructions.

XSeg trainer : now you can choose type of face
XSeg trainer : now you can restart training in “override settings”
Merger: XSeg-* modes now can be used with all types of faces.

Therefore old MaskEditor, FANSEG models, and FAN-x modes have been removed,
because the new XSeg solution is better, simpler and more convenient, which costs only 1 hour of manual masking for regular deepfake.
This commit is contained in:
Colombo 2020-03-30 14:00:40 +04:00
parent e5bad483ca
commit 6d3607a13d
30 changed files with 279 additions and 1520 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

67
DFLIMG/DFLJPG.py
View File

@ -7,7 +7,7 @@ import numpy as np
from core.interact import interact as io
from core.structex import *
from facelib import FaceType
from core.imagelib import SegIEPolys
class DFLJPG(object):
def __init__(self, filename):
@ -151,17 +151,6 @@ class DFLJPG(object):
print (e)
return None
@staticmethod
def embed_dfldict(filename, dfl_dict):
inst = DFLJPG.load_raw (filename)
inst.set_dict (dfl_dict)
try:
with open(filename, "wb") as f:
f.write ( inst.dump() )
except:
raise Exception( 'cannot save %s' % (filename) )
def has_data(self):
return len(self.dfl_dict.keys()) != 0
@ -176,8 +165,10 @@ class DFLJPG(object):
data = b""
dict_data = self.dfl_dict
# Remove None keys
for key in list(dict_data.keys()):
if dict_data[key] is None:
if dict_data[key] is None:
dict_data.pop(key)
for chunk in self.chunks:
@ -251,18 +242,50 @@ class DFLJPG(object):
return None
def set_image_to_face_mat(self, image_to_face_mat): self.dfl_dict['image_to_face_mat'] = image_to_face_mat
def get_ie_polys(self): return self.dfl_dict.get('ie_polys',None)
def set_ie_polys(self, ie_polys):
if ie_polys is not None and \
not isinstance(ie_polys, list):
ie_polys = ie_polys.dump()
self.dfl_dict['ie_polys'] = ie_polys
def get_seg_ie_polys(self): return self.dfl_dict.get('seg_ie_polys',None)
def get_seg_ie_polys(self):
d = self.dfl_dict.get('seg_ie_polys',None)
if d is not None:
d = SegIEPolys.load(d)
else:
d = SegIEPolys()
return d
def set_seg_ie_polys(self, seg_ie_polys):
if seg_ie_polys is not None:
if not isinstance(seg_ie_polys, SegIEPolys):
raise ValueError('seg_ie_polys should be instance of SegIEPolys')
if seg_ie_polys.has_polys():
seg_ie_polys = seg_ie_polys.dump()
else:
seg_ie_polys = None
self.dfl_dict['seg_ie_polys'] = seg_ie_polys
def get_xseg_mask(self):
mask_buf = self.dfl_dict.get('xseg_mask',None)
if mask_buf is None:
return None
img = cv2.imdecode(mask_buf, cv2.IMREAD_UNCHANGED)
if len(img.shape) == 2:
img = img[...,None]
return img.astype(np.float32) / 255.0
def set_xseg_mask(self, mask_a):
if mask_a is None:
self.dfl_dict['xseg_mask'] = None
return
ret, buf = cv2.imencode( '.png', np.clip( mask_a*255, 0, 255 ).astype(np.uint8) )
if not ret:
raise Exception("unable to generate PNG data for set_xseg_mask")
self.dfl_dict['xseg_mask'] = buf

8
XSegEditor/XSegEditor.py
View File

@ -1053,7 +1053,7 @@ class LoaderQSubprocessor(QSubprocessor):
idx, filename = data
dflimg = DFLIMG.load(filename)
if dflimg is not None and dflimg.has_data():
ie_polys = SegIEPolys.load( dflimg.get_seg_ie_polys() )
ie_polys = dflimg.get_seg_ie_polys()
return idx, True, ie_polys.has_polys()
return idx, False, False
@ -1143,7 +1143,7 @@ class MainWindow(QXMainWindow):
return False
dflimg = DFLIMG.load(image_path)
ie_polys = SegIEPolys.load( dflimg.get_seg_ie_polys() )
ie_polys = dflimg.get_seg_ie_polys()
q_img = self.load_QImage(image_path)
self.canvas.op.initialize ( q_img, ie_polys=ie_polys )
@ -1155,12 +1155,12 @@ class MainWindow(QXMainWindow):
def canvas_finalize(self, image_path):
dflimg = DFLIMG.load(image_path)
ie_polys = SegIEPolys.load( dflimg.get_seg_ie_polys() )
ie_polys = dflimg.get_seg_ie_polys()
new_ie_polys = self.canvas.op.get_ie_polys()
if not new_ie_polys.identical(ie_polys):
self.image_paths_has_ie_polys[image_path] = new_ie_polys.has_polys()
dflimg.set_seg_ie_polys( new_ie_polys.dump() )
dflimg.set_seg_ie_polys( new_ie_polys )
dflimg.save()
self.canvas.op.finalize()

109
core/imagelib/IEPolys.py
View File

@ -1,109 +0,0 @@
import numpy as np
import cv2
class IEPolysPoints:
def __init__(self, IEPolys_parent, type):
self.parent = IEPolys_parent
self.type = type
self.points = np.empty( (0,2), dtype=np.int32 )
self.n_max = self.n = 0
def add(self,x,y):
self.points = np.append(self.points[0:self.n], [ (x,y) ], axis=0)
self.n_max = self.n = self.n + 1
self.parent.dirty = True
def n_dec(self):
self.n = max(0, self.n-1)
self.parent.dirty = True
return self.n
def n_inc(self):
self.n = min(len(self.points), self.n+1)
self.parent.dirty = True
return self.n
def n_clip(self):
self.points = self.points[0:self.n]
self.n_max = self.n
def cur_point(self):
return self.points[self.n-1]
def points_to_n(self):
return self.points[0:self.n]
def set_points(self, points):
self.points = np.array(points)
self.n_max = self.n = len(points)
self.parent.dirty = True
class IEPolys:
def __init__(self):
self.list = []
self.n_max = self.n = 0
self.dirty = True
def add(self, type):
self.list = self.list[0:self.n]
l = IEPolysPoints(self, type)
self.list.append ( l )
self.n_max = self.n = self.n + 1
self.dirty = True
return l
def n_dec(self):
self.n = max(0, self.n-1)
self.dirty = True
return self.n
def n_inc(self):
self.n = min(len(self.list), self.n+1)
self.dirty = True
return self.n
def n_list(self):
return self.list[self.n-1]
def n_clip(self):
self.list = self.list[0:self.n]
self.n_max = self.n
if self.n > 0:
self.list[-1].n_clip()
def __iter__(self):
for n in range(self.n):
yield self.list[n]
def switch_dirty(self):
d = self.dirty
self.dirty = False
return d
def overlay_mask(self, mask):
h,w,c = mask.shape
white = (1,)*c
black = (0,)*c
for n in range(self.n):
poly = self.list[n]
if poly.n > 0:
cv2.fillPoly(mask, [poly.points_to_n()], white if poly.type == 1 else black )
def get_total_points(self):
return sum([self.list[n].n for n in range(self.n)])
def dump(self):
result = []
for n in range(self.n):
l = self.list[n]
result += [ (l.type, l.points_to_n().tolist() ) ]
return result
@staticmethod
def load(ie_polys=None):
obj = IEPolys()
if ie_polys is not None and isinstance(ie_polys, list):
for (type, points) in ie_polys:
obj.add(type)
obj.n_list().set_points(points)
return obj

1
core/imagelib/__init__.py
View File

@ -15,7 +15,6 @@ from .color_transfer import color_transfer, color_transfer_mix, color_transfer_s
from .common import normalize_channels, cut_odd_image, overlay_alpha_image
from .IEPolys import IEPolys
from .SegIEPolys import *
from .blursharpen import LinearMotionBlur, blursharpen

92
core/leras/models/Ternaus.py
View File

@ -1,92 +0,0 @@
"""
using https://github.com/ternaus/TernausNet
TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation
"""
from core.leras import nn
tf = nn.tf
class Ternaus(nn.ModelBase):
def on_build(self, in_ch, base_ch):
self.features_0 = nn.Conv2D (in_ch, base_ch, kernel_size=3, padding='SAME')
self.features_3 = nn.Conv2D (base_ch, base_ch*2, kernel_size=3, padding='SAME')
self.features_6 = nn.Conv2D (base_ch*2, base_ch*4, kernel_size=3, padding='SAME')
self.features_8 = nn.Conv2D (base_ch*4, base_ch*4, kernel_size=3, padding='SAME')
self.features_11 = nn.Conv2D (base_ch*4, base_ch*8, kernel_size=3, padding='SAME')
self.features_13 = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
self.features_16 = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
self.features_18 = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
self.blurpool_0 = nn.BlurPool (filt_size=3)
self.blurpool_3 = nn.BlurPool (filt_size=3)
self.blurpool_8 = nn.BlurPool (filt_size=3)
self.blurpool_13 = nn.BlurPool (filt_size=3)
self.blurpool_18 = nn.BlurPool (filt_size=3)
self.conv_center = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
self.conv1_up = nn.Conv2DTranspose (base_ch*8, base_ch*4, kernel_size=3, padding='SAME')
self.conv1 = nn.Conv2D (base_ch*12, base_ch*8, kernel_size=3, padding='SAME')
self.conv2_up = nn.Conv2DTranspose (base_ch*8, base_ch*4, kernel_size=3, padding='SAME')
self.conv2 = nn.Conv2D (base_ch*12, base_ch*8, kernel_size=3, padding='SAME')
self.conv3_up = nn.Conv2DTranspose (base_ch*8, base_ch*2, kernel_size=3, padding='SAME')
self.conv3 = nn.Conv2D (base_ch*6, base_ch*4, kernel_size=3, padding='SAME')
self.conv4_up = nn.Conv2DTranspose (base_ch*4, base_ch, kernel_size=3, padding='SAME')
self.conv4 = nn.Conv2D (base_ch*3, base_ch*2, kernel_size=3, padding='SAME')
self.conv5_up = nn.Conv2DTranspose (base_ch*2, base_ch//2, kernel_size=3, padding='SAME')
self.conv5 = nn.Conv2D (base_ch//2+base_ch, base_ch, kernel_size=3, padding='SAME')
self.out_conv = nn.Conv2D (base_ch, 1, kernel_size=3, padding='SAME')
def forward(self, inp):
x, = inp
x = x0 = tf.nn.relu(self.features_0(x))
x = self.blurpool_0(x)
x = x1 = tf.nn.relu(self.features_3(x))
x = self.blurpool_3(x)
x = tf.nn.relu(self.features_6(x))
x = x2 = tf.nn.relu(self.features_8(x))
x = self.blurpool_8(x)
x = tf.nn.relu(self.features_11(x))
x = x3 = tf.nn.relu(self.features_13(x))
x = self.blurpool_13(x)
x = tf.nn.relu(self.features_16(x))
x = x4 = tf.nn.relu(self.features_18(x))
x = self.blurpool_18(x)
x = self.conv_center(x)
x = tf.nn.relu(self.conv1_up(x))
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], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv2(x))
x = tf.nn.relu(self.conv3_up(x))
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], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv4(x))
x = tf.nn.relu(self.conv5_up(x))
x = tf.concat( [x,x0], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv5(x))
logits = self.out_conv(x)
return logits, tf.nn.sigmoid(logits)
nn.Ternaus = Ternaus

1
core/leras/models/__init__.py
View File

@ -1,5 +1,4 @@
from .ModelBase import *
from .PatchDiscriminator import *
from .CodeDiscriminator import *
from .Ternaus import *
from .XSeg import *

BIN
facelib/FANSeg_full_face_256.npy
View File

Binary file not shown.

18
facelib/LandmarksProcessor.py
View File

@ -9,7 +9,6 @@ import numpy.linalg as npla
from core import imagelib
from core import mathlib
from facelib import FaceType
from core.imagelib import IEPolys
from core.mathlib.umeyama import umeyama
landmarks_2D = np.array([
@ -374,7 +373,7 @@ def expand_eyebrows(lmrks, eyebrows_expand_mod=1.0):
def get_image_hull_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, ie_polys=None ):
def get_image_hull_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0 ):
hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32)
lmrks = expand_eyebrows(image_landmarks, eyebrows_expand_mod)
@ -393,9 +392,6 @@ def get_image_hull_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0,
merged = np.concatenate(item)
cv2.fillConvexPoly(hull_mask, cv2.convexHull(merged), (1,) )
if ie_polys is not None:
ie_polys.overlay_mask(hull_mask)
return hull_mask
def get_image_eye_mask (image_shape, image_landmarks):
@ -647,13 +643,13 @@ def mirror_landmarks (landmarks, val):
result[:,0] = val - result[:,0] - 1
return result
def get_face_struct_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, ie_polys=None, color=(1,) ):
def get_face_struct_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, color=(1,) ):
mask = np.zeros(image_shape[0:2]+( len(color),),dtype=np.float32)
lmrks = expand_eyebrows(image_landmarks, eyebrows_expand_mod)
draw_landmarks (mask, image_landmarks, color=color, draw_circles=False, thickness=2, ie_polys=ie_polys)
draw_landmarks (mask, image_landmarks, color=color, draw_circles=False, thickness=2)
return mask
def draw_landmarks (image, image_landmarks, color=(0,255,0), draw_circles=True, thickness=1, transparent_mask=False, ie_polys=None):
def draw_landmarks (image, image_landmarks, color=(0,255,0), draw_circles=True, thickness=1, transparent_mask=False):
if len(image_landmarks) != 68:
raise Exception('get_image_eye_mask works only with 68 landmarks')
@ -683,11 +679,11 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), draw_circles=True,
cv2.circle(image, (x, y), 2, color, lineType=cv2.LINE_AA)
if transparent_mask:
mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys=ie_polys)
mask = get_image_hull_mask (image.shape, image_landmarks)
image[...] = ( image * (1-mask) + image * mask / 2 )[...]
def draw_rect_landmarks (image, rect, image_landmarks, face_type, face_size=256, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0)):
draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask, ie_polys=ie_polys)
def draw_rect_landmarks (image, rect, image_landmarks, face_type, face_size=256, transparent_mask=False, landmarks_color=(0,255,0)):
draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask)
imagelib.draw_rect (image, rect, (255,0,0), 2 )
image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type)

139
facelib/TernausNet.py
View File

@ -1,139 +0,0 @@
import os
import pickle
from functools import partial
from pathlib import Path
import cv2
import numpy as np
from core.interact import interact as io
from core.leras import nn
class TernausNet(object):
VERSION = 1
def __init__ (self, name, resolution, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False, run_on_cpu=False, optimizer=None, data_format="NHWC"):
nn.initialize(data_format=data_format)
tf = nn.tf
if weights_file_root is not None:
weights_file_root = Path(weights_file_root)
else:
weights_file_root = Path(__file__).parent
self.weights_file_root = weights_file_root
with tf.device ('/CPU:0'):
#Place holders on CPU
self.input_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,3) )
self.target_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,1) )
# Initializing model classes
with tf.device ('/CPU:0' if place_model_on_cpu else '/GPU:0'):
self.net = nn.Ternaus(3, 64, name='Ternaus')
self.net_weights = self.net.get_weights()
model_name = f'{name}_{resolution}'
self.model_filename_list = [ [self.net, f'{model_name}.npy'] ]
if training:
if optimizer is None:
raise ValueError("Optimizer should be provided for traning mode.")
self.opt = optimizer
self.opt.initialize_variables (self.net_weights, vars_on_cpu=place_model_on_cpu)
self.model_filename_list += [ [self.opt, f'{model_name}_opt.npy' ] ]
else:
with tf.device ('/CPU:0' if run_on_cpu else '/GPU:0'):
_, pred = self.net([self.input_t])
def net_run(input_np):
return nn.tf_sess.run ( [pred], feed_dict={self.input_t :input_np})[0]
self.net_run = net_run
# Loading/initializing all models/optimizers weights
for model, filename in self.model_filename_list:
do_init = not load_weights
if not do_init:
do_init = not model.load_weights( self.weights_file_root / filename )
if do_init:
model.init_weights()
if model == self.net:
try:
with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
d = pickle.loads (f.read())
for i in [0,3,6,8,11,13,16,18]:
model.get_layer_by_name ('features_%d' % i).set_weights ( d['features.%d' % i] )
except:
io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy")
def save_weights(self):
for model, filename in io.progress_bar_generator(self.model_filename_list, "Saving", leave=False):
model.save_weights( self.weights_file_root / filename )
def extract (self, input_image):
input_shape_len = len(input_image.shape)
if input_shape_len == 3:
input_image = input_image[None,...]
result = np.clip ( self.net_run(input_image), 0, 1.0 )
result[result < 0.1] = 0 #get rid of noise
if input_shape_len == 3:
result = result[0]
return result
"""
if load_weights:
self.net.load_weights (self.weights_path)
else:
self.net.init_weights()
if load_weights:
self.opt.load_weights (self.opt_path)
else:
self.opt.init_weights()
"""
"""
if training:
try:
with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
d = pickle.loads (f.read())
for i in [0,3,6,8,11,13,16,18]:
s = 'features.%d' % i
self.model.get_layer (s).set_weights ( d[s] )
except:
io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy")
conv_weights_list = []
for layer in self.model.layers:
if 'CA.' in layer.name:
conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
CAInitializerMP ( conv_weights_list )
"""
"""
if training:
inp_t = Input ( (resolution, resolution, 3) )
real_t = Input ( (resolution, resolution, 1) )
out_t = self.model(inp_t)
loss = K.mean(10*K.binary_crossentropy(real_t,out_t) )
out_t_diff1 = out_t[:, 1:, :, :] - out_t[:, :-1, :, :]
out_t_diff2 = out_t[:, :, 1:, :] - out_t[:, :, :-1, :]
total_var_loss = K.mean( 0.1*K.abs(out_t_diff1), axis=[1, 2, 3] ) + K.mean( 0.1*K.abs(out_t_diff2), axis=[1, 2, 3] )
opt = Adam(lr=0.0001, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)
self.train_func = K.function ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss], self.model.trainable_weights) )
"""

24
facelib/XSegNet.py
View File

@ -14,20 +14,22 @@ class XSegNet(object):
VERSION = 1
def __init__ (self, name,
resolution,
resolution=256,
load_weights=True,
weights_file_root=None,
training=False,
place_model_on_cpu=False,
run_on_cpu=False,
optimizer=None,
data_format="NHWC"):
data_format="NHWC",
raise_on_no_model_files=False):
self.resolution = resolution
self.weights_file_root = Path(weights_file_root) if weights_file_root is not None else Path(__file__).parent
nn.initialize(data_format=data_format)
tf = nn.tf
self.weights_file_root = Path(weights_file_root) if weights_file_root is not None else Path(__file__).parent
with tf.device ('/CPU:0'):
#Place holders on CPU
self.input_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,3) )
@ -62,11 +64,17 @@ class XSegNet(object):
do_init = not load_weights
if not do_init:
do_init = not model.load_weights( self.weights_file_root / filename )
model_file_path = self.weights_file_root / filename
do_init = not model.load_weights( model_file_path )
if do_init and raise_on_no_model_files:
raise Exception(f'{model_file_path} does not exists.')
if do_init:
model.init_weights()
def get_resolution(self):
return self.resolution
def flow(self, x):
return self.model(x)
@ -78,7 +86,7 @@ class XSegNet(object):
model.save_weights( self.weights_file_root / filename )
def extract (self, input_image):
input_shape_len = len(input_image.shape)
input_shape_len = len(input_image.shape)
if input_shape_len == 3:
input_image = input_image[None,...]

1
facelib/__init__.py
View File

@ -2,5 +2,4 @@ from .FaceType import FaceType
from .S3FDExtractor import S3FDExtractor
from .FANExtractor import FANExtractor
from .FaceEnhancer import FaceEnhancer
from .TernausNet import TernausNet
from .XSegNet import XSegNet

BIN
facelib/vgg11_enc_weights.npy
View File

Binary file not shown.

48
main.py
View File

@ -224,19 +224,6 @@ if __name__ == "__main__":
p.set_defaults(func=process_videoed_video_from_sequence)
def process_labelingtool_edit_mask(arguments):
from mainscripts import MaskEditorTool
MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir, no_default_mask=arguments.no_default_mask)
labeling_parser = subparsers.add_parser( "labelingtool", help="Labeling tool.").add_subparsers()
p = labeling_parser.add_parser ( "edit_mask", help="")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
p.add_argument('--confirmed-dir', required=True, action=fixPathAction, dest="confirmed_dir", help="This is where the labeled faces will be stored.")
p.add_argument('--skipped-dir', required=True, action=fixPathAction, dest="skipped_dir", help="This is where the labeled faces will be stored.")
p.add_argument('--no-default-mask', action="store_true", dest="no_default_mask", default=False, help="Don't use default mask.")
p.set_defaults(func=process_labelingtool_edit_mask)
facesettool_parser = subparsers.add_parser( "facesettool", help="Faceset tools.").add_subparsers()
def process_faceset_enhancer(arguments):
@ -263,8 +250,10 @@ if __name__ == "__main__":
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_dev_test)
# ========== XSeg util
p = subparsers.add_parser( "xsegeditor", help="XSegEditor.")
# ========== XSeg
xseg_parser = subparsers.add_parser( "xseg", help="XSeg tools.").add_subparsers()
p = xseg_parser.add_parser( "editor", help="XSeg editor.")
def process_xsegeditor(arguments):
osex.set_process_lowest_prio()
@ -274,7 +263,36 @@ if __name__ == "__main__":
p.set_defaults (func=process_xsegeditor)
p = xseg_parser.add_parser( "apply", help="Apply trained XSeg model to the extracted faces.")
def process_xsegapply(arguments):
osex.set_process_lowest_prio()
from mainscripts import XSegUtil
XSegUtil.apply_xseg (Path(arguments.input_dir), Path(arguments.model_dir))
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir")
p.set_defaults (func=process_xsegapply)
p = xseg_parser.add_parser( "remove", help="Remove XSeg from the extracted faces.")
def process_xsegremove(arguments):
osex.set_process_lowest_prio()
from mainscripts import XSegUtil
XSegUtil.remove_xseg (Path(arguments.input_dir) )
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_xsegremove)
p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in <input_dir>_xseg dir.")
def process_xsegfetch(arguments):
osex.set_process_lowest_prio()
from mainscripts import XSegUtil
XSegUtil.fetch_xseg (Path(arguments.input_dir) )
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_xsegfetch)
def bad_args(arguments):
parser.print_help()
exit(0)

2
mainscripts/Extractor.py
View File

@ -14,7 +14,7 @@ import numpy as np
import facelib
from core import imagelib
from core import mathlib
from facelib import FaceType, LandmarksProcessor, TernausNet
from facelib import FaceType, LandmarksProcessor
from core.interact import interact as io
from core.joblib import Subprocessor
from core.leras import nn

571
mainscripts/MaskEditorTool.py
View File

@ -1,571 +0,0 @@
import os
import sys
import time
import traceback
from pathlib import Path
import cv2
import numpy as np
import numpy.linalg as npl
from core import imagelib
from DFLIMG import *
from facelib import LandmarksProcessor
from core.imagelib import IEPolys
from core.interact import interact as io
from core import pathex
from core.cv2ex import *
class MaskEditor:
STATE_NONE=0
STATE_MASKING=1
def __init__(self, img, prev_images, next_images, mask=None, ie_polys=None, get_status_lines_func=None):
self.img = imagelib.normalize_channels (img,3)
h, w, c = img.shape
if h != w and w != 256:
#to support any square res, scale img,mask and ie_polys to 256, then scale ie_polys back on .get_ie_polys()
raise Exception ("MaskEditor does not support image size != 256x256")
ph, pw = h // 4, w // 4 #pad wh
self.prev_images = prev_images
self.next_images = next_images
if mask is not None:
self.mask = imagelib.normalize_channels (mask,3)
else:
self.mask = np.zeros ( (h,w,3) )
self.get_status_lines_func = get_status_lines_func
self.state_prop = self.STATE_NONE
self.w, self.h = w, h
self.pw, self.ph = pw, ph
self.pwh = np.array([self.pw, self.ph])
self.pwh2 = np.array([self.pw*2, self.ph*2])
self.sw, self.sh = w+pw*2, h+ph*2
self.prwh = 64 #preview wh
if ie_polys is None:
ie_polys = IEPolys()
self.ie_polys = ie_polys
self.polys_mask = None
self.preview_images = None
self.mouse_x = self.mouse_y = 9999
self.screen_status_block = None
self.screen_status_block_dirty = True
self.screen_changed = True
def set_state(self, state):
self.state = state
@property
def state(self):
return self.state_prop
@state.setter
def state(self, value):
self.state_prop = value
if value == self.STATE_MASKING:
self.ie_polys.dirty = True
def get_mask(self):
if self.ie_polys.switch_dirty():
self.screen_status_block_dirty = True
self.ie_mask = img = self.mask.copy()
self.ie_polys.overlay_mask(img)
return img
return self.ie_mask
def get_screen_overlay(self):
img = np.zeros ( (self.sh, self.sw, 3) )
if self.state == self.STATE_MASKING:
mouse_xy = self.mouse_xy.copy() + self.pwh
l = self.ie_polys.n_list()
if l.n > 0:
p = l.cur_point().copy() + self.pwh
color = (0,1,0) if l.type == 1 else (0,0,1)
cv2.line(img, tuple(p), tuple(mouse_xy), color )
return img
def undo_to_begin_point(self):
while not self.undo_point():
pass
def undo_point(self):
self.screen_changed = True
if self.state == self.STATE_NONE:
if self.ie_polys.n > 0:
self.state = self.STATE_MASKING
if self.state == self.STATE_MASKING:
if self.ie_polys.n_list().n_dec() == 0 and \
self.ie_polys.n_dec() == 0:
self.state = self.STATE_NONE
else:
return False
return True
def redo_to_end_point(self):
while not self.redo_point():
pass
def redo_point(self):
self.screen_changed = True
if self.state == self.STATE_NONE:
if self.ie_polys.n_max > 0:
self.state = self.STATE_MASKING
if self.ie_polys.n == 0:
self.ie_polys.n_inc()
if self.state == self.STATE_MASKING:
while True:
l = self.ie_polys.n_list()
if l.n_inc() == l.n_max:
if self.ie_polys.n == self.ie_polys.n_max:
break
self.ie_polys.n_inc()
else:
return False
return True
def combine_screens(self, screens):
screens_len = len(screens)
new_screens = []
for screen, padded_overlay in screens:
screen_img = np.zeros( (self.sh, self.sw, 3), dtype=np.float32 )
screen = imagelib.normalize_channels (screen, 3)
h,w,c = screen.shape
screen_img[self.ph:-self.ph, self.pw:-self.pw, :] = screen
if padded_overlay is not None:
screen_img = screen_img + padded_overlay
screen_img = np.clip(screen_img*255, 0, 255).astype(np.uint8)
new_screens.append(screen_img)
return np.concatenate (new_screens, axis=1)
def get_screen_status_block(self, w, c):
if self.screen_status_block_dirty:
self.screen_status_block_dirty = False
lines = [
'Polys current/max = %d/%d' % (self.ie_polys.n, self.ie_polys.n_max),
]
if self.get_status_lines_func is not None:
lines += self.get_status_lines_func()
lines_count = len(lines)
h_line = 21
h = lines_count * h_line
img = np.ones ( (h,w,c) ) * 0.1
for i in range(lines_count):
img[ i*h_line:(i+1)*h_line, 0:w] += \
imagelib.get_text_image ( (h_line,w,c), lines[i], color=[0.8]*c )
self.screen_status_block = np.clip(img*255, 0, 255).astype(np.uint8)
return self.screen_status_block
def set_screen_status_block_dirty(self):
self.screen_status_block_dirty = True
def set_screen_changed(self):
self.screen_changed = True
def switch_screen_changed(self):
result = self.screen_changed
self.screen_changed = False
return result
def make_screen(self):
screen_overlay = self.get_screen_overlay()
final_mask = self.get_mask()
masked_img = self.img*final_mask*0.5 + self.img*(1-final_mask)
pink = np.full ( (self.h, self.w, 3), (1,0,1) )
pink_masked_img = self.img*final_mask + pink*(1-final_mask)
screens = [ (self.img, screen_overlay),
(masked_img, screen_overlay),
(pink_masked_img, screen_overlay),
]
screens = self.combine_screens(screens)
if self.preview_images is None:
sh,sw,sc = screens.shape
prh, prw = self.prwh, self.prwh
total_w = sum ([ img.shape[1] for (t,img) in self.prev_images ]) + \
sum ([ img.shape[1] for (t,img) in self.next_images ])
total_images_len = len(self.prev_images) + len(self.next_images)
max_hor_images_count = sw // prw
max_side_images_count = (max_hor_images_count - 1) // 2
prev_images = self.prev_images[-max_side_images_count:]
next_images = self.next_images[:max_side_images_count]
border = 2
max_wh_bordered = (prw-border*2, prh-border*2)
prev_images = [ (t, cv2.resize( imagelib.normalize_channels(img, 3), max_wh_bordered )) for t,img in prev_images ]
next_images = [ (t, cv2.resize( imagelib.normalize_channels(img, 3), max_wh_bordered )) for t,img in next_images ]
for images in [prev_images, next_images]:
for i, (t, img) in enumerate(images):
new_img = np.zeros ( (prh,prw, sc) )
new_img[border:-border,border:-border] = img
if t == 2:
cv2.line (new_img, ( prw//2, int(prh//1.5) ), (int(prw/1.5), prh ) , (0,1,0), thickness=2 )
cv2.line (new_img, ( int(prw/1.5), prh ), ( prw, prh // 2 ) , (0,1,0), thickness=2 )
elif t == 1:
cv2.line (new_img, ( prw//2, prh//2 ), ( prw, prh ) , (0,0,1), thickness=2 )
cv2.line (new_img, ( prw//2, prh ), ( prw, prh // 2 ) , (0,0,1), thickness=2 )
images[i] = new_img
preview_images = []
if len(prev_images) > 0:
preview_images += [ np.concatenate (prev_images, axis=1) ]
img = np.full ( (prh,prw, sc), (0,0,1), dtype=np.float )
img[border:-border,border:-border] = cv2.resize( self.img, max_wh_bordered )
preview_images += [ img ]
if len(next_images) > 0:
preview_images += [ np.concatenate (next_images, axis=1) ]
preview_images = np.concatenate ( preview_images, axis=1 )
left_pad = sw // 2 - len(prev_images) * prw - prw // 2
right_pad = sw // 2 - len(next_images) * prw - prw // 2
preview_images = np.concatenate ([np.zeros ( (preview_images.shape[0], left_pad, preview_images.shape[2]) ),
preview_images,
np.zeros ( (preview_images.shape[0], right_pad, preview_images.shape[2]) )
], axis=1)
self.preview_images = np.clip(preview_images * 255, 0, 255 ).astype(np.uint8)
status_img = self.get_screen_status_block( screens.shape[1], screens.shape[2] )
result = np.concatenate ( [self.preview_images, screens, status_img], axis=0 )
return result
def mask_finish(self, n_clip=True):
if self.state == self.STATE_MASKING:
self.screen_changed = True
if self.ie_polys.n_list().n <= 2:
self.ie_polys.n_dec()
self.state = self.STATE_NONE
if n_clip:
self.ie_polys.n_clip()
def set_mouse_pos(self,x,y):
if self.preview_images is not None:
y -= self.preview_images.shape[0]
mouse_x = x % (self.sw) - self.pw
mouse_y = y % (self.sh) - self.ph
if mouse_x != self.mouse_x or mouse_y != self.mouse_y:
self.mouse_xy = np.array( [mouse_x, mouse_y] )
self.mouse_x, self.mouse_y = self.mouse_xy
self.screen_changed = True
def mask_point(self, type):
self.screen_changed = True
if self.state == self.STATE_MASKING and \
self.ie_polys.n_list().type != type:
self.mask_finish()
elif self.state == self.STATE_NONE:
self.state = self.STATE_MASKING
self.ie_polys.add(type)
if self.state == self.STATE_MASKING:
self.ie_polys.n_list().add (self.mouse_x, self.mouse_y)
def get_ie_polys(self):
return self.ie_polys
def set_ie_polys(self, saved_ie_polys):
self.state = self.STATE_NONE
self.ie_polys = saved_ie_polys
self.redo_to_end_point()
self.mask_finish()
def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default_mask=False):
input_path = Path(input_dir)
confirmed_path = Path(confirmed_dir)
skipped_path = Path(skipped_dir)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
if not confirmed_path.exists():
confirmed_path.mkdir(parents=True)
if not skipped_path.exists():
skipped_path.mkdir(parents=True)
if not no_default_mask:
eyebrows_expand_mod = np.clip ( io.input_int ("Default eyebrows expand modifier?", 100, add_info="0..400"), 0, 400 ) / 100.0
else:
eyebrows_expand_mod = None
wnd_name = "MaskEditor tool"
io.named_window (wnd_name)
io.capture_mouse(wnd_name)
io.capture_keys(wnd_name)
cached_images = {}
image_paths = [ Path(x) for x in pathex.get_image_paths(input_path)]
done_paths = []
done_images_types = {}
image_paths_total = len(image_paths)
saved_ie_polys = IEPolys()
zoom_factor = 1.0
preview_images_count = 9
target_wh = 256
do_prev_count = 0
do_save_move_count = 0
do_save_count = 0
do_skip_move_count = 0
do_skip_count = 0
def jobs_count():
return do_prev_count + do_save_move_count + do_save_count + do_skip_move_count + do_skip_count
is_exit = False
while not is_exit:
if len(image_paths) > 0:
filepath = image_paths.pop(0)
else:
filepath = None
next_image_paths = image_paths[0:preview_images_count]
next_image_paths_names = [ path.name for path in next_image_paths ]
prev_image_paths = done_paths[-preview_images_count:]
prev_image_paths_names = [ path.name for path in prev_image_paths ]
for key in list( cached_images.keys() ):
if key not in prev_image_paths_names and \
key not in next_image_paths_names:
cached_images.pop(key)
for paths in [prev_image_paths, next_image_paths]:
for path in paths:
if path.name not in cached_images:
cached_images[path.name] = cv2_imread(str(path)) / 255.0
if filepath is not None:
dflimg = DFLIMG.load (filepath)
if dflimg is None or not dflimg.has_data():
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
else:
lmrks = dflimg.get_landmarks()
ie_polys = IEPolys.load(dflimg.get_ie_polys())
if filepath.name in cached_images:
img = cached_images[filepath.name]
else:
img = cached_images[filepath.name] = cv2_imread(str(filepath)) / 255.0
if no_default_mask:
mask = np.zeros ( (target_wh,target_wh,3) )
else:
mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks, eyebrows_expand_mod=eyebrows_expand_mod)
else:
img = np.zeros ( (target_wh,target_wh,3) )
mask = np.ones ( (target_wh,target_wh,3) )
ie_polys = None
def get_status_lines_func():
return ['Progress: %d / %d . Current file: %s' % (len(done_paths), image_paths_total, str(filepath.name) if filepath is not None else "end" ),
'[Left mouse button] - mark include mask.',
'[Right mouse button] - mark exclude mask.',
'[Middle mouse button] - finish current poly.',
'[Mouse wheel] - undo/redo poly or point. [+ctrl] - undo to begin/redo to end',
'[r] - applies edits made to last saved image.',
'[q] - prev image. [w] - skip and move to %s. [e] - save and move to %s. ' % (skipped_path.name, confirmed_path.name),
'[z] - prev image. [x] - skip. [c] - save. ',
'hold [shift] - speed up the frame counter by 10.',
'[-/+] - window zoom [esc] - quit',
]
try:
ed = MaskEditor(img,
[ (done_images_types[name], cached_images[name]) for name in prev_image_paths_names ],
[ (0, cached_images[name]) for name in next_image_paths_names ],
mask, ie_polys, get_status_lines_func)
except Exception as e:
print(e)
continue
next = False
while not next:
io.process_messages(0.005)
if jobs_count() == 0:
for (x,y,ev,flags) in io.get_mouse_events(wnd_name):
x, y = int (x / zoom_factor), int(y / zoom_factor)
ed.set_mouse_pos(x, y)
if filepath is not None:
if ev == io.EVENT_LBUTTONDOWN:
ed.mask_point(1)
elif ev == io.EVENT_RBUTTONDOWN:
ed.mask_point(0)
elif ev == io.EVENT_MBUTTONDOWN:
ed.mask_finish()
elif ev == io.EVENT_MOUSEWHEEL:
if flags & 0x80000000 != 0:
if flags & 0x8 != 0:
ed.undo_to_begin_point()
else:
ed.undo_point()
else:
if flags & 0x8 != 0:
ed.redo_to_end_point()
else:
ed.redo_point()
for key, chr_key, ctrl_pressed, alt_pressed, shift_pressed in io.get_key_events(wnd_name):
if chr_key == 'q' or chr_key == 'z':
do_prev_count = 1 if not shift_pressed else 10
elif chr_key == '-':
zoom_factor = np.clip (zoom_factor-0.1, 0.1, 4.0)
ed.set_screen_changed()
elif chr_key == '+':
zoom_factor = np.clip (zoom_factor+0.1, 0.1, 4.0)
ed.set_screen_changed()
elif key == 27: #esc
is_exit = True
next = True
break
elif filepath is not None:
if chr_key == 'e':
saved_ie_polys = ed.ie_polys
do_save_move_count = 1 if not shift_pressed else 10
elif chr_key == 'c':
saved_ie_polys = ed.ie_polys
do_save_count = 1 if not shift_pressed else 10
elif chr_key == 'w':
do_skip_move_count = 1 if not shift_pressed else 10
elif chr_key == 'x':
do_skip_count = 1 if not shift_pressed else 10
elif chr_key == 'r' and saved_ie_polys != None:
ed.set_ie_polys(saved_ie_polys)
if do_prev_count > 0:
do_prev_count -= 1
if len(done_paths) > 0:
if filepath is not None:
image_paths.insert(0, filepath)
filepath = done_paths.pop(-1)
done_images_types[filepath.name] = 0
if filepath.parent != input_path:
new_filename_path = input_path / filepath.name
filepath.rename ( new_filename_path )
image_paths.insert(0, new_filename_path)
else:
image_paths.insert(0, filepath)
next = True
elif filepath is not None:
if do_save_move_count > 0:
do_save_move_count -= 1
ed.mask_finish()
dflimg.set_ie_polys(ed.get_ie_polys())
dflimg.set_eyebrows_expand_mod(eyebrows_expand_mod)
dflimg.save()
done_paths += [ confirmed_path / filepath.name ]
done_images_types[filepath.name] = 2
filepath.rename(done_paths[-1])
next = True
elif do_save_count > 0:
do_save_count -= 1
ed.mask_finish()
dflimg.set_ie_polys(ed.get_ie_polys())
dflimg.set_eyebrows_expand_mod(eyebrows_expand_mod)
dflimg.save()
done_paths += [ filepath ]
done_images_types[filepath.name] = 2
next = True
elif do_skip_move_count > 0:
do_skip_move_count -= 1
done_paths += [ skipped_path / filepath.name ]
done_images_types[filepath.name] = 1
filepath.rename(done_paths[-1])
next = True
elif do_skip_count > 0:
do_skip_count -= 1
done_paths += [ filepath ]
done_images_types[filepath.name] = 1
next = True
else:
do_save_move_count = do_save_count = do_skip_move_count = do_skip_count = 0
if jobs_count() == 0:
if ed.switch_screen_changed():
screen = ed.make_screen()
if zoom_factor != 1.0:
h,w,c = screen.shape
screen = cv2.resize ( screen, ( int(w*zoom_factor), int(h*zoom_factor) ) )
io.show_image (wnd_name, screen )
io.process_messages(0.005)
io.destroy_all_windows()

9
mainscripts/Merger.py
View File

@ -12,7 +12,7 @@ from core.interact import interact as io
from core.joblib import MPClassFuncOnDemand, MPFunc
from core.leras import nn
from DFLIMG import DFLIMG
from facelib import FaceEnhancer, FaceType, LandmarksProcessor, TernausNet, XSegNet
from facelib import FaceEnhancer, FaceType, LandmarksProcessor, XSegNet
from merger import FrameInfo, MergerConfig, InteractiveMergerSubprocessor
def main (model_class_name=None,
@ -55,12 +55,6 @@ def main (model_class_name=None,
predictor_func = MPFunc(predictor_func)
run_on_cpu = len(nn.getCurrentDeviceConfig().devices) == 0
fanseg_full_face_256_extract_func = MPClassFuncOnDemand(TernausNet, 'extract',
name=f'FANSeg_{FaceType.toString(FaceType.FULL)}',
resolution=256,
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
xseg_256_extract_func = MPClassFuncOnDemand(XSegNet, 'extract',
name='XSeg',
resolution=256,
@ -199,7 +193,6 @@ def main (model_class_name=None,
predictor_func = predictor_func,
predictor_input_shape = predictor_input_shape,
face_enhancer_func = face_enhancer_func,
fanseg_full_face_256_extract_func = fanseg_full_face_256_extract_func,
xseg_256_extract_func = xseg_256_extract_func,
merger_config = cfg,
frames = frames,

27
mainscripts/Sorter.py
View File

@ -6,7 +6,6 @@ import sys
import tempfile
from functools import cmp_to_key
from pathlib import Path
from shutil import copyfile
import cv2
import numpy as np
@ -35,7 +34,7 @@ class BlurEstimatorSubprocessor(Subprocessor):
else:
image = cv2_imread( str(filepath) )
return [ str(filepath), estimate_sharpness(image) ]
#override
def get_data_name (self, data):
@ -146,7 +145,7 @@ def sort_by_face_pitch(input_path):
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
def sort_by_face_source_rect_size(input_path):
io.log_info ("Sorting by face rect size...")
img_list = []
@ -163,15 +162,15 @@ def sort_by_face_source_rect_size(input_path):
source_rect = dflimg.get_source_rect()
rect_area = mathlib.polygon_area(np.array(source_rect[[0,2,2,0]]).astype(np.float32), np.array(source_rect[[1,1,3,3]]).astype(np.float32))
img_list.append( [str(filepath), rect_area ] )
io.log_info ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
return img_list, trash_img_list
class HistSsimSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli):
@ -444,13 +443,13 @@ class FinalLoaderSubprocessor(Subprocessor):
raise Exception ("Unable to load %s" % (filepath.name) )
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
if self.faster:
source_rect = dflimg.get_source_rect()
sharpness = mathlib.polygon_area(np.array(source_rect[[0,2,2,0]]).astype(np.float32), np.array(source_rect[[1,1,3,3]]).astype(np.float32))
else:
sharpness = estimate_sharpness(gray)
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll ( dflimg.get_landmarks(), size=dflimg.get_shape()[1] )
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
@ -586,12 +585,12 @@ class FinalHistDissimSubprocessor(Subprocessor):
def get_result(self):
return self.result
def sort_best_faster(input_path):
def sort_best_faster(input_path):
return sort_best(input_path, faster=True)
def sort_best(input_path, faster=False):
target_count = io.input_int ("Target number of faces?", 2000)
io.log_info ("Performing sort by best faces.")
if faster:
io.log_info("Using faster algorithm. Faces will be sorted by source-rect-area instead of blur.")
@ -630,7 +629,7 @@ def sort_best(input_path, faster=False):
imgs_per_grad += total_lack // grads
sharpned_imgs_per_grad = imgs_per_grad*10
for g in io.progress_bar_generator ( range (grads), "Sort by blur"):
img_list = yaws_sample_list[g]
@ -770,7 +769,7 @@ def sort_by_absdiff(input_path):
outputs_full = []
outputs_remain = []
for i in range(batch_size):
diff_t = tf.reduce_sum( tf.abs(i_t-j_t[i]), axis=[1,2,3] )
outputs_full.append(diff_t)

42
mainscripts/Util.py
View File

@ -5,7 +5,6 @@ import cv2
from DFLIMG import *
from facelib import LandmarksProcessor, FaceType
from core.imagelib import IEPolys
from core.interact import interact as io
from core import pathex
from core.cv2ex import *
@ -100,7 +99,7 @@ def add_landmarks_debug_images(input_path):
rect = dflimg.get_source_rect()
LandmarksProcessor.draw_rect_landmarks(img, rect, face_landmarks, FaceType.FULL )
else:
LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True, ie_polys=IEPolys.load(dflimg.get_ie_polys()) )
LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True )
@ -160,42 +159,3 @@ def recover_original_aligned_filename(input_path):
fs.rename (fd)
except:
io.log_err ('fail to rename %s' % (fs.name) )
"""
def convert_png_to_jpg_file (filepath):
filepath = Path(filepath)
if filepath.suffix != '.png':
return
dflpng = DFLPNG.load (str(filepath) )
if dflpng is None:
io.log_err ("%s is not a dfl png image file" % (filepath.name) )
return
dfl_dict = dflpng.get_dict()
img = cv2_imread (str(filepath))
new_filepath = str(filepath.parent / (filepath.stem + '.jpg'))
cv2_imwrite ( new_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
DFLJPG.x( new_filepath,
face_type=dfl_dict.get('face_type', None),
landmarks=dfl_dict.get('landmarks', None),
ie_polys=dfl_dict.get('ie_polys', None),
source_filename=dfl_dict.get('source_filename', None),
source_rect=dfl_dict.get('source_rect', None),
source_landmarks=dfl_dict.get('source_landmarks', None) )
filepath.unlink()
def convert_png_to_jpg_folder (input_path):
input_path = Path(input_path)
io.log_info ("Converting PNG to JPG...\r\n")
for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Converting"):
filepath = Path(filepath)
convert_png_to_jpg_file(filepath)
"""

173
mainscripts/XSegUtil.py
View File

@ -1,109 +1,96 @@
import traceback
import json
import shutil
import traceback
from pathlib import Path
import numpy as np
from core import pathex
from core.imagelib import IEPolys
from core.cv2ex import *
from core.interact import interact as io
from core.leras import nn
from DFLIMG import *
from facelib import XSegNet
def merge(input_dir):
input_path = Path(input_dir)
def apply_xseg(input_path, model_path):
if not input_path.exists():
raise ValueError('input_dir not found. Please ensure it exists.')
raise ValueError(f'{input_path} not found. Please ensure it exists.')
if not model_path.exists():
raise ValueError(f'{model_path} not found. Please ensure it exists.')
io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.')
device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
nn.initialize(device_config)
xseg = XSegNet(name='XSeg',
load_weights=True,
weights_file_root=model_path,
data_format=nn.data_format,
raise_on_no_model_files=True)
res = xseg.get_resolution()
images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
images_processed = 0
for filepath in io.progress_bar_generator(images_paths, "Processing"):
json_filepath = filepath.parent / (filepath.stem+'.json')
if json_filepath.exists():
dflimg = DFLIMG.load(filepath)
if dflimg is not None and dflimg.has_data():
try:
json_dict = json.loads(json_filepath.read_text())
seg_ie_polys = IEPolys()
total_points = 0
#include polys first
for shape in json_dict['shapes']:
if shape['shape_type'] == 'polygon' and \
shape['label'] != '0':
seg_ie_poly = seg_ie_polys.add(1)
for x,y in shape['points']:
seg_ie_poly.add( int(x), int(y) )
total_points += 1
#exclude polys
for shape in json_dict['shapes']:
if shape['shape_type'] == 'polygon' and \
shape['label'] == '0':
seg_ie_poly = seg_ie_polys.add(0)
for x,y in shape['points']:
seg_ie_poly.add( int(x), int(y) )
total_points += 1
if total_points == 0:
io.log_info(f"No points found in {json_filepath}, skipping.")
continue
dflimg.set_seg_ie_polys ( seg_ie_polys.dump() )
dflimg.save()
json_filepath.unlink()
images_processed += 1
except:
io.log_err(f"err {filepath}, {traceback.format_exc()}")
return
io.log_info(f"Images processed: {images_processed}")
def split(input_dir ):
input_path = Path(input_dir)
if not input_path.exists():
raise ValueError('input_dir not found. Please ensure it exists.')
images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
images_processed = 0
for filepath in io.progress_bar_generator(images_paths, "Processing"):
json_filepath = filepath.parent / (filepath.stem+'.json')
dflimg = DFLIMG.load(filepath)
if dflimg is not None and dflimg.has_data():
try:
seg_ie_polys = dflimg.get_seg_ie_polys()
if seg_ie_polys is not None:
json_dict = {}
json_dict['version'] = "4.2.9"
json_dict['flags'] = {}
json_dict['shapes'] = []
json_dict['imagePath'] = filepath.name
json_dict['imageData'] = None
for poly_type, points_list in seg_ie_polys:
shape_dict = {}
shape_dict['label'] = str(poly_type)
shape_dict['points'] = points_list
shape_dict['group_id'] = None
shape_dict['shape_type'] = 'polygon'
shape_dict['flags'] = {}
json_dict['shapes'].append( shape_dict )
if dflimg is None or not dflimg.has_data():
io.log_info(f'{filepath} is not a DFLIMG')
continue
img = cv2_imread(filepath).astype(np.float32) / 255.0
h,w,c = img.shape
if w != res:
img = cv2.resize( img, (res,res), interpolation=cv2.INTER_CUBIC )
if len(img.shape) == 2:
img = img[...,None]
mask = xseg.extract(img)
mask[mask < 0.5]=0
mask[mask >= 0.5]=1
dflimg.set_xseg_mask(mask)
dflimg.save()
json_filepath.write_text( json.dumps (json_dict,indent=4) )
def remove_xseg(input_path):
if not input_path.exists():
raise ValueError(f'{input_path} not found. Please ensure it exists.')
images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
for filepath in io.progress_bar_generator(images_paths, "Processing"):
dflimg = DFLIMG.load(filepath)
if dflimg is None or not dflimg.has_data():
io.log_info(f'{filepath} is not a DFLIMG')
continue
dflimg.set_xseg_mask(None)
dflimg.save()
def fetch_xseg(input_path):
if not input_path.exists():
raise ValueError(f'{input_path} not found. Please ensure it exists.')
output_path = input_path.parent / (input_path.name + '_xseg')
output_path.mkdir(exist_ok=True, parents=True)
io.log_info(f'Copying faces containing XSeg polygons to {output_path.name}/ folder.')
images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
files_copied = 0
for filepath in io.progress_bar_generator(images_paths, "Processing"):
dflimg = DFLIMG.load(filepath)
if dflimg is None or not dflimg.has_data():
io.log_info(f'{filepath} is not a DFLIMG')
continue
ie_polys = dflimg.get_seg_ie_polys()
dflimg.set_seg_ie_polys(None)
dflimg.save()
images_processed += 1
except:
io.log_err(f"err {filepath}, {traceback.format_exc()}")
return
io.log_info(f"Images processed: {images_processed}")
if ie_polys.has_polys():
files_copied += 1
shutil.copy ( str(filepath), str(output_path / filepath.name) )
io.log_info(f'Files copied: {files_copied}')

29
mainscripts/dev_misc.py
View File

@ -8,7 +8,6 @@ import numpy as np
from core import imagelib, pathex
from core.cv2ex import *
from core.imagelib import IEPolys
from core.interact import interact as io
from core.joblib import Subprocessor
from core.leras import nn
@ -412,31 +411,3 @@ def dev_segmented_trash(input_dir):
except:
io.log_info ('fail to trashing %s' % (src.name) )
"""
#mark only
for data in extract_data:
filepath = data.filepath
output_filepath = output_path / (filepath.stem+'.jpg')
img = cv2_imread(filepath)
img = imagelib.normalize_channels(img, 3)
cv2_imwrite(output_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
json_dict = images_jsons[filepath]
ie_polys = IEPolys()
for shape in json_dict['shapes']:
ie_poly = ie_polys.add(1)
for x,y in shape['points']:
ie_poly.add( int(x), int(y) )
DFLJPG.x(output_filepath, face_type=FaceType.toString(FaceType.MARK_ONLY),
landmarks=data.landmarks[0],
ie_polys=ie_polys,
source_filename=filepath.name,
source_rect=data.rects[0],
source_landmarks=data.landmarks[0]
)
"""

6
merger/InteractiveMergerSubprocessor.py
View File

@ -66,7 +66,6 @@ class InteractiveMergerSubprocessor(Subprocessor):
self.predictor_func = client_dict['predictor_func']
self.predictor_input_shape = client_dict['predictor_input_shape']
self.face_enhancer_func = client_dict['face_enhancer_func']
self.fanseg_full_face_256_extract_func = client_dict['fanseg_full_face_256_extract_func']
self.xseg_256_extract_func = client_dict['xseg_256_extract_func']
@ -103,7 +102,6 @@ class InteractiveMergerSubprocessor(Subprocessor):
try:
final_img = MergeMasked (self.predictor_func, self.predictor_input_shape,
face_enhancer_func=self.face_enhancer_func,
fanseg_full_face_256_extract_func=self.fanseg_full_face_256_extract_func,
xseg_256_extract_func=self.xseg_256_extract_func,
cfg=cfg,
frame_info=frame_info)
@ -137,7 +135,7 @@ class InteractiveMergerSubprocessor(Subprocessor):
#override
def __init__(self, is_interactive, merger_session_filepath, predictor_func, predictor_input_shape, face_enhancer_func, fanseg_full_face_256_extract_func, xseg_256_extract_func, merger_config, frames, frames_root_path, output_path, output_mask_path, model_iter):
def __init__(self, is_interactive, merger_session_filepath, predictor_func, predictor_input_shape, face_enhancer_func, xseg_256_extract_func, merger_config, frames, frames_root_path, output_path, output_mask_path, model_iter):
if len (frames) == 0:
raise ValueError ("len (frames) == 0")
@ -151,7 +149,6 @@ class InteractiveMergerSubprocessor(Subprocessor):
self.predictor_input_shape = predictor_input_shape
self.face_enhancer_func = face_enhancer_func
self.fanseg_full_face_256_extract_func = fanseg_full_face_256_extract_func
self.xseg_256_extract_func = xseg_256_extract_func
self.frames_root_path = frames_root_path
@ -273,7 +270,6 @@ class InteractiveMergerSubprocessor(Subprocessor):
'predictor_func': self.predictor_func,
'predictor_input_shape' : self.predictor_input_shape,
'face_enhancer_func': self.face_enhancer_func,
'fanseg_full_face_256_extract_func' : self.fanseg_full_face_256_extract_func,
'xseg_256_extract_func' : self.xseg_256_extract_func,
'stdin_fd': sys.stdin.fileno() if MERGER_DEBUG else None
}

59
merger/MergeMasked.py
View File

@ -8,12 +8,10 @@ from facelib import FaceType, LandmarksProcessor
from core.interact import interact as io
from core.cv2ex import *
fanseg_input_size = 256
xseg_input_size = 256
def MergeMaskedFace (predictor_func, predictor_input_shape,
face_enhancer_func,
fanseg_full_face_256_extract_func,
xseg_256_extract_func,
cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks):
img_size = img_bgr.shape[1], img_bgr.shape[0]
@ -73,61 +71,27 @@ def MergeMaskedFace (predictor_func, predictor_input_shape,
if cfg.mask_mode == 2: #dst
prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
elif cfg.mask_mode >= 3 and cfg.mask_mode <= 7:
elif cfg.mask_mode >= 3 and cfg.mask_mode <= 6: #XSeg modes
if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6:
prd_face_fanseg_bgr = cv2.resize (prd_face_bgr, (fanseg_input_size,)*2 )
prd_face_fanseg_mask = fanseg_full_face_256_extract_func(prd_face_fanseg_bgr)
FAN_prd_face_mask_a_0 = cv2.resize ( prd_face_fanseg_mask, (output_size, output_size), cv2.INTER_CUBIC)
if cfg.mask_mode >= 4 and cfg.mask_mode <= 7:
full_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, fanseg_input_size, face_type=FaceType.FULL)
dst_face_fanseg_bgr = cv2.warpAffine(img_bgr, full_face_fanseg_mat, (fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
dst_face_fanseg_mask = fanseg_full_face_256_extract_func(dst_face_fanseg_bgr )
if cfg.face_type == FaceType.FULL:
FAN_dst_face_mask_a_0 = cv2.resize (dst_face_fanseg_mask, (output_size,output_size), cv2.INTER_CUBIC)
else:
face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, fanseg_input_size, face_type=cfg.face_type)
fanseg_rect_corner_pts = np.array ( [ [0,0], [fanseg_input_size-1,0], [0,fanseg_input_size-1] ], dtype=np.float32 )
a = LandmarksProcessor.transform_points (fanseg_rect_corner_pts, face_fanseg_mat, invert=True )
b = LandmarksProcessor.transform_points (a, full_face_fanseg_mat )
m = cv2.getAffineTransform(b, fanseg_rect_corner_pts)
FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanseg_mask, m, (fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
if cfg.mask_mode == 3: #FAN-prd
prd_face_mask_a_0 = FAN_prd_face_mask_a_0
elif cfg.mask_mode == 4: #FAN-dst
prd_face_mask_a_0 = FAN_dst_face_mask_a_0
elif cfg.mask_mode == 5:
prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
elif cfg.mask_mode == 6:
prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
elif cfg.mask_mode == 7:
prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0
elif cfg.mask_mode >= 8 and cfg.mask_mode <= 11:
if cfg.mask_mode == 8 or cfg.mask_mode == 10 or cfg.mask_mode == 11:
# obtain XSeg-prd
prd_face_xseg_bgr = cv2.resize (prd_face_bgr, (xseg_input_size,)*2, cv2.INTER_CUBIC)
prd_face_xseg_mask = xseg_256_extract_func(prd_face_xseg_bgr)
X_prd_face_mask_a_0 = cv2.resize ( prd_face_xseg_mask, (output_size, output_size), cv2.INTER_CUBIC)
if cfg.mask_mode >= 9 and cfg.mask_mode <= 11:
whole_face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, xseg_input_size, face_type=FaceType.WHOLE_FACE)
dst_face_xseg_bgr = cv2.warpAffine(img_bgr, whole_face_mat, (xseg_input_size,)*2, flags=cv2.INTER_CUBIC )
if cfg.mask_mode >= 4 and cfg.mask_mode <= 6:
# obtain XSeg-dst
xseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, xseg_input_size, face_type=cfg.face_type)
dst_face_xseg_bgr = cv2.warpAffine(img_bgr, xseg_mat, (xseg_input_size,)*2, flags=cv2.INTER_CUBIC )
dst_face_xseg_mask = xseg_256_extract_func(dst_face_xseg_bgr)
X_dst_face_mask_a_0 = cv2.resize (dst_face_xseg_mask, (output_size,output_size), cv2.INTER_CUBIC)
if cfg.mask_mode == 8: #'XSeg-prd',
if cfg.mask_mode == 3: #'XSeg-prd',
prd_face_mask_a_0 = X_prd_face_mask_a_0
elif cfg.mask_mode == 9: #'XSeg-dst',
elif cfg.mask_mode == 4: #'XSeg-dst',
prd_face_mask_a_0 = X_dst_face_mask_a_0
elif cfg.mask_mode == 10: #'XSeg-prd*XSeg-dst',
elif cfg.mask_mode == 5: #'XSeg-prd*XSeg-dst',
prd_face_mask_a_0 = X_prd_face_mask_a_0 * X_dst_face_mask_a_0
elif cfg.mask_mode == 11: #learned*XSeg-prd*XSeg-dst'
elif cfg.mask_mode == 6: #learned*XSeg-prd*XSeg-dst'
prd_face_mask_a_0 = prd_face_mask_a_0 * X_prd_face_mask_a_0 * X_dst_face_mask_a_0
prd_face_mask_a_0[ prd_face_mask_a_0 < (1.0/255.0) ] = 0.0 # get rid of noise
@ -346,7 +310,6 @@ def MergeMaskedFace (predictor_func, predictor_input_shape,
def MergeMasked (predictor_func,
predictor_input_shape,
face_enhancer_func,
fanseg_full_face_256_extract_func,
xseg_256_extract_func,
cfg,
frame_info):
@ -356,7 +319,7 @@ def MergeMasked (predictor_func,
outs = []
for face_num, img_landmarks in enumerate( frame_info.landmarks_list ):
out_img, out_img_merging_mask = MergeMaskedFace (predictor_func, predictor_input_shape, face_enhancer_func, fanseg_full_face_256_extract_func, xseg_256_extract_func, cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks)
out_img, out_img_merging_mask = MergeMaskedFace (predictor_func, predictor_input_shape, face_enhancer_func, xseg_256_extract_func, cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks)
outs += [ (out_img, out_img_merging_mask) ]
#Combining multiple face outputs

77
merger/MergerConfig.py
View File

@ -83,34 +83,14 @@ mode_str_dict = {}
for key in mode_dict.keys():
mode_str_dict[ mode_dict[key] ] = key
"""
whole_face_mask_mode_dict = {1:'learned',
2:'dst',
3:'FAN-prd',
4:'FAN-dst',
5:'FAN-prd*FAN-dst',
6:'learned*FAN-prd*FAN-dst'
}
"""
whole_face_mask_mode_dict = {1:'learned',
2:'dst',
8:'XSeg-prd',
9:'XSeg-dst',
10:'XSeg-prd*XSeg-dst',
11:'learned*XSeg-prd*XSeg-dst'
}
mask_mode_dict = {1:'learned',
2:'dst',
3:'XSeg-prd',
4:'XSeg-dst',
5:'XSeg-prd*XSeg-dst',
6:'learned*XSeg-prd*XSeg-dst'
}
full_face_mask_mode_dict = {1:'learned',
2:'dst',
3:'FAN-prd',
4:'FAN-dst',
5:'FAN-prd*FAN-dst',
6:'learned*FAN-prd*FAN-dst'}
half_face_mask_mode_dict = {1:'learned',
2:'dst',
4:'FAN-dst',
7:'learned*FAN-dst'}
ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m", 7:"sot-m", 8:"mix-m" }
ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6, "sot-m":7, "mix-m":8 }
@ -176,12 +156,7 @@ class MergerConfigMasked(MergerConfig):
self.hist_match_threshold = np.clip ( self.hist_match_threshold+diff , 0, 255)
def toggle_mask_mode(self):
if self.face_type == FaceType.WHOLE_FACE:
a = list( whole_face_mask_mode_dict.keys() )
elif self.face_type == FaceType.FULL:
a = list( full_face_mask_mode_dict.keys() )
else:
a = list( half_face_mask_mode_dict.keys() )
a = list( mask_mode_dict.keys() )
self.mask_mode = a[ (a.index(self.mask_mode)+1) % len(a) ]
def add_erode_mask_modifier(self, diff):
@ -227,26 +202,11 @@ class MergerConfigMasked(MergerConfig):
if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold", 255, add_info="0..255"), 0, 255)
if self.face_type == FaceType.WHOLE_FACE:
s = """Choose mask mode: \n"""
for key in whole_face_mask_mode_dict.keys():
s += f"""({key}) {whole_face_mask_mode_dict[key]}\n"""
io.log_info(s)
self.mask_mode = io.input_int ("", 1, valid_list=whole_face_mask_mode_dict.keys() )
elif self.face_type == FaceType.FULL:
s = """Choose mask mode: \n"""
for key in full_face_mask_mode_dict.keys():
s += f"""({key}) {full_face_mask_mode_dict[key]}\n"""
io.log_info(s)
self.mask_mode = io.input_int ("", 1, valid_list=full_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks.")
else:
s = """Choose mask mode: \n"""
for key in half_face_mask_mode_dict.keys():
s += f"""({key}) {half_face_mask_mode_dict[key]}\n"""
io.log_info(s)
self.mask_mode = io.input_int ("", 1, valid_list=half_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images.")
s = """Choose mask mode: \n"""
for key in mask_mode_dict.keys():
s += f"""({key}) {mask_mode_dict[key]}\n"""
io.log_info(s)
self.mask_mode = io.input_int ("", 1, valid_list=mask_mode_dict.keys() )
if 'raw' not in self.mode:
self.erode_mask_modifier = np.clip ( io.input_int ("Choose erode mask modifier", 0, add_info="-400..400"), -400, 400)
@ -302,14 +262,9 @@ class MergerConfigMasked(MergerConfig):
if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
r += f"""hist_match_threshold: {self.hist_match_threshold}\n"""
if self.face_type == FaceType.WHOLE_FACE:
r += f"""mask_mode: { whole_face_mask_mode_dict[self.mask_mode] }\n"""
elif self.face_type == FaceType.FULL:
r += f"""mask_mode: { full_face_mask_mode_dict[self.mask_mode] }\n"""
else:
r += f"""mask_mode: { half_face_mask_mode_dict[self.mask_mode] }\n"""
r += f"""mask_mode: { mask_mode_dict[self.mask_mode] }\n"""
if 'raw' not in self.mode:
r += (f"""erode_mask_modifier: {self.erode_mask_modifier}\n"""
f"""blur_mask_modifier: {self.blur_mask_modifier}\n"""

188
models/Model_FANSeg/Model.py
View File

@ -1,188 +0,0 @@
import multiprocessing
import operator
from functools import partial
import numpy as np
from core import mathlib
from core.interact import interact as io
from core.leras import nn
from facelib import FaceType, TernausNet
from models import ModelBase
from samplelib import *
class FANSegModel(ModelBase):
def __init__(self, *args, **kwargs):
super().__init__(*args, force_model_class_name='FANSeg', **kwargs)
#override
def on_initialize_options(self):
device_config = nn.getCurrentDeviceConfig()
yn_str = {True:'y',False:'n'}
ask_override = self.ask_override()
if self.is_first_run() or ask_override:
self.ask_autobackup_hour()
self.ask_target_iter()
self.ask_batch_size(24)
default_lr_dropout = self.options['lr_dropout'] = self.load_or_def_option('lr_dropout', False)
if self.is_first_run() or ask_override:
self.options['lr_dropout'] = io.input_bool ("Use learning rate dropout", default_lr_dropout, help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations.")
#override
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
nn.initialize(data_format="NHWC")
tf = nn.tf
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.resolution = resolution = 256
self.face_type = FaceType.FULL
place_model_on_cpu = len(devices) == 0
models_opt_device = '/CPU:0' if place_model_on_cpu else '/GPU:0'
bgr_shape = nn.get4Dshape(resolution,resolution,3)
mask_shape = nn.get4Dshape(resolution,resolution,1)
# Initializing model classes
self.model = TernausNet(f'FANSeg_{FaceType.toString(self.face_type)}',
resolution,
load_weights=not self.is_first_run(),
weights_file_root=self.get_model_root_path(),
training=True,
place_model_on_cpu=place_model_on_cpu,
optimizer=nn.RMSprop(lr=0.0001, lr_dropout=0.3 if self.options['lr_dropout'] else 1.0,name='opt') )
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_pred_list = []
gpu_losses = []
gpu_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
gpu_input_t = self.model.input_t [batch_slice,:,:,:]
gpu_target_t = self.model.target_t [batch_slice,:,:,:]
# process model tensors
gpu_pred_logits_t, gpu_pred_t = self.model.net([gpu_input_t])
gpu_pred_list.append(gpu_pred_t)
gpu_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=gpu_target_t, logits=gpu_pred_logits_t), axis=[1,2,3])
gpu_losses += [gpu_loss]
gpu_loss_gvs += [ nn.gradients ( gpu_loss, self.model.net_weights ) ]
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
pred = nn.concat(gpu_pred_list, 0)
loss = tf.reduce_mean(gpu_losses)
loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))
# Initializing training and view functions
def train(input_np, target_np):
l, _ = nn.tf_sess.run ( [loss, loss_gv_op], feed_dict={self.model.input_t :input_np, self.model.target_t :target_np })
return l
self.train = train
def view(input_np):
return nn.tf_sess.run ( [pred], feed_dict={self.model.input_t :input_np})
self.view = view
# initializing sample generators
training_data_src_path = self.training_data_src_path
training_data_dst_path = self.training_data_dst_path
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
src_generators_count = int(src_generators_count * 1.5)
src_generator = SampleGeneratorFace(training_data_src_path, random_ct_samples_path=training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE, 'ct_mode':'lct', 'warp':True, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'random_motion_blur':(25, 5), 'random_gaussian_blur':(25,5), 'data_format':nn.data_format, 'resolution': resolution},
{'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':True, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
generators_count=src_generators_count )
dst_generator = SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE, 'warp':False, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
generators_count=dst_generators_count,
raise_on_no_data=False )
if not dst_generator.is_initialized():
io.log_info(f"\nTo view the model on unseen faces, place any aligned faces in {training_data_dst_path}.\n")
self.set_training_data_generators ([src_generator, dst_generator])
#override
def get_model_filename_list(self):
return self.model.model_filename_list
#override
def onSave(self):
self.model.save_weights()
#override
def onTrainOneIter(self):
source_np, target_np = self.generate_next_samples()[0]
loss = self.train (source_np, target_np)
return ( ('loss', loss ), )
#override
def onGetPreview(self, samples):
n_samples = min(4, self.get_batch_size(), 800 // self.resolution )
src_samples, dst_samples = samples
source_np, target_np = src_samples
S, TM, SM, = [ np.clip(x, 0.0, 1.0) for x in ([source_np,target_np] + self.view (source_np) ) ]
TM, SM, = [ np.repeat (x, (3,), -1) for x in [TM, SM] ]
green_bg = np.tile( np.array([0,1,0], dtype=np.float32)[None,None,...], (self.resolution,self.resolution,1) )
result = []
st = []
for i in range(n_samples):
ar = S[i]*TM[i] + 0.5*S[i]*(1-TM[i]) + 0.5*green_bg*(1-TM[i]), SM[i], S[i]*SM[i] + green_bg*(1-SM[i])
st.append ( np.concatenate ( ar, axis=1) )
result += [ ('FANSeg training faces', np.concatenate (st, axis=0 )), ]
if len(dst_samples) != 0:
dst_np, = dst_samples
D, DM, = [ np.clip(x, 0.0, 1.0) for x in ([dst_np] + self.view (dst_np) ) ]
DM, = [ np.repeat (x, (3,), -1) for x in [DM] ]
st = []
for i in range(n_samples):
ar = D[i], DM[i], D[i]*DM[i]+ green_bg*(1-DM[i])
st.append ( np.concatenate ( ar, axis=1) )
result += [ ('FANSeg unseen faces', np.concatenate (st, axis=0 )), ]
return result
Model = FANSegModel

1
models/Model_FANSeg/__init__.py
View File

@ -1 +0,0 @@
from .Model import Model

26
models/Model_XSeg/Model.py
View File

@ -7,7 +7,7 @@ import numpy as np
from core import mathlib
from core.interact import interact as io
from core.leras import nn
from facelib import FaceType, TernausNet, XSegNet
from facelib import FaceType, XSegNet
from models import ModelBase
from samplelib import *
@ -20,6 +20,19 @@ class XSegModel(ModelBase):
def on_initialize_options(self):
self.set_batch_size(4)
ask_override = self.ask_override()
default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'wf')
if not self.is_first_run() and ask_override:
self.restart_training = io.input_bool(f"Restart training?", False, help_message="Reset model weights and start training from scratch.")
else:
self.restart_training = False
if self.is_first_run():
self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf'], help_message="Half / mid face / full face / whole face. Choose the same as your deepfake model.").lower()
#override
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
@ -31,7 +44,14 @@ class XSegModel(ModelBase):
devices = device_config.devices
self.resolution = resolution = 256
self.face_type = FaceType.WHOLE_FACE
if self.restart_training:
self.set_iter(0)
self.face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE}[ self.options['face_type'] ]
place_model_on_cpu = len(devices) == 0
models_opt_device = '/CPU:0' if place_model_on_cpu else '/GPU:0'
@ -40,7 +60,7 @@ class XSegModel(ModelBase):
mask_shape = nn.get4Dshape(resolution,resolution,1)
# Initializing model classes
self.model = XSegNet(name=f'XSeg',
self.model = XSegNet(name='XSeg',
resolution=resolution,
load_weights=not self.is_first_run(),
weights_file_root=self.get_model_root_path(),

33
samplelib/Sample.py
View File

@ -7,7 +7,7 @@ import numpy as np
from core.cv2ex import *
from DFLIMG import *
from facelib import LandmarksProcessor
from core.imagelib import IEPolys, SegIEPolys
from core.imagelib import SegIEPolys
class SampleType(IntEnum):
IMAGE = 0 #raw image
@ -26,8 +26,8 @@ class Sample(object):
'face_type',
'shape',
'landmarks',
'ie_polys',
'seg_ie_polys',
'xseg_mask',
'eyebrows_expand_mod',
'source_filename',
'person_name',
@ -40,8 +40,8 @@ class Sample(object):
face_type=None,
shape=None,
landmarks=None,
ie_polys=None,
seg_ie_polys=None,
xseg_mask=None,
eyebrows_expand_mod=None,
source_filename=None,
person_name=None,
@ -53,8 +53,13 @@ class Sample(object):
self.face_type = face_type
self.shape = shape
self.landmarks = np.array(landmarks) if landmarks is not None else None
self.ie_polys = IEPolys.load(ie_polys)
self.seg_ie_polys = SegIEPolys.load(seg_ie_polys)
if isinstance(seg_ie_polys, SegIEPolys):
self.seg_ie_polys = seg_ie_polys
else:
self.seg_ie_polys = SegIEPolys.load(seg_ie_polys)
self.xseg_mask = xseg_mask
self.eyebrows_expand_mod = eyebrows_expand_mod if eyebrows_expand_mod is not None else 1.0
self.source_filename = source_filename
self.person_name = person_name
@ -90,25 +95,9 @@ class Sample(object):
'face_type': self.face_type,
'shape': self.shape,
'landmarks': self.landmarks.tolist(),
'ie_polys': self.ie_polys.dump(),
'seg_ie_polys': self.seg_ie_polys.dump(),
'xseg_mask' : self.xseg_mask,
'eyebrows_expand_mod': self.eyebrows_expand_mod,
'source_filename': self.source_filename,
'person_name': self.person_name
}
"""
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, eyebrows_expand_mod=None, source_filename=None, fanseg_mask=None, person_name=None):
return Sample(
sample_type=sample_type if sample_type is not None else self.sample_type,
filename=filename if filename is not None else self.filename,
face_type=face_type if face_type is not None else self.face_type,
shape=shape if shape is not None else self.shape,
landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
ie_polys=ie_polys if ie_polys is not None else self.ie_polys,
pitch_yaw_roll=pitch_yaw_roll if pitch_yaw_roll is not None else self.pitch_yaw_roll,
eyebrows_expand_mod=eyebrows_expand_mod if eyebrows_expand_mod is not None else self.eyebrows_expand_mod,
source_filename=source_filename if source_filename is not None else self.source_filename,
person_name=person_name if person_name is not None else self.person_name)
"""

28
samplelib/SampleLoader.py
View File

@ -74,8 +74,8 @@ class SampleLoader:
( face_type,
shape,
landmarks,
ie_polys,
seg_ie_polys,
xseg_mask,
eyebrows_expand_mod,
source_filename,
) in result:
@ -84,35 +84,13 @@ class SampleLoader:
face_type=FaceType.fromString (face_type),
shape=shape,
landmarks=landmarks,
ie_polys=ie_polys,
seg_ie_polys=seg_ie_polys,
xseg_mask=xseg_mask,
eyebrows_expand_mod=eyebrows_expand_mod,
source_filename=source_filename,
))
return sample_list
"""
@staticmethod
def load_face_samples ( image_paths):
sample_list = []
for filename in io.progress_bar_generator (image_paths, desc="Loading"):
dflimg = DFLIMG.load (Path(filename))
if dflimg is None:
io.log_err (f"{filename} is not a dfl image file.")
else:
sample_list.append( Sample(filename=filename,
sample_type=SampleType.FACE,
face_type=FaceType.fromString ( dflimg.get_face_type() ),
shape=dflimg.get_shape(),
landmarks=dflimg.get_landmarks(),
ie_polys=dflimg.get_ie_polys(),
eyebrows_expand_mod=dflimg.get_eyebrows_expand_mod(),
source_filename=dflimg.get_source_filename(),
))
return sample_list
"""
@staticmethod
def upgradeToFaceTemporalSortedSamples( samples ):
new_s = [ (s, s.source_filename) for s in samples]
@ -178,8 +156,8 @@ class FaceSamplesLoaderSubprocessor(Subprocessor):
data = (dflimg.get_face_type(),
dflimg.get_shape(),
dflimg.get_landmarks(),
dflimg.get_ie_polys(),
dflimg.get_seg_ie_polys(),
dflimg.get_xseg_mask(),
dflimg.get_eyebrows_expand_mod(),
dflimg.get_source_filename() )

19
samplelib/SampleProcessor.py
View File

@ -56,8 +56,14 @@ class SampleProcessor(object):
ct_sample_bgr = None
h,w,c = sample_bgr.shape
def get_full_face_mask():
full_face_mask = LandmarksProcessor.get_image_hull_mask (sample_bgr.shape, sample_landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod )
def get_full_face_mask():
if sample.xseg_mask is not None:
full_face_mask = sample.xseg_mask
if full_face_mask.shape[0] != h or full_face_mask.shape[1] != w:
full_face_mask = cv2.resize(full_face_mask, (w,h), interpolation=cv2.INTER_CUBIC)
full_face_mask = imagelib.normalize_channels(full_face_mask, 1)
else:
full_face_mask = LandmarksProcessor.get_image_hull_mask (sample_bgr.shape, sample_landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod )
return np.clip(full_face_mask, 0, 1)
def get_eyes_mask():
@ -125,19 +131,18 @@ class SampleProcessor(object):
raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, face_type) )
if sample_type == SPST.FACE_MASK:
if sample_type == SPST.FACE_MASK:
if face_mask_type == SPFMT.FULL_FACE:
img = get_full_face_mask()
elif face_mask_type == SPFMT.EYES:
img = get_eyes_mask()
elif face_mask_type == SPFMT.FULL_FACE_EYES:
img = get_full_face_mask() + get_eyes_mask()
img = get_full_face_mask()
img += get_eyes_mask()*img
else:
img = np.zeros ( sample_bgr.shape[0:2]+(1,), dtype=np.float32)
if sample.ie_polys is not None:
sample.ie_polys.overlay_mask(img)
if sample_face_type == FaceType.MARK_ONLY:
mat = LandmarksProcessor.get_transform_mat (sample_landmarks, warp_resolution, face_type)

1
samplelib/__init__.py
View File

@ -10,4 +10,5 @@ from .SampleGeneratorImage import SampleGeneratorImage
from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal
from .SampleGeneratorFaceCelebAMaskHQ import SampleGeneratorFaceCelebAMaskHQ
from .SampleGeneratorFaceXSeg import SampleGeneratorFaceXSeg
from .SampleGeneratorFaceAvatarOperator import SampleGeneratorFaceAvatarOperator
from .PackedFaceset import PackedFaceset