diff --git a/core/imagelib/__init__.py b/core/imagelib/__init__.py index 4e8b83d..e74d427 100644 --- a/core/imagelib/__init__.py +++ b/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 diff --git a/core/imagelib/color_transfer.py b/core/imagelib/color_transfer.py index 883cecb..59be668 100644 --- a/core/imagelib/color_transfer.py +++ b/core/imagelib/color_transfer.py @@ -365,4 +365,53 @@ def color_transfer(ct_mode, img_src, img_trg): out = np.clip( out, 0.0, 1.0) else: raise ValueError(f"unknown ct_mode {ct_mode}") - return out \ No newline at end of file + 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 \ No newline at end of file diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index aa2bc2f..fe61f7f 100644 --- a/models/Model_SAEHD/Model.py +++ b/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}, - {'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}, + 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, '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, diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index a951df4..38e64dd 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -197,10 +197,13 @@ 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_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 ) ) + 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 ) ) img = imagelib.warp_by_params (params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate)