diff --git a/converters/ConvertMasked.py b/converters/ConvertMasked.py index 7772ed7..9460a18 100644 --- a/converters/ConvertMasked.py +++ b/converters/ConvertMasked.py @@ -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 ) diff --git a/converters/ConverterConfig.py b/converters/ConverterConfig.py index 6b63630..327a242 100644 --- a/converters/ConverterConfig.py +++ b/converters/ConverterConfig.py @@ -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 diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index c865bbf..671129a 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -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) @@ -310,8 +359,8 @@ def alpha_to_color (img_alpha, color): result[:,:] = color return result * img_alpha - - + + def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): h,w,c = image_shape @@ -361,7 +410,7 @@ def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): s,e = d[name] result = dists[...,s:e] if thickness != 0: - result = np.abs(result)-thickness + result = np.abs(result)-thickness return np.min (result, axis=-1) return get_dists @@ -371,7 +420,7 @@ def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): l_brow = lmrks[22:27] r_brow = lmrks[17:22] mouth = lmrks[48:60] - + up_nose = np.concatenate( (lmrks[27:31], lmrks[33:34]) ) down_nose = lmrks[31:36] nose = np.concatenate ( (up_nose, down_nose) ) @@ -400,7 +449,7 @@ def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): mouth_fall_dist = w // 32 mouth_thickness = max( w // 64, 1 ) - + eyes_mask = gdf('eyes',eyes_thickness) eyes_mask = 1-np.clip( eyes_mask/ eyes_fall_dist, 0, 1) #eyes_mask = np.clip ( 1- ( np.sqrt( np.maximum(eyes_mask,0) ) / eyes_fall_dist ), 0, 1) @@ -409,15 +458,15 @@ def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): brows_mask = gdf('brows', brows_thickness) brows_mask = 1-np.clip( brows_mask / brows_fall_dist, 0, 1) #brows_mask = np.clip ( 1- ( np.sqrt( np.maximum(brows_mask,0) ) / brows_fall_dist ), 0, 1) - + mouth_mask = gdf('mouth', mouth_thickness) mouth_mask = 1-np.clip( mouth_mask / mouth_fall_dist, 0, 1) #mouth_mask = np.clip ( 1- ( np.sqrt( np.maximum(mouth_mask,0) ) / mouth_fall_dist ), 0, 1) - + def blend(a,b,k): x = np.clip ( 0.5+0.5*(b-a)/k, 0.0, 1.0 ) return (a-b)*x+b - k*x*(1.0-x) - + #nose_mask = (a-b)*x+b - k*x*(1.0-x) @@ -426,7 +475,7 @@ def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): nose_mask = blend ( gdf('up_nose', nose_thickness), gdf('down_nose', nose_thickness), nose_thickness*3 ) nose_mask = 1-np.clip( nose_mask / nose_fall_dist, 0, 1) - + up_nose_mask = gdf('up_nose', nose_thickness) up_nose_mask = 1-np.clip( up_nose_mask / nose_fall_dist, 0, 1) #up_nose_mask = np.clip ( 1- ( np.cbrt( np.maximum(up_nose_mask,0) ) / nose_fall_dist ), 0, 1) @@ -441,17 +490,17 @@ def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): #nose_mask = down_nose_mask #nose_mask = np.zeros_like(nose_mask) - + eyes_mask = eyes_mask * (1-mouth_mask) nose_mask = nose_mask * (1-eyes_mask) - + hull_mask = hull[...,0].copy() hull_mask = hull_mask * (1-eyes_mask) * (1-brows_mask) * (1-nose_mask) * (1-mouth_mask) #eyes_mask = eyes_mask * (1-nose_mask) - + mouth_mask= mouth_mask * (1-nose_mask) - + brows_mask = brows_mask * (1-nose_mask)* (1-eyes_mask ) hull_mask = alpha_to_color(hull_mask, (0,1,0) ) @@ -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 diff --git a/interact/interact.py b/interact/interact.py index 25d0549..dce7c37 100644 --- a/interact/interact.py +++ b/interact/interact.py @@ -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 diff --git a/main.py b/main.py index f44f498..8401aee 100644 --- a/main.py +++ b/main.py @@ -49,7 +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_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() from mainscripts import Extractor @@ -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.") diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 0a03093..1968567 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -87,22 +87,26 @@ 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 sharpen_mode == 1: #box - kernel = np.zeros( (kernel_size, kernel_size), dtype=np.float32) - kernel[ kernel_size//2, kernel_size//2] = 1.0 - box_filter = np.ones( (kernel_size, kernel_size), dtype=np.float32) / (kernel_size**2) - kernel = kernel + (kernel - box_filter) * amount - return cv2.filter2D(img, -1, kernel) - elif sharpen_mode == 2: #gaussian + 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 + box_filter = np.ones( (kernel_size, kernel_size), dtype=np.float32) / (kernel_size**2) + kernel = kernel + (kernel - box_filter) * amount + return cv2.filter2D(img, -1, kernel) + elif sharpen_mode == 2: #gaussian + 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 + (0.5 * amount), blur, -(0.5 * amount), 0) - return img + img = cv2.addWeighted(img, 1.0 - a / 50.0, blur, a /50.0, 0) + return img return img - self.sharpen_func = sharpen_func + 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,11 +225,13 @@ class ConvertSubprocessor(Subprocessor): session_data = None if self.is_interactive and self.converter_session_filepath.exists(): - try: - with open( str(self.converter_session_filepath), "rb") as f: - session_data = pickle.loads(f.read()) - except Exception as e: - pass + + 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()) + except Exception as e: + pass self.frames = frames self.frames_idxs = [ *range(len(self.frames)) ] @@ -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() diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 30e464a..f74f174 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -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 @@ -154,6 +156,13 @@ class ExtractSubprocessor(Subprocessor): rects = data.rects = self.e.extract (rotated_image, is_bgr=True) 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 @@ -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) ) diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py index 85a6d68..839f374 100644 --- a/mainscripts/Trainer.py +++ b/mainscripts/Trainer.py @@ -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) diff --git a/mainscripts/dev_misc.py b/mainscripts/dev_misc.py new file mode 100644 index 0000000..fc3f18b --- /dev/null +++ b/mainscripts/dev_misc.py @@ -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)} ") \ No newline at end of file diff --git a/mainscripts/gfx/help_converter_face_avatar.jpg b/mainscripts/gfx/help_converter_face_avatar.jpg index 6d848c3..29b7e72 100644 Binary files a/mainscripts/gfx/help_converter_face_avatar.jpg and b/mainscripts/gfx/help_converter_face_avatar.jpg differ diff --git a/mainscripts/gfx/help_converter_face_avatar_source.psd b/mainscripts/gfx/help_converter_face_avatar_source.psd index c892610..b9d0a88 100644 Binary files a/mainscripts/gfx/help_converter_face_avatar_source.psd and b/mainscripts/gfx/help_converter_face_avatar_source.psd differ diff --git a/mainscripts/gfx/help_converter_masked.jpg b/mainscripts/gfx/help_converter_masked.jpg index 03048da..99fc66f 100644 Binary files a/mainscripts/gfx/help_converter_masked.jpg and b/mainscripts/gfx/help_converter_masked.jpg differ diff --git a/mainscripts/gfx/help_converter_masked_source.psd b/mainscripts/gfx/help_converter_masked_source.psd index c618db6..cb2c16c 100644 Binary files a/mainscripts/gfx/help_converter_masked_source.psd and b/mainscripts/gfx/help_converter_masked_source.psd differ diff --git a/models/ModelBase.py b/models/ModelBase.py index 0197a75..9b5b557 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -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 ] - model.load_weights(filename) + + 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 diff --git a/models/Model_DEV_FUNIT/Model.py b/models/Model_DEV_FUNIT/Model.py new file mode 100644 index 0000000..a9cfb67 --- /dev/null +++ b/models/Model_DEV_FUNIT/Model.py @@ -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 diff --git a/models/Model_DEV_FUNIT/__init__.py b/models/Model_DEV_FUNIT/__init__.py new file mode 100644 index 0000000..0188f11 --- /dev/null +++ b/models/Model_DEV_FUNIT/__init__.py @@ -0,0 +1 @@ +from .Model import Model diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index ccd7e58..ad3e1b3 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -51,7 +51,7 @@ class SAEModel(ModelBase): default_e_ch_dims = 42 default_d_ch_dims = default_e_ch_dims // 2 def_ca_weights = False - + if is_first_run: self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) self.options['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 ) @@ -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,37 +151,37 @@ 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: x = ResidualBlock(dims*8)(x) x = ResidualBlock(dims*8)(x) x = upscale(dims*4)(x) - + if add_residual_blocks: x = ResidualBlock(dims*4)(x) x = ResidualBlock(dims*4)(x) x = upscale(dims*2)(x) - + if add_residual_blocks: 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) - x = ResidualBlock(d_dims*8)(x) - x = ResidualBlock(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) - x = ResidualBlock(d_dims*4)(x) - x = ResidualBlock(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) - x = ResidualBlock(d_dims*2)(x) - x = ResidualBlock(d_dims*2)(x) - return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(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='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 @@ -349,17 +356,17 @@ class SAEModel(ModelBase): loaded, not_loaded = self.load_weights_safe(not_loaded) CA_models = [] - if self.options.get('ca_weights', False): + if self.options.get('ca_weights', False): CA_models += [ model for model, _ in not_loaded ] - + CA_conv_weights_list = [] for model in CA_models: for layer in model.layers: if type(layer) == keras.layers.Conv2D: CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights - + if len(CA_conv_weights_list) != 0: - CAInitializerMP ( CA_conv_weights_list ) + CAInitializerMP ( CA_conv_weights_list ) warped_src = self.model.warped_src target_src = Input ( (resolution, resolution, 3) ) @@ -501,7 +508,7 @@ class SAEModel(ModelBase): if self.options['learn_mask']: feed = [ warped_src, warped_dst, target_srcm, target_dstm ] src_mask_loss, dst_mask_loss, = self.src_dst_mask_train (feed) - + return ( ('src_loss', src_loss), ('dst_loss', dst_loss), ) #override diff --git a/models/Model_TrueFace/Model.py b/models/Model_TrueFace/Model.py new file mode 100644 index 0000000..b4c605b --- /dev/null +++ b/models/Model_TrueFace/Model.py @@ -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 diff --git a/models/Model_TrueFace/__init__.py b/models/Model_TrueFace/__init__.py new file mode 100644 index 0000000..0188f11 --- /dev/null +++ b/models/Model_TrueFace/__init__.py @@ -0,0 +1 @@ +from .Model import Model diff --git a/nnlib/FUNIT.py b/nnlib/FUNIT.py new file mode 100644 index 0000000..37e7652 --- /dev/null +++ b/nnlib/FUNIT.py @@ -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 \ No newline at end of file diff --git a/nnlib/__init__.py b/nnlib/__init__.py index 14793f7..60fc709 100644 --- a/nnlib/__init__.py +++ b/nnlib/__init__.py @@ -1 +1,2 @@ from .nnlib import nnlib +from .FUNIT import FUNIT \ No newline at end of file diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index b5caf30..2cc5dd1 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -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. @@ -687,7 +852,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator nnlib.Adam = Adam def CAInitializerMP( conv_weights_list ): - #Convolution Aware Initialization https://arxiv.org/abs/1702.06295 + #Convolution Aware Initialization https://arxiv.org/abs/1702.06295 data = [ (i, K.int_shape(conv_weights)) for i, conv_weights in enumerate(conv_weights_list) ] data = sorted(data, key=lambda data: np.prod(data[1]) ) result = CAInitializerMPSubprocessor (data, K.floatx(), K.image_data_format() ).run() @@ -814,8 +979,8 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator x = ReflectionPadding2D( self.pad ) (x) return self.func(x) nnlib.Conv2DTranspose = Conv2DTranspose - - class EqualConv2D(KL.Conv2D): + + class EqualConv2D(KL.Conv2D): def __init__(self, filters, kernel_size, strides=(1, 1), @@ -844,16 +1009,16 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator bias_constraint=None, **kwargs) self.gain = gain - + def build(self, input_shape): super().build(input_shape) - + self.wscale = self.gain / np.sqrt( np.prod( K.int_shape(self.kernel)[:-1]) ) self.wscale_t = K.constant (self.wscale, dtype=K.floatx() ) - + def call(self, inputs): k = self.kernel * self.wscale_t - + outputs = K.conv2d( inputs, k, @@ -872,12 +1037,12 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator return self.activation(outputs) return outputs nnlib.EqualConv2D = EqualConv2D - + class PixelNormalization(KL.Layer): # initialize the layer def __init__(self, **kwargs): super(PixelNormalization, self).__init__(**kwargs) - + # perform the operation def call(self, inputs): # calculate square pixel values @@ -891,12 +1056,12 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator # normalize values by the l2 norm normalized = inputs / l2 return normalized - + # define the output shape of the layer def compute_output_shape(self, input_shape): - return input_shape + return input_shape nnlib.PixelNormalization = PixelNormalization - + @staticmethod def import_keras_contrib(device_config): if nnlib.keras_contrib is not None: diff --git a/samplelib/Sample.py b/samplelib/Sample.py index 34f3a58..28f89ab 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -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(), diff --git a/samplelib/SampleGeneratorBase.py b/samplelib/SampleGeneratorBase.py index b89c506..cf98d8d 100644 --- a/samplelib/SampleGeneratorBase.py +++ b/samplelib/SampleGeneratorBase.py @@ -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 diff --git a/samplelib/SampleGeneratorFace.py b/samplelib/SampleGeneratorFace.py index 593f8e2..1eef85d 100644 --- a/samplelib/SampleGeneratorFace.py +++ b/samplelib/SampleGeneratorFace.py @@ -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 @@ -35,9 +47,10 @@ class SampleGeneratorFace(SampleGeneratorBase): raise ValueError("len(generators_random_seed) != generators_count") 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 @@ -49,7 +62,11 @@ class SampleGeneratorFace(SampleGeneratorBase): self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (i, samples[i::self.generators_count], ct_samples ) ) for i in range(self.generators_count) ] self.generator_counter = -1 - + + #overridable + def get_total_sample_count(self): + return self.total_samples_count + def __iter__(self): return self @@ -58,7 +75,7 @@ class SampleGeneratorFace(SampleGeneratorBase): generator = self.generators[self.generator_counter % len(self.generators) ] return next(generator) - def batch_func(self, param ): + def batch_func(self, param ): generator_id, samples, ct_samples = param if self.generators_random_seed is not None: @@ -82,7 +99,7 @@ class SampleGeneratorFace(SampleGeneratorBase): shuffle_idxs = [] shuffle_idxs_2D = [[]]*samples_len - while True: + while True: batches = None for n_batch in range(self.batch_size): while True: @@ -131,12 +148,24 @@ class SampleGeneratorFace(SampleGeneratorBase): if self.add_sample_idx: 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) \ No newline at end of file diff --git a/samplelib/SampleLoader.py b/samplelib/SampleLoader.py index 2d7d467..c86b0e3 100644 --- a/samplelib/SampleLoader.py +++ b/samplelib/SampleLoader.py @@ -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,10 +34,16 @@ 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: - datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] ) + 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: if datas[sample_type] is None: @@ -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': @@ -68,13 +78,13 @@ class SampleLoader: if dflimg is None: print ("%s is not a dfl image file required for training" % (s_filename_path.name) ) continue - + landmarks = dflimg.get_landmarks() pitch_yaw_roll = dflimg.get_pitch_yaw_roll() eyebrows_expand_mod = dflimg.get_eyebrows_expand_mod() - + if pitch_yaw_roll is None: - pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks) + pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks) sample_list.append( s.copy_and_set(sample_type=SampleType.FACE, face_type=FaceType.fromString (dflimg.get_face_type()),