Merge pull request #150 from faceshiftlabs/feat/image-degradation

Feat/image degradation

CHANGELOG.md

@@ -4,6 +4,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.7.2] - 2021-06-15
+### Added
+- New sample degradation options (only affects input, similar to random warp): 
+  - Random noise (gaussian/laplace/poisson)
+  - Random blur (gaussian/motion)
+  - Random jpeg compression
+  - Random downsampling
+- New "warped" preview(s): Shows the input samples with any/all distortions.     
+
 ## [1.7.1] - 2021-06-15
 ### Added
 - New autobackup options:

models/Model_SAEHD/Model.py

@@ -56,6 +56,11 @@ class SAEHDModel(ModelBase):
         default_loss_function      = self.options['loss_function']      = self.load_or_def_option('loss_function', 'SSIM')
 
         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_random_jpeg        = self.options['random_jpeg']        = self.load_or_def_option('random_jpeg', 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)
@@ -162,6 +167,11 @@ 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", default_random_noise, help_message="")
+            self.options['random_blur'] = io.input_bool("Enable random blur of samples", default_random_blur, help_message="")
+            self.options['random_jpeg'] = io.input_bool("Enable random jpeg compression of samples", default_random_jpeg, 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)
 
             if self.options['gan_version'] == 2:
@@ -755,7 +765,13 @@ 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=random_src_flip),
-                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, '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'],
+                                                 'random_jpeg': self.options['random_jpeg'],
+                                                 '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},
@@ -765,7 +781,13 @@ 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=random_dst_flip),
-                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, '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'],
+                                                 'random_jpeg': self.options['random_jpeg'],
+                                                 '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},
@@ -877,6 +899,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
           (warped_dst, target_dst, target_dstm, target_dstm_em) ) = samples
 
         S, D, SS, SSM, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
+        SW, DW = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([warped_src,warped_dst]) ]
         SSM, DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [SSM, DDM, SDM] ]
 
         target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]
@@ -892,6 +915,11 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                 st.append ( np.concatenate ( ar, axis=1) )
             result += [ ('SAEHD', np.concatenate (st, axis=0 )), ]
 
+            wt = []
+            for i in range(n_samples):
+                ar = SW[i], SS[i], DW[i], DD[i], SD[i]
+                wt.append ( np.concatenate ( ar, axis=1) )
+            result += [ ('SAEHD warped', np.concatenate (wt, axis=0 )), ]
 
             st_m = []
             for i in range(n_samples):
@@ -922,6 +950,23 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                 st.append ( np.concatenate ( ar, axis=1) )
             result += [ ('SAEHD pred', np.concatenate (st, axis=0 )), ]
 
+            wt = []
+            for i in range(n_samples):
+                ar = SW[i], SS[i]
+                wt.append ( np.concatenate ( ar, axis=1) )
+            result += [ ('SAEHD warped src-src', np.concatenate (wt, axis=0 )), ]
+
+            wt = []
+            for i in range(n_samples):
+                ar = DW[i], DD[i]
+                wt.append ( np.concatenate ( ar, axis=1) )
+            result += [ ('SAEHD warped dst-dst', np.concatenate (wt, axis=0 )), ]
+
+            wt = []
+            for i in range(n_samples):
+                ar = DW[i], SD[i]
+                wt.append ( np.concatenate ( ar, axis=1) )
+            result += [ ('SAEHD warped pred', np.concatenate (wt, axis=0 )), ]
 
             st_m = []
             for i in range(n_samples):

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
 
@@ -112,6 +112,10 @@ class SampleProcessor(object):
                 nearest_resize_to = opts.get('nearest_resize_to', None)
                 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)
+                random_jpeg = opts.get('random_jpeg', False)
                 motion_blur    = opts.get('motion_blur', None)
                 gaussian_blur  = opts.get('gaussian_blur', None)
                 random_bilinear_resize = opts.get('random_bilinear_resize', None)
@@ -214,6 +218,59 @@ class SampleProcessor(object):
                                 img = imagelib.color_transfer (ct_mode, img, cv2.resize( ct_sample_bgr, (resolution,resolution), interpolation=cv2.INTER_LINEAR ) )
 
 
+                        randomization_order = ['blur', 'noise', 'jpeg', 'down']
+                        np.random.shuffle(randomization_order)
+                        for random_distortion in randomization_order:
+                            # Apply random blur
+                            if random_distortion == 'blur' and 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)
+
+                            # Apply random noise
+                            if random_distortion == 'noise' and 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 jpeg compression
+                            if random_distortion == 'jpeg' and random_jpeg:
+                                img = np.clip(img*255, 0, 255).astype(np.uint8)
+                                jpeg_compression_level = np.random.randint(50, 85)
+                                encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_compression_level]
+                                _, enc_img = cv2.imencode('.jpg', img, encode_param)
+                                img = cv2.imdecode(enc_img, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.0
+
+                            # Apply random downsampling
+                            if random_distortion == 'down' and random_downsample:
+                                down_res = np.random.randint(int(0.125*resolution), int(0.25*resolution))
+                                img = cv2.resize(img, (down_res, down_res), interpolation=cv2.INTER_CUBIC)
+                                img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_CUBIC)
+
                         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)