diff --git a/DFLIMG/DFLJPG.py b/DFLIMG/DFLJPG.py index 4ef16e6..bf39663 100644 --- a/DFLIMG/DFLJPG.py +++ b/DFLIMG/DFLJPG.py @@ -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 diff --git a/XSegEditor/XSegEditor.py b/XSegEditor/XSegEditor.py index 5f66631..44ca9b8 100644 --- a/XSegEditor/XSegEditor.py +++ b/XSegEditor/XSegEditor.py @@ -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() diff --git a/core/imagelib/IEPolys.py b/core/imagelib/IEPolys.py deleted file mode 100644 index aee7333..0000000 --- a/core/imagelib/IEPolys.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/core/imagelib/__init__.py b/core/imagelib/__init__.py index 2a23e60..affdfcf 100644 --- a/core/imagelib/__init__.py +++ b/core/imagelib/__init__.py @@ -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 diff --git a/core/leras/models/Ternaus.py b/core/leras/models/Ternaus.py deleted file mode 100644 index ad5ffc3..0000000 --- a/core/leras/models/Ternaus.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/core/leras/models/__init__.py b/core/leras/models/__init__.py index 9db94fa..2f7e545 100644 --- a/core/leras/models/__init__.py +++ b/core/leras/models/__init__.py @@ -1,5 +1,4 @@ from .ModelBase import * from .PatchDiscriminator import * from .CodeDiscriminator import * -from .Ternaus import * from .XSeg import * \ No newline at end of file diff --git a/facelib/FANSeg_full_face_256.npy b/facelib/FANSeg_full_face_256.npy deleted file mode 100644 index 53a6664..0000000 Binary files a/facelib/FANSeg_full_face_256.npy and /dev/null differ diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index b8c691c..93cdd9f 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -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) diff --git a/facelib/TernausNet.py b/facelib/TernausNet.py deleted file mode 100644 index d955c95..0000000 --- a/facelib/TernausNet.py +++ /dev/null @@ -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) ) -""" diff --git a/facelib/XSegNet.py b/facelib/XSegNet.py index 35f2cef..88acab2 100644 --- a/facelib/XSegNet.py +++ b/facelib/XSegNet.py @@ -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,...] diff --git a/facelib/__init__.py b/facelib/__init__.py index e900019..e46ca51 100644 --- a/facelib/__init__.py +++ b/facelib/__init__.py @@ -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 \ No newline at end of file diff --git a/facelib/vgg11_enc_weights.npy b/facelib/vgg11_enc_weights.npy deleted file mode 100644 index ea9df4e..0000000 Binary files a/facelib/vgg11_enc_weights.npy and /dev/null differ diff --git a/main.py b/main.py index 78aafe0..e5bec0a 100644 --- a/main.py +++ b/main.py @@ -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 _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) diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 8f383d0..e101c46 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -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 diff --git a/mainscripts/MaskEditorTool.py b/mainscripts/MaskEditorTool.py deleted file mode 100644 index 04181fe..0000000 --- a/mainscripts/MaskEditorTool.py +++ /dev/null @@ -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() diff --git a/mainscripts/Merger.py b/mainscripts/Merger.py index 3004c56..1739c9f 100644 --- a/mainscripts/Merger.py +++ b/mainscripts/Merger.py @@ -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, diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index 935830a..f04409a 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -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) diff --git a/mainscripts/Util.py b/mainscripts/Util.py index 86882cc..ba822ce 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -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) -""" \ No newline at end of file diff --git a/mainscripts/XSegUtil.py b/mainscripts/XSegUtil.py index 5f1b3f2..e557e0b 100644 --- a/mainscripts/XSegUtil.py +++ b/mainscripts/XSegUtil.py @@ -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}") \ No newline at end of file + 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}') \ No newline at end of file diff --git a/mainscripts/dev_misc.py b/mainscripts/dev_misc.py index e18bad3..9f9c1fa 100644 --- a/mainscripts/dev_misc.py +++ b/mainscripts/dev_misc.py @@ -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] - ) -""" \ No newline at end of file diff --git a/merger/InteractiveMergerSubprocessor.py b/merger/InteractiveMergerSubprocessor.py index bf92045..5552bce 100644 --- a/merger/InteractiveMergerSubprocessor.py +++ b/merger/InteractiveMergerSubprocessor.py @@ -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 } diff --git a/merger/MergeMasked.py b/merger/MergeMasked.py index c3ee34c..023718d 100644 --- a/merger/MergeMasked.py +++ b/merger/MergeMasked.py @@ -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 diff --git a/merger/MergerConfig.py b/merger/MergerConfig.py index c385769..c9eb868 100644 --- a/merger/MergerConfig.py +++ b/merger/MergerConfig.py @@ -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""" diff --git a/models/Model_FANSeg/Model.py b/models/Model_FANSeg/Model.py deleted file mode 100644 index d41adbd..0000000 --- a/models/Model_FANSeg/Model.py +++ /dev/null @@ -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 diff --git a/models/Model_FANSeg/__init__.py b/models/Model_FANSeg/__init__.py deleted file mode 100644 index 0188f11..0000000 --- a/models/Model_FANSeg/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .Model import Model diff --git a/models/Model_XSeg/Model.py b/models/Model_XSeg/Model.py index b3578fb..6f9bb2f 100644 --- a/models/Model_XSeg/Model.py +++ b/models/Model_XSeg/Model.py @@ -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(), diff --git a/samplelib/Sample.py b/samplelib/Sample.py index f4165b9..058a4a1 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -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) - -""" \ No newline at end of file diff --git a/samplelib/SampleLoader.py b/samplelib/SampleLoader.py index ff4ad15..ce175d0 100644 --- a/samplelib/SampleLoader.py +++ b/samplelib/SampleLoader.py @@ -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() ) diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index f1f961b..effcd97 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -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) diff --git a/samplelib/__init__.py b/samplelib/__init__.py index 9c140ff..87f7e16 100644 --- a/samplelib/__init__.py +++ b/samplelib/__init__.py @@ -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 \ No newline at end of file