mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-22 06:23:20 -07:00
added FANseg extractor for src and dst faces to use it in training
This commit is contained in:
iperov
parent
1a7f108380
commit
62940aade1
11 changed files with 298 additions and 100 deletions
|
|
@ -12,6 +12,9 @@ from nnlib import nnlib
|
|||
"""
|
||||
FANSegmentator is designed to segment faces aligned by 2DFAN-4 landmarks extractor.
|
||||
|
||||
Dataset used to train located in official DFL mega.nz folder
|
||||
https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg
|
||||
|
||||
using https://github.com/ternaus/TernausNet
|
||||
TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation
|
||||
"""
|
||||
|
|
@ -33,18 +36,18 @@ class FANSegmentator(object):
|
|||
if load_weights:
|
||||
self.model.load_weights (str(self.weights_path))
|
||||
else:
|
||||
if training:
|
||||
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")
|
||||
|
||||
|
||||
if training:
|
||||
self.model.compile(loss='mse', optimizer=Adam(tf_cpu_mode=2))
|
||||
|
||||
|
|
@ -64,14 +67,14 @@ class FANSegmentator(object):
|
|||
input_shape_len = len(input_image.shape)
|
||||
if input_shape_len == 3:
|
||||
input_image = input_image[np.newaxis,...]
|
||||
|
||||
|
||||
result = np.clip ( self.model.predict( [input_image] ), 0, 1.0 )
|
||||
|
||||
|
||||
if input_shape_len == 3:
|
||||
result = result[0]
|
||||
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@staticmethod
|
||||
def BuildModel ( resolution, ngf=64, norm='', act='lrelu'):
|
||||
exec( nnlib.import_all(), locals(), globals() )
|
||||
|
|
@ -90,45 +93,45 @@ class FANSegmentator(object):
|
|||
|
||||
x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x)
|
||||
x = MaxPooling2D()(x)
|
||||
|
||||
|
||||
x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x)
|
||||
x = MaxPooling2D()(x)
|
||||
|
||||
|
||||
x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.6')(x)
|
||||
x2 = x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.8')(x)
|
||||
x = MaxPooling2D()(x)
|
||||
|
||||
|
||||
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.11')(x)
|
||||
x3 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.13')(x)
|
||||
x = MaxPooling2D()(x)
|
||||
|
||||
|
||||
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.16')(x)
|
||||
x4 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.18')(x)
|
||||
x = MaxPooling2D()(x)
|
||||
|
||||
|
||||
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same')(x)
|
||||
|
||||
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x4])
|
||||
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x3])
|
||||
|
||||
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x4])
|
||||
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x2])
|
||||
x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x1])
|
||||
x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x0])
|
||||
|
||||
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x3])
|
||||
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x2])
|
||||
x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x1])
|
||||
x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x)
|
||||
x = Concatenate(axis=3)([ x, x0])
|
||||
x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu') (x)
|
||||
|
||||
|
||||
return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid')(x)
|
||||
|
||||
|
||||
|
||||
|
||||
return func
|
||||
|
|
|
|||
22
main.py
22
main.py
|
|
@ -49,6 +49,21 @@ if __name__ == "__main__":
|
|||
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.")
|
||||
p.set_defaults (func=process_extract)
|
||||
|
||||
def process_extract_fanseg(arguments):
|
||||
os_utils.set_process_lowest_prio()
|
||||
from mainscripts import Extractor
|
||||
Extractor.extract_fanseg( arguments.input_dir,
|
||||
device_args={'cpu_only' : arguments.cpu_only,
|
||||
'multi_gpu' : arguments.multi_gpu,
|
||||
}
|
||||
)
|
||||
|
||||
p = subparsers.add_parser( "extract_fanseg", help="Extract fanseg mask from faces.")
|
||||
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
|
||||
p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
|
||||
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU.")
|
||||
p.set_defaults (func=process_extract_fanseg)
|
||||
|
||||
def process_sort(arguments):
|
||||
os_utils.set_process_lowest_prio()
|
||||
from mainscripts import Sorter
|
||||
|
|
@ -71,12 +86,17 @@ if __name__ == "__main__":
|
|||
|
||||
if arguments.recover_original_aligned_filename:
|
||||
Util.recover_original_aligned_filename (input_path=arguments.input_dir)
|
||||
|
||||
|
||||
if arguments.remove_fanseg:
|
||||
Util.remove_fanseg_folder (input_path=arguments.input_dir)
|
||||
|
||||
p = subparsers.add_parser( "util", help="Utilities.")
|
||||
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
|
||||
p.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.")
|
||||
p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.")
|
||||
p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.")
|
||||
p.add_argument('--remove-fanseg', action="store_true", dest="remove_fanseg", default=False, help="Remove fanseg mask from aligned faces.")
|
||||
|
||||
p.set_defaults (func=process_util)
|
||||
|
||||
def process_train(arguments):
|
||||
|
|
|
|||
|
|
@ -10,11 +10,13 @@ import mathlib
|
|||
import imagelib
|
||||
import cv2
|
||||
from utils import Path_utils
|
||||
from utils.DFLPNG import DFLPNG
|
||||
from utils.DFLJPG import DFLJPG
|
||||
from utils.cv2_utils import *
|
||||
import facelib
|
||||
from facelib import FaceType
|
||||
from facelib import LandmarksProcessor
|
||||
from facelib import FANSegmentator
|
||||
from nnlib import nnlib
|
||||
from joblib import Subprocessor
|
||||
from interact import interact as io
|
||||
|
|
@ -79,7 +81,12 @@ class ExtractSubprocessor(Subprocessor):
|
|||
self.second_pass_e.__enter__()
|
||||
else:
|
||||
self.second_pass_e = None
|
||||
|
||||
|
||||
elif self.type == 'fanseg':
|
||||
nnlib.import_all (device_config)
|
||||
self.e = facelib.FANSegmentator(256, FaceType.toString(FaceType.FULL) )
|
||||
self.e.__enter__()
|
||||
|
||||
elif self.type == 'final':
|
||||
pass
|
||||
|
||||
|
|
@ -124,6 +131,8 @@ class ExtractSubprocessor(Subprocessor):
|
|||
h, w, ch = image.shape
|
||||
if h == w:
|
||||
#extracting from already extracted jpg image?
|
||||
if filename_path.suffix == '.png':
|
||||
src_dflimg = DFLPNG.load ( str(filename_path) )
|
||||
if filename_path.suffix == '.jpg':
|
||||
src_dflimg = DFLJPG.load ( str(filename_path) )
|
||||
|
||||
|
|
@ -253,15 +262,22 @@ class ExtractSubprocessor(Subprocessor):
|
|||
cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
|
||||
|
||||
return data
|
||||
|
||||
|
||||
|
||||
elif self.type == 'fanseg':
|
||||
if src_dflimg is not None:
|
||||
fanseg_mask = self.e.extract( image / 255.0 )
|
||||
src_dflimg.embed_and_set( filename_path_str,
|
||||
fanseg_mask=fanseg_mask,
|
||||
#fanseg_mask_ver=FANSegmentator.VERSION,
|
||||
)
|
||||
|
||||
#overridable
|
||||
def get_data_name (self, data):
|
||||
#return string identificator of your data
|
||||
return data.filename
|
||||
|
||||
#override
|
||||
def __init__(self, input_data, type, image_size, face_type, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, final_output_path=None):
|
||||
def __init__(self, input_data, type, image_size=None, face_type=None, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, final_output_path=None):
|
||||
self.input_data = input_data
|
||||
self.type = type
|
||||
self.image_size = image_size
|
||||
|
|
@ -561,7 +577,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
if 'cpu' in backend:
|
||||
cpu_only = True
|
||||
|
||||
if 'rects' in type or type == 'landmarks':
|
||||
if 'rects' in type or type == 'landmarks' or type == 'fanseg':
|
||||
if not cpu_only and type == 'rects-mt' and backend == "plaidML": #plaidML works with MT very slowly
|
||||
cpu_only = True
|
||||
|
||||
|
|
@ -583,7 +599,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
dev_name = nnlib.device.getDeviceName(idx)
|
||||
dev_vram = nnlib.device.getDeviceVRAMTotalGb(idx)
|
||||
|
||||
if not manual and (type == 'rects-dlib' or type == 'rects-mt'):
|
||||
if not manual and (type == 'rects-dlib' or type == 'rects-mt' ):
|
||||
for i in range ( int (max (1, dev_vram / 2) ) ):
|
||||
result += [ (idx, 'GPU', '%s #%d' % (dev_name,i) , dev_vram) ]
|
||||
else:
|
||||
|
|
@ -657,8 +673,33 @@ class DeletedFilesSearcherSubprocessor(Subprocessor):
|
|||
def get_result(self):
|
||||
return self.result
|
||||
|
||||
def extract_fanseg(input_dir, device_args={} ):#ignore_extracted
|
||||
multi_gpu = device_args.get('multi_gpu', False)
|
||||
cpu_only = device_args.get('cpu_only', False)
|
||||
|
||||
input_path = Path(input_dir)
|
||||
if not input_path.exists():
|
||||
raise ValueError('Input directory not found. Please ensure it exists.')
|
||||
|
||||
paths_to_extract = []
|
||||
for filename in Path_utils.get_image_paths(input_path) :
|
||||
filepath = Path(filename)
|
||||
if filepath.suffix == '.png':
|
||||
dflimg = DFLPNG.load( str(filepath) )
|
||||
elif filepath.suffix == '.jpg':
|
||||
dflimg = DFLJPG.load ( str(filepath) )
|
||||
else:
|
||||
dflimg = None
|
||||
|
||||
if dflimg is not None:
|
||||
paths_to_extract.append (filepath)
|
||||
|
||||
paths_to_extract_len = len(paths_to_extract)
|
||||
if paths_to_extract_len > 0:
|
||||
io.log_info ("Performing extract fanseg for %d files..." % (paths_to_extract_len) )
|
||||
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in paths_to_extract ], 'fanseg', multi_gpu=multi_gpu, cpu_only=cpu_only).run()
|
||||
|
||||
|
||||
def main(input_dir,
|
||||
output_dir,
|
||||
debug_dir=None,
|
||||
|
|
|
|||
|
|
@ -397,13 +397,17 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
|
|||
else:
|
||||
lmrks = dflimg.get_landmarks()
|
||||
ie_polys = dflimg.get_ie_polys()
|
||||
fanseg_mask = dflimg.get_fanseg_mask()
|
||||
|
||||
if filepath.name in cached_images:
|
||||
img = cached_images[filepath.name]
|
||||
else:
|
||||
img = cached_images[filepath.name] = cv2_imread(str(filepath)) / 255.0
|
||||
|
||||
mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks)
|
||||
if fanseg_mask is not None:
|
||||
mask = fanseg_mask
|
||||
else:
|
||||
mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks)
|
||||
else:
|
||||
img = np.zeros ( (target_wh,target_wh,3) )
|
||||
mask = np.ones ( (target_wh,target_wh,3) )
|
||||
|
|
|
|||
|
|
@ -7,6 +7,33 @@ from utils.cv2_utils import *
|
|||
from facelib import LandmarksProcessor
|
||||
from interact import interact as io
|
||||
|
||||
def remove_fanseg_file (filepath):
|
||||
filepath = Path(filepath)
|
||||
|
||||
if filepath.suffix == '.png':
|
||||
dflimg = DFLPNG.load( str(filepath) )
|
||||
elif filepath.suffix == '.jpg':
|
||||
dflimg = DFLJPG.load ( str(filepath) )
|
||||
else:
|
||||
return
|
||||
|
||||
if dflimg is None:
|
||||
io.log_err ("%s is not a dfl image file" % (filepath.name) )
|
||||
return
|
||||
|
||||
dflimg.remove_fanseg_mask()
|
||||
dflimg.embed_and_set( str(filepath) )
|
||||
|
||||
|
||||
def remove_fanseg_folder(input_path):
|
||||
input_path = Path(input_path)
|
||||
|
||||
io.log_info ("Removing fanseg mask...\r\n")
|
||||
|
||||
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Removing"):
|
||||
filepath = Path(filepath)
|
||||
remove_fanseg_file(filepath)
|
||||
|
||||
def convert_png_to_jpg_file (filepath):
|
||||
filepath = Path(filepath)
|
||||
|
||||
|
|
|
|||
|
|
@ -385,20 +385,20 @@ class SAEModel(ModelBase):
|
|||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A = sample[0][1][0:4] #first 4 samples
|
||||
test_A_m = sample[0][2][0:4] #first 4 samples
|
||||
test_B = sample[1][1][0:4]
|
||||
test_B_m = sample[1][2][0:4]
|
||||
test_S = sample[0][1][0:4] #first 4 samples
|
||||
test_S_m = sample[0][2][0:4] #first 4 samples
|
||||
test_D = sample[1][1][0:4]
|
||||
test_D_m = sample[1][2][0:4]
|
||||
|
||||
if self.options['learn_mask']:
|
||||
S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ]
|
||||
S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ]
|
||||
DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
|
||||
else:
|
||||
S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ]
|
||||
S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ]
|
||||
|
||||
result = []
|
||||
st = []
|
||||
for i in range(0, len(test_A)):
|
||||
for i in range(0, len(test_S)):
|
||||
ar = S[i], SS[i], D[i], DD[i], SD[i]
|
||||
st.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
|
|
@ -406,8 +406,15 @@ class SAEModel(ModelBase):
|
|||
|
||||
if self.options['learn_mask']:
|
||||
st_m = []
|
||||
for i in range(0, len(test_A)):
|
||||
ar = S[i], SS[i], D[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
|
||||
for i in range(0, len(test_S)):
|
||||
ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
|
||||
st_m.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ]
|
||||
else:
|
||||
st_m = []
|
||||
for i in range(0, len(test_S)):
|
||||
ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i], SD[i]
|
||||
st_m.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ]
|
||||
|
|
|
|||
|
|
@ -1,7 +1,13 @@
|
|||
from enum import IntEnum
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from utils.cv2_utils import *
|
||||
from utils.DFLJPG import DFLJPG
|
||||
from utils.DFLPNG import DFLPNG
|
||||
|
||||
|
||||
class SampleType(IntEnum):
|
||||
IMAGE = 0 #raw image
|
||||
|
|
@ -16,7 +22,7 @@ class SampleType(IntEnum):
|
|||
QTY = 5
|
||||
|
||||
class Sample(object):
|
||||
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None):
|
||||
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None, fanseg_mask_exist=False):
|
||||
self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
|
||||
self.filename = filename
|
||||
self.face_type = face_type
|
||||
|
|
@ -27,8 +33,9 @@ class Sample(object):
|
|||
self.yaw = yaw
|
||||
self.mirror = mirror
|
||||
self.close_target_list = close_target_list
|
||||
self.fanseg_mask_exist = fanseg_mask_exist
|
||||
|
||||
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None):
|
||||
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=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,
|
||||
|
|
@ -39,7 +46,8 @@ class Sample(object):
|
|||
pitch=pitch if pitch is not None else self.pitch,
|
||||
yaw=yaw if yaw is not None else self.yaw,
|
||||
mirror=mirror if mirror is not None else self.mirror,
|
||||
close_target_list=close_target_list if close_target_list is not None else self.close_target_list)
|
||||
close_target_list=close_target_list if close_target_list is not None else self.close_target_list,
|
||||
fanseg_mask_exist=fanseg_mask_exist if fanseg_mask_exist is not None else self.fanseg_mask_exist)
|
||||
|
||||
def load_bgr(self):
|
||||
img = cv2_imread (self.filename).astype(np.float32) / 255.0
|
||||
|
|
@ -47,6 +55,19 @@ class Sample(object):
|
|||
img = img[:,::-1].copy()
|
||||
return img
|
||||
|
||||
def load_fanseg_mask(self):
|
||||
if self.fanseg_mask_exist:
|
||||
filepath = Path(self.filename)
|
||||
if filepath.suffix == '.png':
|
||||
dflimg = DFLPNG.load ( str(filepath) )
|
||||
elif filepath.suffix == '.jpg':
|
||||
dflimg = DFLJPG.load ( str(filepath) )
|
||||
else:
|
||||
dflimg = None
|
||||
return dflimg.get_fanseg_mask()
|
||||
|
||||
return None
|
||||
|
||||
def get_random_close_target_sample(self):
|
||||
if self.close_target_list is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -77,7 +77,8 @@ class SampleLoader:
|
|||
landmarks=dflimg.get_landmarks(),
|
||||
ie_polys=dflimg.get_ie_polys(),
|
||||
pitch=pitch,
|
||||
yaw=yaw) )
|
||||
yaw=yaw,
|
||||
fanseg_mask_exist=dflimg.get_fanseg_mask() is not None, ) )
|
||||
except:
|
||||
print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) )
|
||||
|
||||
|
|
|
|||
|
|
@ -13,18 +13,18 @@ class SampleProcessor(object):
|
|||
WARPED_TRANSFORMED = 0x00000004,
|
||||
TRANSFORMED = 0x00000008,
|
||||
LANDMARKS_ARRAY = 0x00000010, #currently unused
|
||||
|
||||
|
||||
RANDOM_CLOSE = 0x00000020, #currently unused
|
||||
MORPH_TO_RANDOM_CLOSE = 0x00000040, #currently unused
|
||||
|
||||
|
||||
FACE_TYPE_HALF = 0x00000100,
|
||||
FACE_TYPE_FULL = 0x00000200,
|
||||
FACE_TYPE_HEAD = 0x00000400, #currently unused
|
||||
FACE_TYPE_AVATAR = 0x00000800, #currently unused
|
||||
|
||||
|
||||
FACE_MASK_FULL = 0x00001000,
|
||||
FACE_MASK_EYES = 0x00002000, #currently unused
|
||||
|
||||
|
||||
MODE_BGR = 0x00010000, #BGR
|
||||
MODE_G = 0x00020000, #Grayscale
|
||||
MODE_GGG = 0x00040000, #3xGrayscale
|
||||
|
|
@ -35,7 +35,7 @@ class SampleProcessor(object):
|
|||
|
||||
class Options(object):
|
||||
#motion_blur = [chance_int, range] - chance 0..100 to apply to face (not mask), and range [1..3] where 3 is highest power of motion blur
|
||||
|
||||
|
||||
def __init__(self, random_flip = True, normalize_tanh = False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], motion_blur=None ):
|
||||
self.random_flip = random_flip
|
||||
self.normalize_tanh = normalize_tanh
|
||||
|
|
@ -49,11 +49,11 @@ class SampleProcessor(object):
|
|||
chance = np.clip(chance, 0, 100)
|
||||
range = [3,5,7,9][ : np.clip(range, 0, 3)+1 ]
|
||||
self.motion_blur = (chance, range)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def process (sample, sample_process_options, output_sample_types, debug):
|
||||
SPTF = SampleProcessor.TypeFlags
|
||||
|
||||
|
||||
sample_bgr = sample.load_bgr()
|
||||
h,w,c = sample_bgr.shape
|
||||
|
||||
|
|
@ -113,7 +113,7 @@ class SampleProcessor(object):
|
|||
target_face_type = FaceType.HEAD
|
||||
elif f & SPTF.FACE_TYPE_AVATAR != 0:
|
||||
target_face_type = FaceType.AVATAR
|
||||
|
||||
|
||||
apply_motion_blur = f & SPTF.OPT_APPLY_MOTION_BLUR != 0
|
||||
|
||||
if img_type == 4:
|
||||
|
|
@ -170,9 +170,16 @@ class SampleProcessor(object):
|
|||
if np.random.randint(100) < chance :
|
||||
dim = mb_range[ np.random.randint(len(mb_range) ) ]
|
||||
img = imagelib.LinearMotionBlur (img, dim, np.random.randint(180) )
|
||||
|
||||
|
||||
if face_mask_type == 1:
|
||||
mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks, cur_sample.ie_polys)
|
||||
mask = cur_sample.load_fanseg_mask() #using fanseg_mask if exist
|
||||
|
||||
if mask is None:
|
||||
mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks)
|
||||
|
||||
if cur_sample.ie_polys is not None:
|
||||
cur_sample.ie_polys.overlay_mask(mask)
|
||||
|
||||
img = np.concatenate( (img, mask ), -1 )
|
||||
elif face_mask_type == 2:
|
||||
mask = LandmarksProcessor.get_image_eye_mask (img.shape, cur_sample.landmarks)
|
||||
|
|
|
|||
|
|
@ -1,9 +1,13 @@
|
|||
import struct
|
||||
import pickle
|
||||
import struct
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from facelib import FaceType
|
||||
from imagelib import IEPolys
|
||||
from utils.struct_utils import *
|
||||
from interact import interact as io
|
||||
|
||||
class DFLJPG(object):
|
||||
def __init__(self):
|
||||
|
|
@ -137,8 +141,15 @@ class DFLJPG(object):
|
|||
if type(chunk['data']) == bytes:
|
||||
inst.dfl_dict = pickle.loads(chunk['data'])
|
||||
|
||||
if (inst.dfl_dict is not None) and ('face_type' not in inst.dfl_dict.keys()):
|
||||
inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
|
||||
if (inst.dfl_dict is not None):
|
||||
if 'face_type' not in inst.dfl_dict:
|
||||
inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
|
||||
|
||||
if 'fanseg_mask' in inst.dfl_dict:
|
||||
fanseg_mask = inst.dfl_dict['fanseg_mask']
|
||||
if fanseg_mask is not None:
|
||||
numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8)
|
||||
inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
if inst.dfl_dict == None:
|
||||
return None
|
||||
|
|
@ -155,9 +166,21 @@ class DFLJPG(object):
|
|||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None, **kwargs
|
||||
):
|
||||
|
||||
if fanseg_mask is not None:
|
||||
fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 )
|
||||
|
||||
ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] )
|
||||
|
||||
if ret and len(buf) < 60000:
|
||||
fanseg_mask = buf
|
||||
else:
|
||||
io.log_err("Unable to encode fanseg_mask for %s" % (filename) )
|
||||
fanseg_mask = None
|
||||
|
||||
inst = DFLJPG.load_raw (filename)
|
||||
inst.setDFLDictData ({
|
||||
'face_type': face_type,
|
||||
|
|
@ -166,7 +189,8 @@ class DFLJPG(object):
|
|||
'source_filename': source_filename,
|
||||
'source_rect': source_rect,
|
||||
'source_landmarks': source_landmarks,
|
||||
'image_to_face_mat': image_to_face_mat
|
||||
'image_to_face_mat': image_to_face_mat,
|
||||
'fanseg_mask' : fanseg_mask,
|
||||
})
|
||||
|
||||
try:
|
||||
|
|
@ -181,7 +205,8 @@ class DFLJPG(object):
|
|||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None, **kwargs
|
||||
):
|
||||
if face_type is None: face_type = self.get_face_type()
|
||||
if landmarks is None: landmarks = self.get_landmarks()
|
||||
|
|
@ -190,14 +215,18 @@ class DFLJPG(object):
|
|||
if source_rect is None: source_rect = self.get_source_rect()
|
||||
if source_landmarks is None: source_landmarks = self.get_source_landmarks()
|
||||
if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
|
||||
if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask()
|
||||
DFLJPG.embed_data (filename, face_type=face_type,
|
||||
landmarks=landmarks,
|
||||
ie_polys=ie_polys,
|
||||
source_filename=source_filename,
|
||||
source_rect=source_rect,
|
||||
source_landmarks=source_landmarks,
|
||||
image_to_face_mat=image_to_face_mat)
|
||||
|
||||
image_to_face_mat=image_to_face_mat,
|
||||
fanseg_mask=fanseg_mask)
|
||||
def remove_fanseg_mask(self):
|
||||
self.dfl_dict['fanseg_mask'] = None
|
||||
|
||||
def dump(self):
|
||||
data = b""
|
||||
|
||||
|
|
@ -252,8 +281,13 @@ class DFLJPG(object):
|
|||
def get_source_filename(self): return self.dfl_dict['source_filename']
|
||||
def get_source_rect(self): return self.dfl_dict['source_rect']
|
||||
def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] )
|
||||
def get_image_to_face_mat(self):
|
||||
def get_image_to_face_mat(self):
|
||||
mat = self.dfl_dict.get ('image_to_face_mat', None)
|
||||
if mat is not None:
|
||||
return np.array (mat)
|
||||
return None
|
||||
return None
|
||||
def get_fanseg_mask(self):
|
||||
fanseg_mask = self.dfl_dict.get ('fanseg_mask', None)
|
||||
if fanseg_mask is not None:
|
||||
return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis]
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -1,13 +1,16 @@
|
|||
PNG_HEADER = b"\x89PNG\r\n\x1a\n"
|
||||
|
||||
import pickle
|
||||
import string
|
||||
import struct
|
||||
import zlib
|
||||
import pickle
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from facelib import FaceType
|
||||
from imagelib import IEPolys
|
||||
|
||||
PNG_HEADER = b"\x89PNG\r\n\x1a\n"
|
||||
|
||||
class Chunk(object):
|
||||
def __init__(self, name=None, data=None):
|
||||
self.length = 0
|
||||
|
|
@ -219,7 +222,7 @@ class DFLPNG(object):
|
|||
self.data = b""
|
||||
self.length = 0
|
||||
self.chunks = []
|
||||
self.fcwp_dict = None
|
||||
self.dfl_dict = None
|
||||
|
||||
@staticmethod
|
||||
def load_raw(filename):
|
||||
|
|
@ -252,12 +255,19 @@ class DFLPNG(object):
|
|||
def load(filename):
|
||||
try:
|
||||
inst = DFLPNG.load_raw (filename)
|
||||
inst.fcwp_dict = inst.getDFLDictData()
|
||||
inst.dfl_dict = inst.getDFLDictData()
|
||||
|
||||
if (inst.fcwp_dict is not None) and ('face_type' not in inst.fcwp_dict.keys()):
|
||||
inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL)
|
||||
if inst.dfl_dict is not None:
|
||||
if 'face_type' not in inst.dfl_dict:
|
||||
inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
|
||||
|
||||
if inst.fcwp_dict == None:
|
||||
if 'fanseg_mask' in inst.dfl_dict:
|
||||
fanseg_mask = inst.dfl_dict['fanseg_mask']
|
||||
if fanseg_mask is not None:
|
||||
numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8)
|
||||
inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
if inst.dfl_dict == None:
|
||||
return None
|
||||
|
||||
return inst
|
||||
|
|
@ -272,9 +282,21 @@ class DFLPNG(object):
|
|||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None, **kwargs
|
||||
):
|
||||
|
||||
if fanseg_mask is not None:
|
||||
fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 )
|
||||
|
||||
ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] )
|
||||
|
||||
if ret and len(buf) < 60000:
|
||||
fanseg_mask = buf
|
||||
else:
|
||||
io.log_err("Unable to encode fanseg_mask for %s" % (filename) )
|
||||
fanseg_mask = None
|
||||
|
||||
inst = DFLPNG.load_raw (filename)
|
||||
inst.setDFLDictData ({
|
||||
'face_type': face_type,
|
||||
|
|
@ -283,7 +305,8 @@ class DFLPNG(object):
|
|||
'source_filename': source_filename,
|
||||
'source_rect': source_rect,
|
||||
'source_landmarks': source_landmarks,
|
||||
'image_to_face_mat':image_to_face_mat
|
||||
'image_to_face_mat':image_to_face_mat,
|
||||
'fanseg_mask' : fanseg_mask,
|
||||
})
|
||||
|
||||
try:
|
||||
|
|
@ -292,13 +315,14 @@ class DFLPNG(object):
|
|||
except:
|
||||
raise Exception( 'cannot save %s' % (filename) )
|
||||
|
||||
def embed_and_set(self, filename, face_type=None,
|
||||
landmarks=None,
|
||||
ie_polys=None,
|
||||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None
|
||||
def embed_and_set(self, filename, face_type=None,
|
||||
landmarks=None,
|
||||
ie_polys=None,
|
||||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None, **kwargs
|
||||
):
|
||||
if face_type is None: face_type = self.get_face_type()
|
||||
if landmarks is None: landmarks = self.get_landmarks()
|
||||
|
|
@ -307,13 +331,18 @@ class DFLPNG(object):
|
|||
if source_rect is None: source_rect = self.get_source_rect()
|
||||
if source_landmarks is None: source_landmarks = self.get_source_landmarks()
|
||||
if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
|
||||
if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask()
|
||||
DFLPNG.embed_data (filename, face_type=face_type,
|
||||
landmarks=landmarks,
|
||||
ie_polys=ie_polys,
|
||||
source_filename=source_filename,
|
||||
source_rect=source_rect,
|
||||
source_landmarks=source_landmarks,
|
||||
image_to_face_mat=image_to_face_mat)
|
||||
image_to_face_mat=image_to_face_mat,
|
||||
fanseg_mask=fanseg_mask)
|
||||
|
||||
def remove_fanseg_mask(self):
|
||||
self.dfl_dict['fanseg_mask'] = None
|
||||
|
||||
def dump(self):
|
||||
data = PNG_HEADER
|
||||
|
|
@ -352,17 +381,21 @@ class DFLPNG(object):
|
|||
chunk = DFLChunk(dict_data)
|
||||
self.chunks.insert(-1, chunk)
|
||||
|
||||
def get_face_type(self): return self.fcwp_dict['face_type']
|
||||
def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] )
|
||||
def get_ie_polys(self): return IEPolys.load(self.fcwp_dict.get('ie_polys',None))
|
||||
def get_source_filename(self): return self.fcwp_dict['source_filename']
|
||||
def get_source_rect(self): return self.fcwp_dict['source_rect']
|
||||
def get_source_landmarks(self): return np.array ( self.fcwp_dict['source_landmarks'] )
|
||||
def get_face_type(self): return self.dfl_dict['face_type']
|
||||
def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
|
||||
def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None))
|
||||
def get_source_filename(self): return self.dfl_dict['source_filename']
|
||||
def get_source_rect(self): return self.dfl_dict['source_rect']
|
||||
def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] )
|
||||
def get_image_to_face_mat(self):
|
||||
mat = self.fcwp_dict.get ('image_to_face_mat', None)
|
||||
mat = self.dfl_dict.get ('image_to_face_mat', None)
|
||||
if mat is not None:
|
||||
return np.array (mat)
|
||||
return None
|
||||
|
||||
def get_fanseg_mask(self):
|
||||
fanseg_mask = self.dfl_dict.get ('fanseg_mask', None)
|
||||
if fanseg_mask is not None:
|
||||
return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis]
|
||||
return None
|
||||
def __str__(self):
|
||||
return "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue