diff --git a/converters/ConvertMasked.py b/converters/ConvertMasked.py index 9408458..2520c50 100644 --- a/converters/ConvertMasked.py +++ b/converters/ConvertMasked.py @@ -189,8 +189,8 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0: if cfg.color_transfer_mode == 1: #rct - prd_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (prd_face_bgr*255).astype(np.uint8), 0, 255), - np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), + prd_face_bgr = imagelib.reinhard_color_transfer ( (prd_face_bgr*255).astype(np.uint8), + (dst_face_bgr*255).astype(np.uint8), source_mask=prd_face_mask_a, target_mask=prd_face_mask_a) prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) @@ -205,12 +205,10 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 6: #idt-m prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) - - elif cfg.color_transfer_mode == 7: #ebs, currently unused - prd_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8), - np.clip( (prd_face_bgr*255), 0, 255).astype(np.uint8), )#prd_face_mask_a - prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) - + elif cfg.color_transfer_mode == 7: #sot-m + prd_face_bgr = imagelib.color_transfer_sot (prd_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) + prd_face_bgr = np.clip (prd_face_bgr, 0.0, 1.0) + if cfg.mode == 'hist-match-bw': prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY) prd_face_bgr = np.repeat( np.expand_dims (prd_face_bgr, -1), (3,), -1 ) @@ -238,9 +236,6 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i #mask used for cv2.seamlessClone img_face_mask_a = img_face_mask_aaa[...,0:1] - if cfg.mode == 'seamless2': - img_face_mask_a = cv2.warpAffine( img_face_mask_a, face_output_mat, (output_size, output_size), flags=cv2.INTER_CUBIC ) - img_face_seamless_mask_a = None for i in range(1,10): a = img_face_mask_a > i / 10.0 @@ -251,15 +246,11 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i img_face_seamless_mask_a[img_face_seamless_mask_a <= i / 10.0] = 0.0 break - if cfg.mode == 'seamless2': - face_seamless = imagelib.seamless_clone ( prd_face_bgr, dst_face_bgr, img_face_seamless_mask_a ) - out_img = cv2.warpAffine( face_seamless, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) - else: - out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) out_img = np.clip(out_img, 0.0, 1.0) - if 'seamless' in cfg.mode and cfg.mode != 'seamless2': + if 'seamless' in cfg.mode: try: #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering) l,t,w,h = cv2.boundingRect( (img_face_seamless_mask_a*255).astype(np.uint8) ) @@ -284,8 +275,8 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i if cfg.color_transfer_mode == 1: face_mask_aaa = cv2.warpAffine( img_face_mask_aaa, face_mat, (output_size, output_size) ) - out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255), 0, 255).astype(np.uint8), - np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8), + out_face_bgr = imagelib.reinhard_color_transfer ( (out_face_bgr*255).astype(np.uint8), + (dst_face_bgr*255).astype(np.uint8), source_mask=face_mask_aaa, target_mask=face_mask_aaa) out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) elif cfg.color_transfer_mode == 2: #lct @@ -299,10 +290,9 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i out_face_bgr = imagelib.color_transfer_idt (out_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 6: #idt-m out_face_bgr = imagelib.color_transfer_idt (out_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) - elif cfg.color_transfer_mode == 7: #ebs - out_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8), - np.clip( (out_face_bgr*255), 0, 255).astype(np.uint8), ) - out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + elif cfg.color_transfer_mode == 7: #sot-m + out_face_bgr = imagelib.color_transfer_sot (out_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) + out_face_bgr = np.clip (out_face_bgr, 0.0, 1.0) if cfg.mode == 'seamless-hist-match': out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, cfg.hist_match_threshold) diff --git a/converters/ConverterConfig.py b/converters/ConverterConfig.py index e3084d3..4993789 100644 --- a/converters/ConverterConfig.py +++ b/converters/ConverterConfig.py @@ -103,13 +103,12 @@ class ConverterConfig(object): mode_dict = {0:'original', 1:'overlay', 2:'hist-match', - 3:'seamless2', - 4:'seamless', - 5:'seamless-hist-match', - 6:'raw-rgb', - 7:'raw-rgb-mask', - 8:'raw-mask-only', - 9:'raw-predicted-only'} + 3:'seamless', + 4:'seamless-hist-match', + 5:'raw-rgb', + 6:'raw-rgb-mask', + 7:'raw-mask-only', + 8:'raw-predicted-only'} full_face_mask_mode_dict = {1:'learned', 2:'dst', @@ -123,24 +122,24 @@ half_face_mask_mode_dict = {1:'learned', 4:'FAN-dst', 7:'learned*FAN-dst'} -ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m" } -ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6 } +ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m", 7:"sot-m" } +ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6, "sot-m":7 } class ConverterConfigMasked(ConverterConfig): def __init__(self, face_type=FaceType.FULL, - default_mode = 4, + default_mode = 'overlay', clip_hborder_mask_per = 0, mode='overlay', masked_hist_match=True, hist_match_threshold = 238, mask_mode = 1, - erode_mask_modifier = 0, - blur_mask_modifier = 0, + erode_mask_modifier = 50, + blur_mask_modifier = 50, motion_blur_power = 0, output_face_scale = 0, - color_transfer_mode = 0, + color_transfer_mode = ctm_str_dict['rct'], image_denoise_power = 0, bicubic_degrade_power = 0, color_degrade_power = 0, @@ -176,7 +175,7 @@ class ConverterConfigMasked(ConverterConfig): return copy.copy(self) def set_mode (self, mode): - self.mode = mode_dict.get (mode, mode_dict[self.default_mode] ) + self.mode = mode_dict.get (mode, self.default_mode) def toggle_masked_hist_match(self): if self.mode == 'hist-match' or self.mode == 'hist-match-bw': diff --git a/imagelib/__init__.py b/imagelib/__init__.py index 2ab6687..43fad64 100644 --- a/imagelib/__init__.py +++ b/imagelib/__init__.py @@ -11,7 +11,7 @@ from .warp import gen_warp_params, warp_by_params from .reduce_colors import reduce_colors -from .color_transfer import color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer, seamless_clone +from .color_transfer import color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer, seamless_clone from .RankSRGAN import RankSRGAN diff --git a/imagelib/color_transfer.py b/imagelib/color_transfer.py index cb46e9c..37f0af9 100644 --- a/imagelib/color_transfer.py +++ b/imagelib/color_transfer.py @@ -1,10 +1,62 @@ import cv2 import numpy as np +from numpy import linalg as npla import scipy as sp import scipy.sparse from scipy.sparse.linalg import spsolve +def color_transfer_sot(src,trg, steps=10, batch_size=5, reg_sigmaXY=16.0, reg_sigmaV=5.0): + """ + Color Transform via Sliced Optimal Transfer + ported by @iperov from https://github.com/dcoeurjo/OTColorTransfer + + src - any float range any channel image + dst - any float range any channel image, same shape as src + steps - number of solver steps + batch_size - solver batch size + reg_sigmaXY - apply regularization and sigmaXY of filter, otherwise set to 0.0 + reg_sigmaV - sigmaV of filter + + return value - clip it manually + """ + if not np.issubdtype(src.dtype, np.floating): + raise ValueError("src value must be float") + if not np.issubdtype(trg.dtype, np.floating): + raise ValueError("trg value must be float") + + if len(src.shape) != 3: + raise ValueError("src shape must have rank 3 (h,w,c)") + + if src.shape != trg.shape: + raise ValueError("src and trg shapes must be equal") + + src_dtype = src.dtype + h,w,c = src.shape + new_src = src.copy() + + for step in range (steps): + advect = np.zeros ( (h*w,c), dtype=src_dtype ) + for batch in range (batch_size): + dir = np.random.normal(size=c).astype(src_dtype) + dir /= npla.norm(dir) + + projsource = np.sum( new_src*dir, axis=-1).reshape ((h*w)) + projtarget = np.sum( trg*dir, axis=-1).reshape ((h*w)) + + idSource = np.argsort (projsource) + idTarget = np.argsort (projtarget) + + a = projtarget[idTarget]-projsource[idSource] + for i_c in range(c): + advect[idSource,i_c] += a * dir[i_c] + new_src += advect.reshape( (h,w,c) ) / batch_size + + if reg_sigmaXY != 0.0: + src_diff = new_src-src + new_src = src + cv2.bilateralFilter (src_diff, 0, reg_sigmaV, reg_sigmaXY ) + return new_src + def color_transfer_mkl(x0, x1): eps = np.finfo(float).eps diff --git a/mainscripts/gfx/help_converter_masked.jpg b/mainscripts/gfx/help_converter_masked.jpg index 41e9947..e737f54 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 fb2734a..ebb3b0d 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/Model_DF/Model.py b/models/Model_DF/Model.py index c0b55ac..9120951 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -121,7 +121,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode=4) + return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode='seamless') def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 280389a..90e468e 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -129,7 +129,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode=4) + return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode='seamless') def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index d288206..78eb79b 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -130,7 +130,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return self.predictor_func, (64,64,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode=4) + return self.predictor_func, (64,64,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode='seamless') def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index a908677..c6a6dbe 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -127,7 +127,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode=4) + return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode='seamless') def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_Quick96/Model.py b/models/Model_Quick96/Model.py index 42e508c..51d675b 100644 --- a/models/Model_Quick96/Model.py +++ b/models/Model_Quick96/Model.py @@ -257,6 +257,6 @@ class Quick96Model(ModelBase): def get_ConverterConfig(self): import converters return self.predictor_func, (self.resolution, self.resolution, 3), converters.ConverterConfigMasked(face_type=FaceType.FULL, - default_mode = 1, clip_hborder_mask_per=0.0625) + default_mode='seamless', clip_hborder_mask_per=0.0625) Model = Quick96Model diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 9b8cf51..90e8f92 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -83,7 +83,7 @@ class SAEModel(ModelBase): help_message="Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) default_ct_mode = self.options.get('ct_mode', 'none') - self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") + self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt/sot, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301 default_clipgrad = False if is_first_run else self.options.get('clipgrad', False) @@ -564,7 +564,7 @@ class SAEModel(ModelBase): import converters return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type, - default_mode = 1 if self.options['ct_mode'] != 'none' or self.options['face_style_power'] or self.options['bg_style_power'] else 4, + default_mode = 'overlay' if self.options['ct_mode'] != 'none' or self.options['face_style_power'] or self.options['bg_style_power'] else 'seamless', clip_hborder_mask_per=0.0625 if (self.options['face_type'] == 'f') else 0, ) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index f25e2f7..1ee851f 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -77,7 +77,7 @@ class SAEHDModel(ModelBase): help_message="Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) default_ct_mode = self.options.get('ct_mode', 'none') - self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") + self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt/sot, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301 default_clipgrad = False if is_first_run else self.options.get('clipgrad', False) @@ -639,7 +639,7 @@ class SAEHDModel(ModelBase): import converters return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type, - default_mode = 1 if self.options['ct_mode'] != 'none' or self.options['face_style_power'] or self.options['bg_style_power'] else 4, + default_mode = 'overlay' if self.options['ct_mode'] != 'none' or self.options['face_style_power'] or self.options['bg_style_power'] else 'seamless', clip_hborder_mask_per=0.0625 if (face_type != FaceType.HALF) else 0, ) diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index c915076..609a39c 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -261,6 +261,9 @@ class SampleProcessor(object): img_bgr = imagelib.color_transfer_mkl (img_bgr, ct_sample_bgr_resized) elif ct_mode == 'idt': img_bgr = imagelib.color_transfer_idt (img_bgr, ct_sample_bgr_resized) + elif ct_mode == 'sot': + img_bgr = imagelib.color_transfer_sot (img_bgr, ct_sample_bgr_resized) + img_bgr = np.clip( img_bgr, 0.0, 1.0) if random_hsv_shift: rnd_state = np.random.RandomState (sample_rnd_seed)