diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 79ffa30..0f4d2d7 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -57,6 +57,9 @@ class SAEHDModel(ModelBase): default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True) default_random_downsample = self.options['random_downsample'] = self.load_or_def_option('random_downsample', False) + default_random_noise = self.options['random_noise'] = self.load_or_def_option('random_noise', False) + default_random_blur = self.options['random_blur'] = self.load_or_def_option('random_blur', False) + default_background_power = self.options['background_power'] = self.load_or_def_option('background_power', 0.0) default_true_face_power = self.options['true_face_power'] = self.load_or_def_option('true_face_power', 0.0) default_face_style_power = self.options['face_style_power'] = self.load_or_def_option('face_style_power', 0.0) @@ -161,8 +164,8 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.") self.options['random_downsample'] = io.input_bool("Enable random downsample of samples", default_random_downsample, help_message="") - # self.options['random_noise'] = io.input_bool("Enable random noise added to samples", False, help_message="") - # self.options['random_blur'] = io.input_bool("Enable random blur of samples", False, help_message="") + self.options['random_noise'] = io.input_bool("Enable random noise added to samples", default_random_noise, help_message="") + self.options['random_blur'] = io.input_bool("Enable random blur of samples", False, help_message="") # self.options['random_jpeg'] = io.input_bool("Enable random jpeg compression of samples", False, help_message="") self.options['gan_version'] = np.clip (io.input_int("GAN version", default_gan_version, add_info="2 or 3", help_message="Choose GAN version (v2: 7/16/2020, v3: 1/3/2021):"), 2, 3) @@ -751,7 +754,12 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.set_training_data_generators ([ SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_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, 'random_downsample': self.options['random_downsample'], 'transform':True, 'channel_type' : channel_type, 'ct_mode': ct_mode, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, + output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, + 'random_downsample': self.options['random_downsample'], + 'random_noise': self.options['random_noise'], + 'random_blur': self.options['random_blur'], + 'transform':True, 'channel_type' : channel_type, 'ct_mode': ct_mode, + 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : channel_type, 'ct_mode': ct_mode, '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, '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}, @@ -761,7 +769,12 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... 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, 'random_downsample': self.options['random_downsample'], 'transform':True, 'channel_type' : channel_type, 'ct_mode': fs_aug, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, + output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, + 'random_downsample': self.options['random_downsample'], + 'random_noise': self.options['random_noise'], + 'random_blur': self.options['random_blur'], + 'transform':True, 'channel_type' : channel_type, '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' : channel_type, '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, '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}, diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 1f1b96f..8278108 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -7,7 +7,7 @@ import numpy as np from core import imagelib from core.cv2ex import * -from core.imagelib import sd +from core.imagelib import sd, LinearMotionBlur from core.imagelib.color_transfer import random_lab_rotation from facelib import FaceType, LandmarksProcessor @@ -113,6 +113,8 @@ class SampleProcessor(object): warp = opts.get('warp', False) transform = opts.get('transform', False) random_downsample = opts.get('random_downsample', False) + random_noise = opts.get('random_noise', False) + random_blur = opts.get('random_blur', False) motion_blur = opts.get('motion_blur', None) gaussian_blur = opts.get('gaussian_blur', None) random_bilinear_resize = opts.get('random_bilinear_resize', None) @@ -220,6 +222,42 @@ class SampleProcessor(object): img = cv2.resize(img, (down_res, down_res), interpolation=cv2.INTER_CUBIC) img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_CUBIC) + # Apply random noise + if random_noise: + noise_type = np.random.choice(['gaussian', 'laplace', 'poisson']) + noise_scale = (20 * np.random.random() + 20) + + if noise_type == 'gaussian': + noise = np.random.normal(scale=noise_scale, size=img.shape) + img += noise / 255.0 + elif noise_type == 'laplace': + noise = np.random.laplace(scale=noise_scale, size=img.shape) + img += noise / 255.0 + elif noise_type == 'poisson': + noise_lam = (15 * np.random.random() + 15) + noise = np.random.poisson(lam=noise_lam, size=img.shape) + img += noise / 255.0 + + # Apply random blur + if random_blur: + blur_type = np.random.choice(['motion', 'gaussian']) + + if blur_type == 'motion': + blur_k = np.random.randint(10, 20) + blur_angle = 360 * np.random.random() + img = LinearMotionBlur(img, blur_k, blur_angle) + elif blur_type == 'gaussian': + blur_sigma = 5 * np.random.random() + 3 + + if blur_sigma < 5.0: + kernel_size = 2.9 * blur_sigma # 97% of weight + else: + kernel_size = 2.6 * blur_sigma # 95% of weight + kernel_size = int(kernel_size) + kernel_size = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size + + img = cv2.GaussianBlur(img, (kernel_size, kernel_size), blur_sigma) + img = imagelib.warp_by_params (params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate) img = np.clip(img.astype(np.float32), 0, 1)