github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

SAE : WARNING, RETRAIN IS REQUIRED !

Browse source 
fixed model sizes from previous update.
avoided bug in ML framework(keras) that forces to train the model on random noise.

Converter: added blur on the same keys as sharpness

Added new model 'TrueFace'. This is a GAN model ported from https://github.com/NVlabs/FUNIT
Model produces near zero morphing and high detail face.
Model has higher failure rate than other models.
Keep src and dst faceset in same lighting conditions.
This commit is contained in:
Colombo 2019-09-19 11:13:56 +04:00
parent 201b762541
commit dc11ec32be
26 changed files with 1308 additions and 250 deletions
Download patch file Download diff file Expand all files Collapse all files

4
converters/ConvertMasked.py
View file

@ -302,8 +302,8 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i
k_size *= 2
out_face_bgr = imagelib.LinearMotionBlur (out_face_bgr, k_size , frame_info.motion_deg)
if cfg.sharpen_mode != 0 and cfg.sharpen_amount != 0:
out_face_bgr = cfg.sharpen_func ( out_face_bgr, cfg.sharpen_mode, 3, cfg.sharpen_amount)
if cfg.blursharpen_amount != 0:
out_face_bgr = cfg.blursharpen_func ( out_face_bgr, cfg.sharpen_mode, 3, cfg.blursharpen_amount)
new_out = cv2.warpAffine( out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT )
out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (new_out*img_face_mask_aaa) , 0, 1.0 )

15
converters/ConverterConfig.py
View file

@ -18,7 +18,7 @@ class ConverterConfig(object):
self.type = type
self.superres_func = None
self.sharpen_func = None
self.blursharpen_func = None
self.fanseg_input_size = None
self.fanseg_extract_func = None
self.ebs_ct_func = None
@ -29,7 +29,7 @@ class ConverterConfig(object):
#default changeable params
self.super_resolution_mode = 0
self.sharpen_mode = 0
self.sharpen_amount = 0
self.blursharpen_amount = 0
def copy(self):
return copy.copy(self)
@ -43,7 +43,7 @@ class ConverterConfig(object):
self.sharpen_mode = io.input_int (s, 0, valid_list=self.sharpen_dict.keys(), help_message="Enhance details by applying sharpen filter.")
if self.sharpen_mode != 0:
self.sharpen_amount = np.clip ( io.input_int ("Choose sharpen amount [0..100] (skip:%d) : " % 10, 10), 0, 100 )
self.blursharpen_amount = np.clip ( io.input_int ("Choose blur/sharpen amount [-100..100] (skip:0) : ", 0), -100, 100 )
s = """Choose super resolution mode: \n"""
for key in self.super_res_dict.keys():
@ -55,8 +55,8 @@ class ConverterConfig(object):
a = list( self.sharpen_dict.keys() )
self.sharpen_mode = a[ (a.index(self.sharpen_mode)+1) % len(a) ]
def add_sharpen_amount(self, diff):
self.sharpen_amount = np.clip ( self.sharpen_amount+diff, 0, 100)
def add_blursharpen_amount(self, diff):
self.blursharpen_amount = np.clip ( self.blursharpen_amount+diff, -100, 100)
def toggle_super_resolution_mode(self):
a = list( self.super_res_dict.keys() )
@ -68,7 +68,7 @@ class ConverterConfig(object):
if isinstance(other, ConverterConfig):
return self.sharpen_mode == other.sharpen_mode and \
(self.sharpen_mode == 0 or ((self.sharpen_mode == other.sharpen_mode) and (self.sharpen_amount == other.sharpen_amount) )) and \
self.blursharpen_amount == other.blursharpen_amount and \
self.super_resolution_mode == other.super_resolution_mode
return False
@ -77,8 +77,7 @@ class ConverterConfig(object):
def to_string(self, filename):
r = ""
r += f"sharpen_mode : {self.sharpen_dict[self.sharpen_mode]}\n"
if self.sharpen_mode != 0:
r += f"sharpen_amount : {self.sharpen_amount}\n"
r += f"blursharpen_amount : {self.blursharpen_amount}\n"
r += f"super_resolution_mode : {self.super_res_dict[self.super_resolution_mode]}\n"
return r

51
facelib/LandmarksProcessor.py
View file

@ -183,6 +183,55 @@ landmarks_68_3D = np.array( [
[0.205322 , 31.408738 , -21.903670 ],
[-7.198266 , 30.844876 , -20.328022 ] ], dtype=np.float32)
def convert_98_to_68(lmrks):
#jaw
result = [ lmrks[0] ]
for i in range(2,16,2):
result += [ ( lmrks[i] + (lmrks[i-1]+lmrks[i+1])/2 ) / 2 ]
result += [ lmrks[16] ]
for i in range(18,32,2):
result += [ ( lmrks[i] + (lmrks[i-1]+lmrks[i+1])/2 ) / 2 ]
result += [ lmrks[32] ]
#eyebrows averaging
result += [ lmrks[33],
(lmrks[34]+lmrks[41])/2,
(lmrks[35]+lmrks[40])/2,
(lmrks[36]+lmrks[39])/2,
(lmrks[37]+lmrks[38])/2,
]
result += [ (lmrks[42]+lmrks[50])/2,
(lmrks[43]+lmrks[49])/2,
(lmrks[44]+lmrks[48])/2,
(lmrks[45]+lmrks[47])/2,
lmrks[46]
]
#nose
result += list ( lmrks[51:60] )
#left eye (from our view)
result += [ lmrks[60],
lmrks[61],
lmrks[63],
lmrks[64],
lmrks[65],
lmrks[67] ]
#right eye
result += [ lmrks[68],
lmrks[69],
lmrks[71],
lmrks[72],
lmrks[73],
lmrks[75] ]
#mouth
result += list ( lmrks[76:96] )
return np.concatenate (result).reshape ( (68,2) )
def transform_points(points, mat, invert=False):
if invert:
mat = cv2.invertAffineTransform (mat)
@ -613,5 +662,5 @@ def estimate_pitch_yaw_roll(aligned_256px_landmarks):
pitch, yaw, roll = mathlib.rotationMatrixToEulerAngles( cv2.Rodrigues(rotation_vector)[0] )
pitch = np.clip ( pitch/1.30, -1.0, 1.0 )
yaw = np.clip ( yaw / 1.11, -1.0, 1.0 )
roll = np.clip ( roll/3.15, -1.0, 1.0 )
roll = np.clip ( roll/3.15, -1.0, 1.0 ) #todo radians
return -pitch, yaw, roll

2
interact/interact.py
View file

@ -154,7 +154,7 @@ class InteractBase(object):
self.pg_bar = None
else: print("progress_bar not set.")
def progress_bar_generator(self, data, desc, leave=True, initial=0):
def progress_bar_generator(self, data, desc=None, leave=True, initial=0):
self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True, initial=initial )
for x in self.pg_bar:
yield x

18
main.py
View file

@ -49,6 +49,20 @@ 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_dev_extract_vggface2_dataset(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.extract_vggface2_dataset( arguments.input_dir,
device_args={'cpu_only' : arguments.cpu_only,
'multi_gpu' : arguments.multi_gpu,
}
)
p = subparsers.add_parser( "dev_extract_vggface2_dataset", help="")
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_dev_extract_vggface2_dataset)
def process_dev_extract_umd_csv(arguments):
os_utils.set_process_lowest_prio()
@ -152,7 +166,8 @@ if __name__ == "__main__":
def process_convert(arguments):
os_utils.set_process_lowest_prio()
args = {'input_dir' : arguments.input_dir,
args = {'training_data_src_dir' : arguments.training_data_src_dir,
'input_dir' : arguments.input_dir,
'output_dir' : arguments.output_dir,
'aligned_dir' : arguments.aligned_dir,
'model_dir' : arguments.model_dir,
@ -165,6 +180,7 @@ if __name__ == "__main__":
Converter.main (args, device_args)
p = subparsers.add_parser( "convert", help="Converter")
p.add_argument('--training-data-src-dir', action=fixPathAction, dest="training_data_src_dir", help="(optional, may be required by some models) Dir of extracted SRC faceset.")
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('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.")
p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored.")

26
mainscripts/Converter.py
View file

@ -87,10 +87,10 @@ class ConvertSubprocessor(Subprocessor):
#therefore forcing active_DeviceConfig to CPU only
nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True)
def sharpen_func (img, sharpen_mode=0, kernel_size=3, amount=150):
def blursharpen_func (img, sharpen_mode=0, kernel_size=3, amount=100):
if kernel_size % 2 == 0:
kernel_size += 1
if amount > 0:
if sharpen_mode == 1: #box
kernel = np.zeros( (kernel_size, kernel_size), dtype=np.float32)
kernel[ kernel_size//2, kernel_size//2] = 1.0
@ -101,8 +101,12 @@ class ConvertSubprocessor(Subprocessor):
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
img = cv2.addWeighted(img, 1.0 + (0.5 * amount), blur, -(0.5 * amount), 0)
return img
elif amount < 0:
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
img = cv2.addWeighted(img, 1.0 - a / 50.0, blur, a /50.0, 0)
return img
self.sharpen_func = sharpen_func
return img
self.blursharpen_func = blursharpen_func
self.fanseg_by_face_type = {}
self.fanseg_input_size = 256
@ -128,7 +132,7 @@ class ConvertSubprocessor(Subprocessor):
#override
def process_data(self, pf): #pf=ProcessingFrame
cfg = pf.cfg.copy()
cfg.sharpen_func = self.sharpen_func
cfg.blursharpen_func = self.blursharpen_func
cfg.superres_func = self.superres_func
cfg.ebs_ct_func = self.ebs_ct_func
@ -221,6 +225,8 @@ class ConvertSubprocessor(Subprocessor):
session_data = None
if self.is_interactive and self.converter_session_filepath.exists():
if io.input_bool ("Use saved session? (y/n skip:y) : ", True):
try:
with open( str(self.converter_session_filepath), "rb") as f:
session_data = pickle.loads(f.read())
@ -430,9 +436,9 @@ class ConvertSubprocessor(Subprocessor):
elif chr_key == 'g':
cfg.add_color_degrade_power(-1 if not shift_pressed else -5)
elif chr_key == 'y':
cfg.add_sharpen_amount(1 if not shift_pressed else 5)
cfg.add_blursharpen_amount(1 if not shift_pressed else 5)
elif chr_key == 'h':
cfg.add_sharpen_amount(-1 if not shift_pressed else -5)
cfg.add_blursharpen_amount(-1 if not shift_pressed else -5)
elif chr_key == 'u':
cfg.add_output_face_scale(1 if not shift_pressed else 5)
elif chr_key == 'j':
@ -453,9 +459,9 @@ class ConvertSubprocessor(Subprocessor):
else:
if chr_key == 'y':
cfg.add_sharpen_amount(1 if not shift_pressed else 5)
cfg.add_blursharpen_amount(1 if not shift_pressed else 5)
elif chr_key == 'h':
cfg.add_sharpen_amount(-1 if not shift_pressed else -5)
cfg.add_blursharpen_amount(-1 if not shift_pressed else -5)
elif chr_key == 's':
cfg.toggle_add_source_image()
elif chr_key == 'v':
@ -576,6 +582,8 @@ class ConvertSubprocessor(Subprocessor):
def main (args, device_args):
io.log_info ("Running converter.\r\n")
training_data_src_dir = args.get('training_data_src_dir', None)
training_data_src_path = Path(training_data_src_dir) if training_data_src_dir is not None else None
aligned_dir = args.get('aligned_dir', None)
avaperator_aligned_dir = args.get('avaperator_aligned_dir', None)
@ -598,7 +606,7 @@ def main (args, device_args):
is_interactive = io.input_bool ("Use interactive converter? (y/n skip:y) : ", True) if not io.is_colab() else False
import models
model = models.import_model( args['model_name'] )(model_path, device_args=device_args)
model = models.import_model( args['model_name'])(model_path, device_args=device_args, training_data_src_path=training_data_src_path)
converter_session_filepath = model.get_strpath_storage_for_file('converter_session.dat')
predictor_func, predictor_input_shape, cfg = model.get_ConverterConfig()

222
mainscripts/Extractor.py
View file

@ -1,26 +1,27 @@
import traceback
import math
import multiprocessing
import operator
import os
import shutil
import sys
import time
import multiprocessing
import shutil
from pathlib import Path
import numpy as np
import math
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 numpy as np
import facelib
from facelib import FaceType
from facelib import LandmarksProcessor
from facelib import FANSegmentator
from nnlib import nnlib
from joblib import Subprocessor
import imagelib
import mathlib
from facelib import FaceType, FANSegmentator, LandmarksProcessor
from interact import interact as io
from joblib import Subprocessor
from nnlib import nnlib
from utils import Path_utils
from utils.cv2_utils import *
from utils.DFLJPG import DFLJPG
from utils.DFLPNG import DFLPNG
DEBUG = False
@ -43,6 +44,7 @@ class ExtractSubprocessor(Subprocessor):
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.max_faces_from_image = client_dict['max_faces_from_image']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = Path(client_dict['final_output_dir']) if 'final_output_dir' in client_dict.keys() else None
@ -155,6 +157,13 @@ class ExtractSubprocessor(Subprocessor):
if len(rects) != 0:
break
if self.max_faces_from_image != 0 and len(data.rects) > 1:
#sort by largest area first
x = [ [(l,t,r,b), (r-l)*(b-t) ] for (l,t,r,b) in data.rects]
x = sorted(x, key=operator.itemgetter(1), reverse=True )
x = [ a[0] for a in x]
data.rects = x[0:self.max_faces_from_image]
return data
elif self.type == 'landmarks':
@ -283,7 +292,7 @@ class ExtractSubprocessor(Subprocessor):
return data.filename
#override
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):
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, max_faces_from_image=0, final_output_path=None):
self.input_data = input_data
self.type = type
self.image_size = image_size
@ -292,6 +301,7 @@ class ExtractSubprocessor(Subprocessor):
self.final_output_path = final_output_path
self.manual = manual
self.manual_window_size = manual_window_size
self.max_faces_from_image = max_faces_from_image
self.result = []
self.devices = ExtractSubprocessor.get_devices_for_config(self.manual, self.type, multi_gpu, cpu_only)
@ -341,6 +351,7 @@ class ExtractSubprocessor(Subprocessor):
base_dict = {'type' : self.type,
'image_size': self.image_size,
'face_type': self.face_type,
'max_faces_from_image':self.max_faces_from_image,
'debug_dir': self.debug_dir,
'final_output_dir': str(self.final_output_path),
'stdin_fd': sys.stdin.fileno() }
@ -681,8 +692,96 @@ class DeletedFilesSearcherSubprocessor(Subprocessor):
def get_result(self):
return self.result
def main(input_dir,
output_dir,
debug_dir=None,
detector='mt',
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
image_size=256,
face_type='full_face',
max_faces_from_image=0,
device_args={}):
#currently unused
input_path = Path(input_dir)
output_path = Path(output_dir)
face_type = FaceType.fromString(face_type)
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
if output_path.exists():
if not manual_output_debug_fix and input_path != output_path:
output_images_paths = Path_utils.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool("WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False )
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
if debug_dir is None:
raise ValueError('debug-dir must be specified')
detector = 'manual'
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
if debug_dir is not None:
debug_output_path = Path(debug_dir)
if manual_output_debug_fix:
if not debug_output_path.exists():
raise ValueError("%s not found " % ( str(debug_output_path) ))
input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, Path_utils.get_image_paths(debug_output_path) ).run()
input_path_image_paths = sorted (input_path_image_paths)
io.log_info('Found %d images.' % (len(input_path_image_paths)))
else:
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info ('Performing manual extract...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug_dir, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).run()
else:
io.log_info ('Performing 1st pass...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in input_path_image_paths ], 'rects-'+detector, image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, max_faces_from_image=max_faces_from_image).run()
io.log_info ('Performing 2nd pass...')
data = ExtractSubprocessor (data, 'landmarks', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
io.log_info ('Performing 3rd pass...')
data = ExtractSubprocessor (data, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all ( np.array ( [ d.faces_detected > 0 for d in data] ) == True ):
io.log_info ('All faces are detected, manual fix not needed.')
else:
fix_data = [ ExtractSubprocessor.Data(d.filename) for d in data if d.faces_detected == 0 ]
io.log_info ('Performing manual fix for %d images...' % (len(fix_data)) )
fix_data = ExtractSubprocessor (fix_data, 'landmarks', image_size, face_type, debug_dir, manual=True, manual_window_size=manual_window_size).run()
fix_data = ExtractSubprocessor (fix_data, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info ('-------------------------')
io.log_info ('Images found: %d' % (images_found) )
io.log_info ('Faces detected: %d' % (faces_detected) )
io.log_info ('-------------------------')
#unused in end user workflow
def extract_fanseg(input_dir, device_args={} ):
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
@ -709,6 +808,7 @@ def extract_fanseg(input_dir, device_args={} ):
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()
#unused in end user workflow
def extract_umd_csv(input_file_csv,
image_size=256,
face_type='full_face',
@ -781,94 +881,6 @@ def extract_umd_csv(input_file_csv,
faces_detected += sum([d.faces_detected for d in data])
io.log_info ('-------------------------')
io.log_info ('Images found: %d' % (images_found) )
io.log_info ('Faces detected: %d' % (faces_detected) )
io.log_info ('-------------------------')
def main(input_dir,
output_dir,
debug_dir=None,
detector='mt',
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
image_size=256,
face_type='full_face',
device_args={}):
input_path = Path(input_dir)
output_path = Path(output_dir)
face_type = FaceType.fromString(face_type)
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
if output_path.exists():
if not manual_output_debug_fix and input_path != output_path:
output_images_paths = Path_utils.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool("WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False )
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
if debug_dir is None:
raise ValueError('debug-dir must be specified')
detector = 'manual'
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
if debug_dir is not None:
debug_output_path = Path(debug_dir)
if manual_output_debug_fix:
if not debug_output_path.exists():
raise ValueError("%s not found " % ( str(debug_output_path) ))
input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, Path_utils.get_image_paths(debug_output_path) ).run()
input_path_image_paths = sorted (input_path_image_paths)
io.log_info('Found %d images.' % (len(input_path_image_paths)))
else:
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info ('Performing manual extract...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug_dir, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).run()
else:
io.log_info ('Performing 1st pass...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in input_path_image_paths ], 'rects-'+detector, image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
io.log_info ('Performing 2nd pass...')
data = ExtractSubprocessor (data, 'landmarks', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
io.log_info ('Performing 3rd pass...')
data = ExtractSubprocessor (data, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all ( np.array ( [ d.faces_detected > 0 for d in data] ) == True ):
io.log_info ('All faces are detected, manual fix not needed.')
else:
fix_data = [ ExtractSubprocessor.Data(d.filename) for d in data if d.faces_detected == 0 ]
io.log_info ('Performing manual fix for %d images...' % (len(fix_data)) )
fix_data = ExtractSubprocessor (fix_data, 'landmarks', image_size, face_type, debug_dir, manual=True, manual_window_size=manual_window_size).run()
fix_data = ExtractSubprocessor (fix_data, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info ('-------------------------')
io.log_info ('Images found: %d' % (images_found) )
io.log_info ('Faces detected: %d' % (faces_detected) )

1
mainscripts/Trainer.py
View file

@ -45,6 +45,7 @@ def trainerThread (s2c, c2s, e, args, device_args):
training_data_src_path=training_data_src_path,
training_data_dst_path=training_data_dst_path,
pretraining_data_path=pretraining_data_path,
is_training=True,
debug=debug,
device_args=device_args)

50
mainscripts/dev_misc.py Normal file
View file

@ -0,0 +1,50 @@
from . import Extractor
from . import Sorter
from pathlib import Path
from utils import Path_utils
import shutil
from interact import interact as io
def extract_vggface2_dataset(input_dir, device_args={} ):
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.')
output_path = input_path.parent / (input_path.name + '_out')
dir_names = Path_utils.get_all_dir_names(input_path)
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
for dir_name in dir_names:
cur_input_path = input_path / dir_name
cur_output_path = output_path / dir_name
io.log_info (f"Processing: {str(cur_input_path)} ")
if not cur_output_path.exists():
cur_output_path.mkdir(parents=True, exist_ok=True)
Extractor.main( str(cur_input_path),
str(cur_output_path),
detector='s3fd',
image_size=256,
face_type='full_face',
max_faces_from_image=1,
device_args=device_args )
io.log_info (f"Sorting: {str(cur_input_path)} ")
Sorter.main (input_path=str(cur_output_path), sort_by_method='hist')
try:
io.log_info (f"Removing: {str(cur_input_path)} ")
shutil.rmtree(cur_input_path)
except:
io.log_info (f"unable to remove: {str(cur_input_path)} ")

BIN
mainscripts/gfx/help_converter_face_avatar.jpg
View file

Binary file not shown.
Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 159 KiB

After

Width:  |  Height:  |  Size: 222 KiB

Before After
Before After

BIN
mainscripts/gfx/help_converter_face_avatar_source.psd
View file

Binary file not shown.

BIN
mainscripts/gfx/help_converter_masked.jpg
View file

Binary file not shown.
Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 288 KiB

After

Width:  |  Height:  |  Size: 385 KiB

Before After
Before After

BIN
mainscripts/gfx/help_converter_masked_source.psd
View file

Binary file not shown.

66
models/ModelBase.py
View file

@ -23,7 +23,7 @@ You can implement your own model. Check examples.
class ModelBase(object):
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, pretraining_data_path=None, debug = False, device_args = None,
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, pretraining_data_path=None, is_training=False, debug = False, device_args = None,
ask_enable_autobackup=True,
ask_write_preview_history=True,
ask_target_iter=True,
@ -56,14 +56,8 @@ class ModelBase(object):
self.training_data_dst_path = training_data_dst_path
self.pretraining_data_path = pretraining_data_path
self.src_images_paths = None
self.dst_images_paths = None
self.src_yaw_images_paths = None
self.dst_yaw_images_paths = None
self.src_data_generator = None
self.dst_data_generator = None
self.debug = debug
self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
self.is_training_mode = is_training
self.iter = 0
self.options = {}
@ -412,40 +406,60 @@ class ModelBase(object):
cv2_imwrite (filepath, img )
def load_weights_safe(self, model_filename_list, optimizer_filename_list=[]):
exec(nnlib.code_import_all, locals(), globals())
loaded = []
not_loaded = []
for mf in model_filename_list:
model, filename = mf
filename = self.get_strpath_storage_for_file(filename)
if Path(filename).exists():
loaded += [ mf ]
if issubclass(model.__class__, keras.optimizers.Optimizer):
opt = model
try:
with open(filename, "rb") as f:
fd = pickle.loads(f.read())
weights = fd.get('weights', None)
if weights is not None:
opt.set_weights(weights)
except Exception as e:
print ("Unable to load ", filename)
else:
model.load_weights(filename)
else:
not_loaded += [ mf ]
if len(optimizer_filename_list) != 0:
opt_filename = self.get_strpath_storage_for_file('opt.h5')
if Path(opt_filename).exists():
try:
with open(opt_filename, "rb") as f:
d = pickle.loads(f.read())
for x in optimizer_filename_list:
opt, filename = x
if filename in d:
weights = d[filename].get('weights', None)
if weights:
opt.set_weights(weights)
print("set ok")
except Exception as e:
print ("Unable to load ", opt_filename)
return loaded, not_loaded
def save_weights_safe(self, model_filename_list):
exec(nnlib.code_import_all, locals(), globals())
for model, filename in model_filename_list:
filename = self.get_strpath_storage_for_file(filename)
model.save_weights( filename + '.tmp' )
filename = self.get_strpath_storage_for_file(filename) + '.tmp'
if issubclass(model.__class__, keras.optimizers.Optimizer):
opt = model
try:
fd = {}
symbolic_weights = getattr(opt, 'weights')
if symbolic_weights:
fd['weights'] = self.K.batch_get_value(symbolic_weights)
with open(filename, 'wb') as f:
f.write( pickle.dumps(fd) )
except Exception as e:
print ("Unable to save ", filename)
else:
model.save_weights( filename)
rename_list = model_filename_list

166
models/Model_DEV_FUNIT/Model.py Normal file
View file

@ -0,0 +1,166 @@
from functools import partial
import cv2
import numpy as np
from facelib import FaceType
from interact import interact as io
from mathlib import get_power_of_two
from models import ModelBase
from nnlib import nnlib, FUNIT
from samplelib import *
class FUNITModel(ModelBase):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs,
ask_sort_by_yaw=False,
ask_random_flip=False,
ask_src_scale_mod=False)
#override
def onInitializeOptions(self, is_first_run, ask_override):
default_face_type = 'f'
if is_first_run:
self.options['resolution'] = io.input_int("Resolution ( 128,224 ?:help skip:128) : ", 128, [128,224])
else:
self.options['resolution'] = self.options.get('resolution', 128)
if is_first_run:
self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="").lower()
else:
self.options['face_type'] = self.options.get('face_type', default_face_type)
#override
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({4:16})
resolution = self.options['resolution']
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
person_id_max_count = SampleGeneratorFace.get_person_id_max_count(self.training_data_src_path)
self.model = FUNIT( face_type_str=FaceType.toString(face_type),
batch_size=self.batch_size,
encoder_nf=64,
encoder_downs=2,
encoder_res_blk=2,
class_downs=4,
class_nf=64,
class_latent=64,
mlp_nf=256,
mlp_blks=2,
dis_nf=64,
dis_res_blks=10,
num_classes=person_id_max_count,
subpixel_decoder=True,
initialize_weights=self.is_first_run(),
is_training=self.is_training_mode
)
if not self.is_first_run():
self.load_weights_safe(self.model.get_model_filename_list())
if self.is_training_mode:
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':128, 'normalize_tanh':True} ]
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types, person_id_mode=True ),
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types, person_id_mode=True ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types, person_id_mode=True ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types, person_id_mode=True ),
])
#override
def get_model_filename_list(self):
return self.model.get_model_filename_list()
#override
def onSave(self):
self.save_weights_safe(self.model.get_model_filename_list())
#override
def onTrainOneIter(self, generators_samples, generators_list):
xa,la = generators_samples[0]
xb,lb = generators_samples[1]
G_loss, D_loss = self.model.train(xa,la,xb,lb)
return ( ('G_loss', G_loss), ('D_loss', D_loss), )
#override
def onGetPreview(self, generators_samples):
xa = generators_samples[0][0]
xb = generators_samples[1][0]
ta = generators_samples[2][0]
tb = generators_samples[3][0]
view_samples = min(4, xa.shape[0])
lines_train = []
lines_test = []
for i in range(view_samples):
s_xa = self.model.get_average_class_code([ xa[i:i+1] ])[0][None,...]
s_xb = self.model.get_average_class_code([ xb[i:i+1] ])[0][None,...]
s_ta = self.model.get_average_class_code([ ta[i:i+1] ])[0][None,...]
s_tb = self.model.get_average_class_code([ tb[i:i+1] ])[0][None,...]
xaxa = self.model.convert ([ xa[i:i+1], s_xa ] )[0][0]
xbxb = self.model.convert ([ xb[i:i+1], s_xb ] )[0][0]
xaxb = self.model.convert ([ xa[i:i+1], s_xb ] )[0][0]
xbxa = self.model.convert ([ xb[i:i+1], s_xa ] )[0][0]
tata = self.model.convert ([ ta[i:i+1], s_ta ] )[0][0]
tbtb = self.model.convert ([ tb[i:i+1], s_tb ] )[0][0]
tatb = self.model.convert ([ ta[i:i+1], s_tb ] )[0][0]
tbta = self.model.convert ([ tb[i:i+1], s_ta ] )[0][0]
line_train = [ xa[i], xaxa, xb[i], xbxb, xaxb, xbxa ]
line_test = [ ta[i], tata, tb[i], tbtb, tatb, tbta ]
lines_train += [ np.concatenate([ np.clip(x/2+0.5,0,1) for x in line_train], axis=1) ]
lines_test += [ np.concatenate([ np.clip(x/2+0.5,0,1) for x in line_test ], axis=1) ]
lines_train = np.concatenate ( lines_train, axis=0 )
lines_test = np.concatenate ( lines_test, axis=0 )
return [ ('TRAIN', lines_train ), ('TEST', lines_test) ]
def predictor_func (self, face=None, dummy_predict=False):
if dummy_predict:
self.model.convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ), self.average_class_code ])
else:
bgr, = self.model.convert ([ face[np.newaxis,...]*2-1, self.average_class_code ])
return bgr[0] / 2 + 0.5
#override
def get_ConverterConfig(self):
face_type = FaceType.FULL
import converters
return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type,
default_mode = 1,
clip_hborder_mask_per=0.0625 if (face_type == FaceType.FULL) else 0,
)
Model = FUNITModel

1
models/Model_DEV_FUNIT/__init__.py Normal file
View file

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

49
models/Model_SAE/Model.py
View file

@ -133,15 +133,15 @@ class SAEModel(ModelBase):
def upscale (dim):
def func(x):
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x)))
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))))
return func
def enc_flow(e_dims, ae_dims, lowest_dense_res):
def func(x):
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = Dense(ae_dims)(Flatten()(x))
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
@ -151,18 +151,18 @@ class SAEModel(ModelBase):
return func
def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True):
dims = output_nc * d_ch_dims
def ResidualBlock(dim):
def func(inp):
x = Conv2D(dim, kernel_size=3, padding='same')(inp)
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(inp))
x = LeakyReLU(0.2)(x)
x = Conv2D(dim, kernel_size=3, padding='same')(x)
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(x))
x = Add()([x, inp])
x = LeakyReLU(0.2)(x)
return x
return func
def func(x):
dims = output_nc * d_ch_dims
x = upscale(dims*8)(x)
if add_residual_blocks:
@ -181,7 +181,7 @@ class SAEModel(ModelBase):
x = ResidualBlock(dims*2)(x)
x = ResidualBlock(dims*2)(x)
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)
return Conv2D(output_nc, kernel_size=5, padding='valid', activation='sigmoid')(ZeroPadding2D(2)(x))
return func
self.encoder = modelify(enc_flow(e_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) )
@ -232,20 +232,20 @@ class SAEModel(ModelBase):
mask_shape = (resolution, resolution, 1)
e_dims = output_nc*e_ch_dims
d_dims = output_nc*d_ch_dims
lowest_dense_res = resolution // 16
def upscale (dim):
def func(x):
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x)))
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))))
return func
def enc_flow(e_dims):
def func(x):
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
x = Flatten()(x)
return x
return func
@ -259,12 +259,13 @@ class SAEModel(ModelBase):
return x
return func
def dec_flow(output_nc, d_dims):
def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True):
d_dims = output_nc*d_ch_dims
def ResidualBlock(dim):
def func(inp):
x = Conv2D(dim, kernel_size=3, padding='same')(inp)
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(inp))
x = LeakyReLU(0.2)(x)
x = Conv2D(dim, kernel_size=3, padding='same')(x)
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(inp))
x = Add()([x, inp])
x = LeakyReLU(0.2)(x)
return x
@ -272,18 +273,24 @@ class SAEModel(ModelBase):
def func(x):
x = upscale(d_dims*8)(x)
if add_residual_blocks:
x = ResidualBlock(d_dims*8)(x)
x = ResidualBlock(d_dims*8)(x)
x = upscale(d_dims*4)(x)
if add_residual_blocks:
x = ResidualBlock(d_dims*4)(x)
x = ResidualBlock(d_dims*4)(x)
x = upscale(d_dims*2)(x)
if add_residual_blocks:
x = ResidualBlock(d_dims*2)(x)
x = ResidualBlock(d_dims*2)(x)
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)
return Conv2D(output_nc, kernel_size=5, padding='valid', activation='sigmoid')(ZeroPadding2D(2)(x))
return func
self.encoder = modelify(enc_flow(e_dims)) ( Input(bgr_shape) )
@ -293,10 +300,10 @@ class SAEModel(ModelBase):
self.inter_AB = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) )
sh = np.array(K.int_shape( self.inter_B.outputs[0] )[1:])*(1,1,2)
self.decoder = modelify(dec_flow(output_nc, d_dims)) ( Input(sh) )
self.decoder = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )
if learn_mask:
self.decoderm = modelify(dec_flow(1, d_dims)) ( Input(sh) )
self.decoderm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) )
self.src_dst_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights

180
models/Model_TrueFace/Model.py Normal file
View file

@ -0,0 +1,180 @@
import numpy as np
from facelib import FaceType
from interact import interact as io
from models import ModelBase
from nnlib import nnlib, FUNIT
from samplelib import *
class TrueFaceModel(ModelBase):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs,
ask_sort_by_yaw=False,
ask_random_flip=False,
ask_src_scale_mod=False)
#override
def onInitializeOptions(self, is_first_run, ask_override):
default_resolution = 128
default_face_type = 'f'
if is_first_run:
resolution = self.options['resolution'] = io.input_int(f"Resolution ( 64-256 ?:help skip:{default_resolution}) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
resolution = np.clip (resolution, 64, 256)
while np.modf(resolution / 16)[0] != 0.0:
resolution -= 1
else:
self.options['resolution'] = self.options.get('resolution', default_resolution)
if is_first_run:
self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="").lower()
else:
self.options['face_type'] = self.options.get('face_type', default_face_type)
#override
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({2:1,3:1,4:4,5:8,6:16})
resolution = self.options['resolution']
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
self.model = FUNIT( face_type_str=FaceType.toString(face_type),
batch_size=self.batch_size,
encoder_nf=64,
encoder_downs=2,
encoder_res_blk=2,
class_downs=4,
class_nf=64,
class_latent=64,
mlp_nf=256,
mlp_blks=2,
dis_nf=64,
dis_res_blks=10,
num_classes=2,
subpixel_decoder=True,
initialize_weights=self.is_first_run(),
is_training=self.is_training_mode
)
if not self.is_first_run():
self.load_weights_safe(self.model.get_model_filename_list())
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
if self.is_training_mode:
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True},
]
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types )
])
else:
generator = SampleGeneratorFace(self.training_data_src_path, batch_size=1,
sample_process_options=SampleProcessor.Options(),
output_sample_types=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ] )
io.log_info("Calculating average src face style...")
codes = []
for i in io.progress_bar_generator(range(generator.get_total_sample_count())):
codes += self.model.get_average_class_code( generator.generate_next() )
self.average_class_code = np.mean ( np.array(codes), axis=0 )[None,...]
#override
def get_model_filename_list(self):
return self.model.get_model_filename_list()
#override
def onSave(self):
self.save_weights_safe(self.model.get_model_filename_list())
#override
def onTrainOneIter(self, generators_samples, generators_list):
bs = self.batch_size
lbs = bs // 2
hbs = bs - lbs
src, = generators_samples[0]
dst, = generators_samples[1]
xa = np.concatenate ( [src[0:lbs], dst[0:lbs]], axis=0 )
la = np.concatenate ( [ np.array ([0]*lbs, np.int32),
np.array ([1]*lbs, np.int32) ] )
xb = np.concatenate ( [src[lbs:], dst[lbs:]], axis=0 )
lb = np.concatenate ( [ np.array ([0]*hbs, np.int32),
np.array ([1]*hbs, np.int32) ] )
rnd_list = np.arange(lbs*2)
np.random.shuffle(rnd_list)
xa = xa[rnd_list,...]
la = la[rnd_list,...]
la = la[...,None]
rnd_list = np.arange(hbs*2)
np.random.shuffle(rnd_list)
xb = xb[rnd_list,...]
lb = lb[rnd_list,...]
lb = lb[...,None]
G_loss, D_loss = self.model.train(xa,la,xb,lb)
return ( ('G_loss', G_loss), ('D_loss', D_loss), )
#override
def onGetPreview(self, generators_samples):
xa = generators_samples[0][0]
xb = generators_samples[1][0]
view_samples = min(4, xa.shape[0])
s_xa_mean = self.model.get_average_class_code([xa])[0][None,...]
s_xb_mean = self.model.get_average_class_code([xb])[0][None,...]
s_xab_mean = self.model.get_average_class_code([ np.concatenate( [xa,xb], axis=0) ])[0][None,...]
lines = []
for i in range(view_samples):
xaxa, = self.model.convert ([ xa[i:i+1], s_xa_mean ] )
xbxb, = self.model.convert ([ xb[i:i+1], s_xb_mean ] )
xbxa, = self.model.convert ([ xb[i:i+1], s_xa_mean ] )
xa_i,xb_i,xaxa,xbxb,xbxa = [ np.clip(x/2+0.5, 0, 1) for x in [xa[i], xb[i], xaxa[0],xbxb[0],xbxa[0]] ]
lines += [ np.concatenate( (xa_i, xaxa, xb_i, xbxb, xbxa), axis=1) ]
r = np.concatenate ( lines, axis=0 )
return [ ('TrueFace', r ) ]
def predictor_func (self, face=None, dummy_predict=False):
if dummy_predict:
self.model.convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ), self.average_class_code ])
else:
bgr, = self.model.convert ([ face[np.newaxis,...]*2-1, self.average_class_code ])
return bgr[0] / 2 + 0.5
#override
def get_ConverterConfig(self):
face_type = FaceType.FULL
import converters
return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type,
default_mode = 1,
clip_hborder_mask_per=0.0625 if (face_type == FaceType.FULL) else 0,
)
Model = TrueFaceModel

1
models/Model_TrueFace/__init__.py Normal file
View file

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

343
nnlib/FUNIT.py Normal file
View file

@ -0,0 +1,343 @@
from pathlib import Path
import numpy as np
from interact import interact as io
from nnlib import nnlib
"""
My port of FUNIT: Few-Shot Unsupervised Image-to-Image Translation to pure keras.
original repo: https://github.com/NVlabs/FUNIT/
"""
class FUNIT(object):
VERSION = 1
def __init__ (self, face_type_str,
batch_size,
encoder_nf=64,
encoder_downs=2,
encoder_res_blk=2,
class_downs=4,
class_nf=64,
class_latent=64,
mlp_nf=256,
mlp_blks=2,
dis_nf=64,
dis_res_blks=10,
num_classes=2,
subpixel_decoder=True,
initialize_weights=True,
load_weights_locally=False,
weights_file_root=None,
is_training=True
):
exec( nnlib.import_all(), locals(), globals() )
self.batch_size = batch_size
bgr_shape = (None, None, 3)
label_shape = (1,)
self.enc_content = modelify ( FUNIT.ContentEncoderFlow(downs=encoder_downs, nf=encoder_nf, n_res_blks=encoder_res_blk) ) ( Input(bgr_shape) )
self.enc_class_model = modelify ( FUNIT.ClassModelEncoderFlow(downs=class_downs, nf=class_nf, latent_dim=class_latent) ) ( Input(bgr_shape) )
self.decoder = modelify ( FUNIT.DecoderFlow(ups=encoder_downs, n_res_blks=encoder_res_blk, mlp_nf=mlp_nf, mlp_blks=mlp_blks, subpixel_decoder=subpixel_decoder ) ) \
( [ Input(K.int_shape(self.enc_content.outputs[0])[1:], name="decoder_input_1"),
Input(K.int_shape(self.enc_class_model.outputs[0])[1:], name="decoder_input_2")
] )
self.dis = modelify ( FUNIT.DiscriminatorFlow(nf=dis_nf, n_res_blks=dis_res_blks, num_classes=num_classes) ) (Input(bgr_shape))
self.G_opt = RMSprop(lr=0.0001, decay=0.0001, tf_cpu_mode=2 if 'tensorflow' in nnlib.active_DeviceConfig.backend else 0)
self.D_opt = RMSprop(lr=0.0001, decay=0.0001, tf_cpu_mode=2 if 'tensorflow' in nnlib.active_DeviceConfig.backend else 0)
xa = Input(bgr_shape, name="xa")
la = Input(label_shape, dtype=np.int32, name="la")
xb = Input(bgr_shape, name="xb")
lb = Input(label_shape, dtype=np.int32, name="lb")
s_xa_one = Input( (self.enc_class_model.outputs[0].shape[-1].value,), name="s_xa_input")
c_xa = self.enc_content(xa)
s_xa = self.enc_class_model(xa)
s_xb = self.enc_class_model(xb)
s_xa_mean = K.mean(s_xa, axis=0)
xr = self.decoder ([c_xa,s_xa])
xt = self.decoder ([c_xa,s_xb])
xr_one = self.decoder ([c_xa,s_xa_one])
d_xr, d_xr_feat = self.dis(xr)
d_xt, d_xt_feat = self.dis(xt)
d_xa, d_xa_feat = self.dis(xa)
d_xb, d_xb_feat = self.dis(xb)
def dis_gather(x,l):
tensors = []
for i in range(self.batch_size):
t = x[i:i+1,:,:, l[i,0]]
tensors += [t]
return tensors
def dis_gather_batch_mean(x,l, func=None):
x_shape = K.shape(x)
b,h,w,c = x_shape[0],x_shape[1],x_shape[2],x_shape[3]
b,h,w,c = [ K.cast(x, K.floatx()) for x in [b,h,w,c] ]
tensors = dis_gather(x,l)
if func is not None:
tensors = [func(t) for t in tensors]
return K.sum(tensors, axis=[1,2,3]) / (h*w)
def dis_gather_mean(x,l, func=None, acc_func=None):
x_shape = K.shape(x)
b,h,w,c = x_shape[0],x_shape[1],x_shape[2],x_shape[3]
b,h,w,c = [ K.cast(x, K.floatx()) for x in [b,h,w,c] ]
tensors = dis_gather(x,l)
if acc_func is not None:
acc = []
for t in tensors:
acc += [ K.sum( K.cast( acc_func(t), K.floatx() )) ]
acc = K.cast( K.sum(acc), K.floatx() ) / (b*h*w)
else:
acc = None
if func is not None:
tensors = [func(t) for t in tensors]
return K.sum(tensors) / (b*h*w), acc
d_xr_la, d_xr_la_acc = dis_gather_mean(d_xr, la, acc_func=lambda x: x >= 0)
d_xt_lb, d_xt_lb_acc = dis_gather_mean(d_xt, lb, acc_func=lambda x: x >= 0)
d_xb_lb = dis_gather_batch_mean(d_xb, lb)
d_xb_lb_real, d_xb_lb_real_acc = dis_gather_mean(d_xb, lb, lambda x: K.relu(1.0-x), acc_func=lambda x: x >= 0)
d_xt_lb_fake, d_xt_lb_fake_acc = dis_gather_mean(d_xt, lb, lambda x: K.relu(1.0+x), acc_func=lambda x: x < 0)
G_c_rec = K.mean(K.abs(K.mean(d_xr_feat, axis=[1,2]) - K.mean(d_xa_feat, axis=[1,2]))) #* 1.0
G_m_rec = K.mean(K.abs(K.mean(d_xt_feat, axis=[1,2]) - K.mean(d_xb_feat, axis=[1,2]))) #* 1.0
G_x_rec = 0.1 * K.mean(K.abs(xr-xa))
G_loss = (-d_xr_la-d_xt_lb)*0.5 + G_x_rec + G_c_rec + G_m_rec
G_acc = (d_xr_la_acc+d_xt_lb_acc)*0.5
G_weights = self.enc_class_model.trainable_weights + self.enc_content.trainable_weights + self.decoder.trainable_weights
######
D_real = d_xb_lb_real #1.0 *
D_fake = d_xt_lb_fake #1.0 *
l_reg = 10 * K.sum( K.gradients( d_xb_lb, xb )[0] ** 2 ) # , axis=[1,2,3] / self.batch_size )
D_loss = D_real + D_fake + l_reg
D_acc = (d_xb_lb_real_acc+d_xt_lb_fake_acc)*0.5
D_weights = self.dis.trainable_weights
self.G_train = K.function ([xa, la, xb, lb],[G_loss], self.G_opt.get_updates(G_loss, G_weights) )
self.D_train = K.function ([xa, la, xb, lb],[D_loss], self.D_opt.get_updates(D_loss, D_weights) )
self.get_average_class_code = K.function ([xa],[s_xa_mean])
self.G_convert = K.function ([xa,s_xa_one],[xr_one])
if initialize_weights:
#gather weights from layers for initialization
weights_list = []
for model, _ in self.get_model_filename_list():
if type(model) == keras.models.Model:
for layer in model.layers:
if type(layer) == FUNITAdain:
weights_list += [ x for x in layer.weights if 'kernel' in x.name ]
elif type(layer) == keras.layers.Conv2D or type(layer) == keras.layers.Dense:
weights_list += [ layer.weights[0] ]
initer = keras.initializers.he_normal()
for w in weights_list:
K.set_value( w, K.get_value(initer(K.int_shape(w))) )
#if not self.is_first_run():
# self.load_weights_safe(self.get_model_filename_list())
if load_weights_locally:
pass
#f weights_file_root is not None:
# weights_file_root = Path(weights_file_root)
#lse:
# weights_file_root = Path(__file__).parent
#elf.weights_path = weights_file_root / ('FUNIT_%s.h5' % (face_type_str) )
#f load_weights:
# self.model.load_weights (str(self.weights_path))
def get_model_filename_list(self):
return [[self.enc_class_model, 'enc_class_model.h5'],
[self.enc_content, 'enc_content.h5'],
[self.decoder, 'decoder.h5'],
[self.dis, 'dis.h5'],
[self.G_opt, 'G_opt.h5'],
[self.D_opt, 'D_opt.h5'],
]
#def save_weights(self):
# self.model.save_weights (str(self.weights_path))
def train(self, xa,la,xb,lb):
D_loss, = self.D_train ([xa,la,xb,lb])
G_loss, = self.G_train ([xa,la,xb,lb])
return G_loss, D_loss
def get_average_class_code(self, *args, **kwargs):
return self.get_average_class_code(*args, **kwargs)
def convert(self, *args, **kwargs):
return self.G_convert(*args, **kwargs)
@staticmethod
def ContentEncoderFlow(downs=2, nf=64, n_res_blks=2):
exec (nnlib.import_all(), locals(), globals())
def ResBlock(dim):
def func(input):
x = input
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = InstanceNormalization()(x)
x = ReLU()(x)
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = InstanceNormalization()(x)
return Add()([x,input])
return func
def func(x):
x = Conv2D (nf, kernel_size=7, strides=1, padding='valid')(ZeroPadding2D(3)(x))
x = InstanceNormalization()(x)
x = ReLU()(x)
for i in range(downs):
x = Conv2D (nf * 2**(i+1), kernel_size=4, strides=2, padding='valid')(ZeroPadding2D(1)(x))
x = InstanceNormalization()(x)
x = ReLU()(x)
for i in range(n_res_blks):
x = ResBlock( nf * 2**downs )(x)
return x
return func
@staticmethod
def ClassModelEncoderFlow(downs=4, nf=64, latent_dim=64):
exec (nnlib.import_all(), locals(), globals())
def func(x):
x = Conv2D (nf, kernel_size=7, strides=1, padding='valid', activation='relu')(ZeroPadding2D(3)(x))
for i in range(downs):
x = Conv2D (nf * min ( 4, 2**(i+1) ), kernel_size=4, strides=2, padding='valid', activation='relu')(ZeroPadding2D(1)(x))
x = GlobalAveragePooling2D()(x)
x = Dense(nf)(x)
return x
return func
@staticmethod
def DecoderFlow(ups, n_res_blks=2, mlp_nf=256, mlp_blks=2, subpixel_decoder=False ):
exec (nnlib.import_all(), locals(), globals())
def ResBlock(dim):
def func(input):
inp, mlp = input
x = inp
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = FUNITAdain()([x,mlp])
x = ReLU()(x)
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = FUNITAdain()([x,mlp])
return Add()([x,inp])
return func
def func(inputs):
x , class_code = inputs
nf = x.shape[-1].value
### MLP block inside decoder
mlp = class_code
for i in range(mlp_blks):
mlp = Dense(mlp_nf, activation='relu')(mlp)
for i in range(n_res_blks):
x = ResBlock(nf)( [x,mlp] )
for i in range(ups):
if subpixel_decoder:
x = Conv2D (4* (nf // 2**(i+1)), kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = SubpixelUpscaler()(x)
else:
x = UpSampling2D()(x)
x = Conv2D (nf // 2**(i+1), kernel_size=5, strides=1, padding='valid')(ZeroPadding2D(2)(x))
x = InstanceNormalization()(x)
x = ReLU()(x)
rgb = Conv2D (3, kernel_size=7, strides=1, padding='valid', activation='tanh')(ZeroPadding2D(3)(x))
return rgb
return func
@staticmethod
def DiscriminatorFlow(nf, n_res_blks, num_classes ):
exec (nnlib.import_all(), locals(), globals())
n_layers = n_res_blks // 2
def ActFirstResBlock(fout):
def func(x):
fin = K.int_shape(x)[-1]
fhid = min(fin, fout)
if fin != fout:
x_s = Conv2D (fout, kernel_size=1, strides=1, padding='valid', use_bias=False)(x)
else:
x_s = x
x = LeakyReLU(0.2)(x)
x = Conv2D (fhid, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = LeakyReLU(0.2)(x)
x = Conv2D (fout, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
return Add()([x_s, x])
return func
def func( x ):
l_nf = nf
x = Conv2D (l_nf, kernel_size=7, strides=1, padding='valid')(ZeroPadding2D(3)(x))
for i in range(n_layers-1):
l_nf_out = min( l_nf*2, 1024 )
x = ActFirstResBlock(l_nf)(x)
x = ActFirstResBlock(l_nf_out)(x)
x = AveragePooling2D( pool_size=3, strides=2, padding='valid' )(ZeroPadding2D(1)(x))
l_nf = min( l_nf*2, 1024 )
l_nf_out = min( l_nf*2, 1024 )
x = ActFirstResBlock(l_nf)(x)
feat = x = ActFirstResBlock(l_nf_out)(x)
x = LeakyReLU(0.2)(x)
x = Conv2D (num_classes, kernel_size=1, strides=1, padding='valid')(x)
return x, feat
return func

1
nnlib/__init__.py
View file

@ -1 +1,2 @@
from .nnlib import nnlib
from .FUNIT import FUNIT

169
nnlib/nnlib.py
View file

@ -51,10 +51,11 @@ KL = keras.layers
Input = KL.Input
Dense = KL.Dense
Conv2D = nnlib.Conv2D
Conv2DTranspose = nnlib.Conv2DTranspose
Conv2D = KL.Conv2D
Conv2DTranspose = KL.Conv2DTranspose
EqualConv2D = nnlib.EqualConv2D
SeparableConv2D = KL.SeparableConv2D
DepthwiseConv2D = KL.DepthwiseConv2D
MaxPooling2D = KL.MaxPooling2D
AveragePooling2D = KL.AveragePooling2D
GlobalAveragePooling2D = KL.GlobalAveragePooling2D
@ -86,6 +87,7 @@ RandomNormal = keras.initializers.RandomNormal
Model = keras.models.Model
Adam = nnlib.Adam
RMSprop = nnlib.RMSprop
modelify = nnlib.modelify
gaussian_blur = nnlib.gaussian_blur
@ -96,6 +98,7 @@ PixelShuffler = nnlib.PixelShuffler
SubpixelUpscaler = nnlib.SubpixelUpscaler
Scale = nnlib.Scale
BlurPool = nnlib.BlurPool
FUNITAdain = nnlib.FUNITAdain
SelfAttention = nnlib.SelfAttention
CAInitializerMP = nnlib.CAInitializerMP
@ -512,6 +515,82 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
nnlib.BlurPool = BlurPool
class FUNITAdain(KL.Layer):
"""
differents from NVLabs/FUNIT:
I moved two dense blocks inside this layer,
so we don't need to slice outter MLP block and assign weights every call, just pass MLP inside.
also size of dense blocks is calculated automatically
"""
def __init__(self, axis=-1, epsilon=1e-5, momentum=0.99, **kwargs):
self.axis = axis
self.epsilon = epsilon
self.momentum = momentum
super(FUNITAdain, self).__init__(**kwargs)
def build(self, input_shape):
self.input_spec = None
x, mlp = input_shape
units = x[self.axis]
self.kernel1 = self.add_weight(shape=(units, units), initializer='he_normal', name='kernel1')
self.bias1 = self.add_weight(shape=(units,), initializer='zeros', name='bias1')
self.kernel2 = self.add_weight(shape=(units, units), initializer='he_normal', name='kernel2')
self.bias2 = self.add_weight(shape=(units,), initializer='zeros', name='bias2')
self.built = True
def call(self, inputs, training=None):
x, mlp = inputs
gamma = K.dot(mlp, self.kernel1)
gamma = K.bias_add(gamma, self.bias1, data_format='channels_last')
beta = K.dot(mlp, self.kernel2)
beta = K.bias_add(beta, self.bias2, data_format='channels_last')
input_shape = K.int_shape(x)
reduction_axes = list(range(len(input_shape)))
del reduction_axes[self.axis]
#broadcast_shape = [1] * len(input_shape)
#broadcast_shape[self.axis] = input_shape[self.axis]
#normed = x# (x - K.reshape(self.moving_mean,broadcast_shape) ) / ( K.sqrt( K.reshape(self.moving_variance,broadcast_shape)) +self.epsilon)
#normed *= K.reshape(gamma,[-1]+broadcast_shape[1:] )
#normed += K.reshape(beta, [-1]+broadcast_shape[1:] )
#mean = K.mean(x, axis=reduction_axes)
#self.moving_mean = self.add_weight(shape=(units,), name='moving_mean', initializer='zeros',trainable=False)
#self.moving_variance = self.add_weight(shape=(units,), name='moving_variance',initializer='ones', trainable=False)
#variance = K.var(x, axis=reduction_axes)
#sample_size = K.prod([ K.shape(x)[axis] for axis in reduction_axes ])
#sample_size = K.cast(sample_size, dtype=K.dtype(x))
#variance *= sample_size / (sample_size - (1.0 + self.epsilon))
#self.add_update([K.moving_average_update(self.moving_mean, mean, self.momentum),
# K.moving_average_update(self.moving_variance, variance, self.momentum)], None)
#return normed
del reduction_axes[0]
broadcast_shape = [1] * len(input_shape)
broadcast_shape[self.axis] = input_shape[self.axis]
mean = K.mean(x, reduction_axes, keepdims=True)
stddev = K.std(x, reduction_axes, keepdims=True) + self.epsilon
normed = (x - mean) / stddev
normed *= K.reshape(gamma,[-1]+broadcast_shape[1:] )
normed += K.reshape(beta, [-1]+broadcast_shape[1:] )
return normed
def get_config(self):
config = {'axis': self.axis, 'epsilon': self.epsilon }
base_config = super(FUNITAdain, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
def compute_output_shape(self, input_shape):
return input_shape
nnlib.FUNITAdain = FUNITAdain
class Scale(KL.Layer):
"""
@ -581,6 +660,92 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
return out
nnlib.SelfAttention = SelfAttention
class RMSprop(keras.optimizers.Optimizer):
"""RMSProp optimizer.
It is recommended to leave the parameters of this optimizer
at their default values
(except the learning rate, which can be freely tuned).
# Arguments
learning_rate: float >= 0. Learning rate.
rho: float >= 0.
# References
- [rmsprop: Divide the gradient by a running average of its recent magnitude
](http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
tf_cpu_mode: only for tensorflow backend
0 - default, no changes.
1 - allows to train x2 bigger network on same VRAM consuming RAM
2 - allows to train x3 bigger network on same VRAM consuming RAM*2 and CPU power.
"""
def __init__(self, learning_rate=0.001, rho=0.9, tf_cpu_mode=0, **kwargs):
self.initial_decay = kwargs.pop('decay', 0.0)
self.epsilon = kwargs.pop('epsilon', K.epsilon())
self.tf_cpu_mode = tf_cpu_mode
learning_rate = kwargs.pop('lr', learning_rate)
super(RMSprop, self).__init__(**kwargs)
with K.name_scope(self.__class__.__name__):
self.learning_rate = K.variable(learning_rate, name='learning_rate')
self.rho = K.variable(rho, name='rho')
self.decay = K.variable(self.initial_decay, name='decay')
self.iterations = K.variable(0, dtype='int64', name='iterations')
def get_updates(self, loss, params):
grads = self.get_gradients(loss, params)
e = K.tf.device("/cpu:0") if self.tf_cpu_mode > 0 else None
if e: e.__enter__()
accumulators = [K.zeros(K.int_shape(p),
dtype=K.dtype(p),
name='accumulator_' + str(i))
for (i, p) in enumerate(params)]
if e: e.__exit__(None, None, None)
self.weights = [self.iterations] + accumulators
self.updates = [K.update_add(self.iterations, 1)]
lr = self.learning_rate
if self.initial_decay > 0:
lr = lr * (1. / (1. + self.decay * K.cast(self.iterations,
K.dtype(self.decay))))
for p, g, a in zip(params, grads, accumulators):
# update accumulator
e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None
if e: e.__enter__()
new_a = self.rho * a + (1. - self.rho) * K.square(g)
new_p = p - lr * g / (K.sqrt(new_a) + self.epsilon)
if e: e.__exit__(None, None, None)
self.updates.append(K.update(a, new_a))
# Apply constraints.
if getattr(p, 'constraint', None) is not None:
new_p = p.constraint(new_p)
self.updates.append(K.update(p, new_p))
return self.updates
def set_weights(self, weights):
params = self.weights
# Override set_weights for backward compatibility of Keras 2.2.4 optimizer
# since it does not include iteration at head of the weight list. Set
# iteration to 0.
if len(params) == len(weights) + 1:
weights = [np.array(0)] + weights
super(RMSprop, self).set_weights(weights)
def get_config(self):
config = {'learning_rate': float(K.get_value(self.learning_rate)),
'rho': float(K.get_value(self.rho)),
'decay': float(K.get_value(self.decay)),
'epsilon': self.epsilon}
base_config = super(RMSprop, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
nnlib.RMSprop = RMSprop
class Adam(keras.optimizers.Optimizer):
"""Adam optimizer.

6
samplelib/Sample.py
View file

@ -22,9 +22,10 @@ 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_yaw_roll=None, eyebrows_expand_mod=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False):
def __init__(self, sample_type=None, filename=None, person_id=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, eyebrows_expand_mod=None, source_filename=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.person_id = person_id
self.face_type = face_type
self.shape = shape
self.landmarks = np.array(landmarks) if landmarks is not None else None
@ -36,10 +37,11 @@ class Sample(object):
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_yaw_roll=None, eyebrows_expand_mod=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None):
def copy_and_set(self, sample_type=None, filename=None, person_id=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, eyebrows_expand_mod=None, source_filename=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,
person_id=person_id if person_id is not None else self.person_id,
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(),

6
samplelib/SampleGeneratorBase.py
View file

@ -6,7 +6,7 @@ You can implement your own SampleGenerator
class SampleGeneratorBase(object):
def __init__ (self, samples_path, debug, batch_size):
def __init__ (self, samples_path, debug=False, batch_size=1):
if samples_path is None:
raise Exception('samples_path is None')
@ -25,6 +25,10 @@ class SampleGeneratorBase(object):
self.last_generation = next(self)
return self.last_generation
#overridable
def get_total_sample_count(self):
return 0
#overridable
def __iter__(self):
#implement your own iterator

33
samplelib/SampleGeneratorFace.py
View file

@ -18,11 +18,23 @@ output_sample_types = [
]
'''
class SampleGeneratorFace(SampleGeneratorBase):
def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, random_ct_samples_path=None, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, generators_count=2, generators_random_seed=None, **kwargs):
def __init__ (self, samples_path, debug=False, batch_size=1,
sort_by_yaw=False,
sort_by_yaw_target_samples_path=None,
random_ct_samples_path=None,
sample_process_options=SampleProcessor.Options(),
output_sample_types=[],
person_id_mode=False,
add_sample_idx=False,
generators_count=2,
generators_random_seed=None,
**kwargs):
super().__init__(samples_path, debug, batch_size)
self.sample_process_options = sample_process_options
self.output_sample_types = output_sample_types
self.add_sample_idx = add_sample_idx
self.person_id_mode = person_id_mode
if sort_by_yaw_target_samples_path is not None:
self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET
@ -36,7 +48,8 @@ class SampleGeneratorFace(SampleGeneratorBase):
self.generators_random_seed = generators_random_seed
samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path, person_id_mode=person_id_mode)
self.total_samples_count = len(samples)
ct_samples = SampleLoader.load (SampleType.FACE, random_ct_samples_path) if random_ct_samples_path is not None else None
self.random_ct_sample_chance = 100
@ -50,6 +63,10 @@ class SampleGeneratorFace(SampleGeneratorBase):
self.generator_counter = -1
#overridable
def get_total_sample_count(self):
return self.total_samples_count
def __iter__(self):
return self
@ -132,11 +149,23 @@ class SampleGeneratorFace(SampleGeneratorBase):
batches += [ [] ]
i_sample_idx = len(batches)-1
if self.person_id_mode:
batches += [ [] ]
i_person_id = len(batches)-1
for i in range(len(x)):
batches[i].append ( x[i] )
if self.add_sample_idx:
batches[i_sample_idx].append (idx)
if self.person_id_mode:
batches[i_person_id].append ( np.array([sample.person_id]) )
break
yield [ np.array(batch) for batch in batches]
@staticmethod
def get_person_id_max_count(samples_path):
return SampleLoader.get_person_id_max_count(samples_path)

18
samplelib/SampleLoader.py
View file

@ -19,7 +19,11 @@ class SampleLoader:
cache = dict()
@staticmethod
def load(sample_type, samples_path, target_samples_path=None):
def get_person_id_max_count(samples_path):
return len ( Path_utils.get_all_dir_names(samples_path) )
@staticmethod
def load(sample_type, samples_path, target_samples_path=None, person_id_mode=False):
cache = SampleLoader.cache
if str(samples_path) not in cache.keys():
@ -30,9 +34,15 @@ class SampleLoader:
if sample_type == SampleType.IMAGE:
if datas[sample_type] is None:
datas[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( Path_utils.get_image_paths(samples_path), "Loading") ]
elif sample_type == SampleType.FACE:
if datas[sample_type] is None:
if person_id_mode:
dir_names = Path_utils.get_all_dir_names(samples_path)
all_samples = []
for i, dir_name in io.progress_bar_generator( [*enumerate(dir_names)] , "Loading"):
all_samples += SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename, person_id=i) for filename in Path_utils.get_image_paths( samples_path / dir_name ) ], silent=True )
datas[sample_type] = all_samples
else:
datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] )
elif sample_type == SampleType.FACE_TEMPORAL_SORTED:
@ -52,10 +62,10 @@ class SampleLoader:
return datas[sample_type]
@staticmethod
def upgradeToFaceSamples ( samples ):
def upgradeToFaceSamples ( samples, silent=False ):
sample_list = []
for s in io.progress_bar_generator(samples, "Loading"):
for s in (samples if silent else io.progress_bar_generator(samples, "Loading")):
s_filename_path = Path(s.filename)
try:
if s_filename_path.suffix == '.png':