SAE,SAEHD,Converter:

added sot-m color transfer

Converter:
removed seamless2 mode

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),
+                    prd_face_bgr = imagelib.reinhard_color_transfer ( (prd_face_bgr*255).astype(np.uint8),
-                                                                      np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255),
+                                                                      (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: #sot-m                    
-                elif cfg.color_transfer_mode == 7: #ebs, currently unused
+                    prd_face_bgr = imagelib.color_transfer_sot (prd_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a)
-                    prd_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8),
+                    prd_face_bgr = np.clip (prd_face_bgr, 0.0, 1.0)
-                                                     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)
-
             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':
+            out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT )
-                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 = 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),
+                    out_face_bgr = imagelib.reinhard_color_transfer ( (out_face_bgr*255).astype(np.uint8),
-                                                                      np.clip( (dst_face_bgr*255), 0, 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
+                elif cfg.color_transfer_mode == 7: #sot-m                    
-                    out_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8),
+                    out_face_bgr = imagelib.color_transfer_sot (out_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a)
-                                                     np.clip( (out_face_bgr*255), 0, 255).astype(np.uint8),  )
+                    out_face_bgr = np.clip (out_face_bgr, 0.0, 1.0)
-                    out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 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)

converters/ConverterConfig.py

@@ -103,13 +103,12 @@ class ConverterConfig(object):
 mode_dict = {0:'original',
              1:'overlay',
              2:'hist-match',
-             3:'seamless2',
+             3:'seamless',
-             4:'seamless',
+             4:'seamless-hist-match',
-             5:'seamless-hist-match',
+             5:'raw-rgb',
-             6:'raw-rgb',
+             6:'raw-rgb-mask',
-             7:'raw-rgb-mask',
+             7:'raw-mask-only',
-             8:'raw-mask-only',
+             8:'raw-predicted-only'}
-             9:'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_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 }
+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,
+                       erode_mask_modifier = 50,
-                       blur_mask_modifier = 0,
+                       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':

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
 

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
     

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())

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())

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())

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())

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

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,
                                     )
 

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,
                                     )
 

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)