fs augmentation

core/imagelib/__init__.py

@@ -12,7 +12,7 @@ from .warp import gen_warp_params, warp_by_params
 
 from .reduce_colors import reduce_colors
 
-from .color_transfer import color_transfer, color_transfer_mix, color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer
+from .color_transfer import color_transfer, color_transfer_mix, color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer, color_augmentation 
 
 from .common import normalize_channels, cut_odd_image, overlay_alpha_image
 

core/imagelib/color_transfer.py

@@ -366,3 +366,52 @@ def color_transfer(ct_mode, img_src, img_trg):
     else:
         raise ValueError(f"unknown ct_mode {ct_mode}")
     return out
+
+
+# imported from faceswap
+def color_augmentation(img):
+    """ Color adjust RGB image """
+    face = img
+    face = (face * 255.0).astype("uint8")
+    face = random_clahe(face)
+    face = random_lab(face)
+    img[:, :, :3] = face
+    return img.astype('float32') / 255.0
+
+
+def random_lab(image):
+    """ Perform random color/lightness adjustment in L*a*b* colorspace """
+    amount_l = 30 / 100
+    amount_ab = 8 / 100
+    randoms = [(random() * amount_l * 2) - amount_l,  # L adjust
+                (random() * amount_ab * 2) - amount_ab,  # A adjust
+                (random() * amount_ab * 2) - amount_ab]  # B adjust
+    image = cv2.cvtColor(  # pylint:disable=no-member
+    image, cv2.COLOR_BGR2LAB).astype("float32") / 255.0  # pylint:disable=no-member
+
+    for idx, adjustment in enumerate(randoms):
+        if adjustment >= 0:
+            image[:, :, idx] = ((1 - image[:, :, idx]) * adjustment) + image[:, :, idx]
+        else:
+            image[:, :, idx] = image[:, :, idx] * (1 + adjustment)
+    image = cv2.cvtColor((image * 255.0).astype("uint8"),  # pylint:disable=no-member
+                        cv2.COLOR_LAB2BGR)  # pylint:disable=no-member
+    return image
+
+def random_clahe(image):
+    """ Randomly perform Contrast Limited Adaptive Histogram Equalization """
+    contrast_random = random()
+    if contrast_random > 50 / 100:
+        return image
+
+    # base_contrast = image.shape[0] // 128
+    base_contrast = 1 # testing because it breaks on small sizes
+    grid_base = random() * 4
+    contrast_adjustment = int(grid_base * (base_contrast / 2))
+    grid_size = base_contrast + contrast_adjustment
+
+    clahe = cv2.createCLAHE(clipLimit=2.0,  # pylint: disable=no-member
+                            tileGridSize=(grid_size, grid_size))
+    for chan in range(3):
+        image[:, :, chan] = clahe.apply(image[:, :, chan])
+    return image

models/Model_SAEHD/Model.py

@@ -152,7 +152,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
             self.options['face_style_power'] = np.clip ( io.input_number("Face style power", default_face_style_power, add_info="0.0..100.0", help_message="Learn the color of the predicted face to be the same as dst inside mask. If you want to use this option with 'whole_face' you have to use XSeg trained mask. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.001 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )
             self.options['bg_style_power'] = np.clip ( io.input_number("Background style power", default_bg_style_power, add_info="0.0..100.0", help_message="Learn the area outside mask of the predicted face to be the same as dst. If you want to use this option with 'whole_face' you have to use XSeg trained mask. For whole_face you have to use XSeg trained mask. This can make face more like dst. Enabling this option increases the chance of model collapse. Typical value is 2.0"), 0.0, 100.0 )
 
-            self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", 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.")
+            self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot', 'fs-aug'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best. FS aug adds random color to dst and src")
             self.options['clipgrad'] = io.input_bool ("Enable gradient clipping", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.")
 
             self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly.")
@@ -615,10 +615,13 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                         uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
                         generators_count=src_generators_count ),
 
+                    fs_aug = None
+                    if ct_mode == 'fs-aug':
+                        fs_aug = 'fs-aug'
                     SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                         sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
-                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                                'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
+                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': fs_aug,                                                                 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
-                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                                'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
+                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': fs_aug,                                                                 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                 {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE_EYES, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                               ],
                         uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,

samplelib/SampleProcessor.py

@@ -197,7 +197,10 @@ class SampleProcessor(object):
                                 img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_CUBIC )
                                 
                         # Apply random color transfer                        
-                        if ct_mode is not None and ct_sample is not None:
+                        if ct_mode is not None and ct_sample is not None or ct_mode == 'fs-aug':
+                            if 'fs-aug':
+                                img = imagelib.color_augmentation(img)
+                            elif:
                                 if ct_sample_bgr is None:
                                     ct_sample_bgr = ct_sample.load_bgr()
                                 img = imagelib.color_transfer (ct_mode, img, cv2.resize( ct_sample_bgr, (resolution,resolution), interpolation=cv2.INTER_LINEAR ) )