moving some files

converters/ConvertMasked.py

@@ -56,14 +56,14 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i
 
     if cfg.mask_mode == 2: #dst
         prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
-    elif cfg.mask_mode >= 3 and cfg.mask_mode <= 7:
+    elif cfg.mask_mode >= 3 and cfg.mask_mode <= 8:
 
         if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6:
             prd_face_fanseg_bgr = cv2.resize (prd_face_bgr, (cfg.fanseg_input_size,)*2 )
             prd_face_fanseg_mask = cfg.fanseg_extract_func(FaceType.FULL, prd_face_fanseg_bgr)
             FAN_prd_face_mask_a_0 = cv2.resize ( prd_face_fanseg_mask, (output_size, output_size), cv2.INTER_CUBIC)
 
-        if cfg.mask_mode >= 4 or cfg.mask_mode <= 7:
+        if cfg.mask_mode >= 4 and cfg.mask_mode <= 7:
 
             full_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=FaceType.FULL)
             dst_face_fanseg_bgr = cv2.warpAffine(img_bgr, full_face_fanseg_mat, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
@@ -80,9 +80,24 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i
                 m = cv2.getAffineTransform(b, fanseg_rect_corner_pts)
                 FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanseg_mask, m, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
                 FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
-            #else:
+        """
-            #    raise ValueError ("cfg.face_type unsupported")
+        if cfg.mask_mode == 8: #FANCHQ-dst
+            full_face_fanchq_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanchq_input_size, face_type=FaceType.FULL)
+            dst_face_fanchq_bgr = cv2.warpAffine(img_bgr, full_face_fanchq_mat, (cfg.fanchq_input_size,)*2, flags=cv2.INTER_CUBIC )
+            dst_face_fanchq_mask = cfg.fanchq_extract_func( FaceType.FULL, dst_face_fanchq_bgr )
+            
+            if cfg.face_type == FaceType.FULL:
+                FANCHQ_dst_face_mask_a_0 = cv2.resize (dst_face_fanchq_mask, (output_size,output_size), cv2.INTER_CUBIC)
+            else:
+                face_fanchq_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanchq_input_size, face_type=cfg.face_type)
 
+                fanchq_rect_corner_pts = np.array ( [ [0,0], [cfg.fanchq_input_size-1,0], [0,cfg.fanchq_input_size-1] ], dtype=np.float32 )
+                a = LandmarksProcessor.transform_points (fanchq_rect_corner_pts, face_fanchq_mat, invert=True )
+                b = LandmarksProcessor.transform_points (a, full_face_fanchq_mat )
+                m = cv2.getAffineTransform(b, fanchq_rect_corner_pts)
+                FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanchq_mask, m, (cfg.fanchq_input_size,)*2, flags=cv2.INTER_CUBIC )
+                FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
+        """
         if cfg.mask_mode == 3:   #FAN-prd
             prd_face_mask_a_0 = FAN_prd_face_mask_a_0
         elif cfg.mask_mode == 4: #FAN-dst
@@ -93,7 +108,9 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i
             prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
         elif cfg.mask_mode == 7:
             prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0
-
+        #elif cfg.mask_mode == 8: #FANCHQ-dst
+        #    prd_face_mask_a_0 = FANCHQ_dst_face_mask_a_0
+            
     prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0
 
     prd_face_mask_a   = prd_face_mask_a_0[...,np.newaxis]

converters/ConverterConfig.py

@@ -21,6 +21,9 @@ class ConverterConfig(object):
         self.blursharpen_func = None
         self.fanseg_input_size = None
         self.fanseg_extract_func = None
+        
+        self.fanchq_input_size = None
+        self.fanchq_extract_func = None
         self.ebs_ct_func = None
 
         self.super_res_dict = {0:"None", 1:'RankSRGAN'}

facelib/__init__.py

@@ -3,5 +3,4 @@ from .DLIBExtractor import DLIBExtractor
 from .MTCExtractor import MTCExtractor
 from .S3FDExtractor import S3FDExtractor
 from .FANExtractor import FANExtractor
-from .FANSegmentator import FANSegmentator
 from .PoseEstimator import PoseEstimator

main.py

@@ -78,6 +78,16 @@ if __name__ == "__main__":
     p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
     p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU.")
     p.set_defaults (func=process_dev_extract_umd_csv)
+    
+    
+    def process_dev_apply_celebamaskhq(arguments):
+        os_utils.set_process_lowest_prio()
+        from mainscripts import dev_misc
+        dev_misc.apply_celebamaskhq( arguments.input_dir )
+
+    p = subparsers.add_parser( "dev_apply_celebamaskhq", help="")
+    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
+    p.set_defaults (func=process_dev_apply_celebamaskhq)
     """
     def process_extract_fanseg(arguments):
         os_utils.set_process_lowest_prio()

mainscripts/Converter.py

@@ -16,7 +16,9 @@ import numpy.linalg as npla
 import imagelib
 from converters import (ConverterConfig, ConvertFaceAvatar, ConvertMasked,
                         FrameInfo)
-from facelib import FaceType, FANSegmentator, LandmarksProcessor
+from facelib import FaceType, LandmarksProcessor
+from nnlib import TernausNet
+
 from interact import interact as io
 from joblib import SubprocessFunctionCaller, Subprocessor
 from utils import Path_utils
@@ -114,13 +116,25 @@ class ConvertSubprocessor(Subprocessor):
             def fanseg_extract(face_type, *args, **kwargs):
                 fanseg = self.fanseg_by_face_type.get(face_type, None)
                 if self.fanseg_by_face_type.get(face_type, None) is None:
-                    fanseg = FANSegmentator( self.fanseg_input_size , FaceType.toString( face_type ) )
+                    fanseg = TernausNet("FANSeg", self.fanseg_input_size , FaceType.toString( face_type ) )
                     self.fanseg_by_face_type[face_type] = fanseg
 
                 return fanseg.extract(*args, **kwargs)
 
             self.fanseg_extract_func = fanseg_extract
             
+            self.fanchq_by_face_type = {}
+            self.fanchq_input_size = 256
+            def fanchq_extract(face_type, *args, **kwargs):
+                fanchq = self.fanchq_by_face_type.get(face_type, None)
+                if self.fanchq_by_face_type.get(face_type, None) is None:
+                    fanchq = TernausNet("FANCHQ", self.fanchq_input_size , FaceType.toString( face_type ) )
+                    self.fanchq_by_face_type[face_type] = fanchq
+
+                return fanchq.extract(*args, **kwargs)
+
+            self.fanchq_extract_func = fanchq_extract
+            
             import ebsynth
             def ebs_ct(*args, **kwargs):                    
                 return ebsynth.color_transfer(*args, **kwargs)
@@ -161,6 +175,8 @@ class ConvertSubprocessor(Subprocessor):
                 if cfg.type == ConverterConfig.TYPE_MASKED:
                     cfg.fanseg_input_size = self.fanseg_input_size
                     cfg.fanseg_extract_func = self.fanseg_extract_func
+                    cfg.fanchq_input_size = self.fanchq_input_size
+                    cfg.fanchq_extract_func = self.fanchq_extract_func
 
                     try:
                         final_img = ConvertMasked (self.predictor_func, self.predictor_input_shape, cfg, frame_info)

mainscripts/Extractor.py

@@ -14,10 +14,10 @@ import numpy as np
 import facelib
 import imagelib
 import mathlib
-from facelib import FaceType, FANSegmentator, LandmarksProcessor
+from facelib import FaceType, LandmarksProcessor
 from interact import interact as io
 from joblib import Subprocessor
-from nnlib import nnlib
+from nnlib import TernausNet, nnlib
 from utils import Path_utils
 from utils.cv2_utils import *
 from utils.DFLJPG import DFLJPG
@@ -95,7 +95,7 @@ class ExtractSubprocessor(Subprocessor):
                     
             elif self.type == 'fanseg':
                 nnlib.import_all (device_config)
-                self.e = facelib.FANSegmentator(256, FaceType.toString(FaceType.FULL) )
+                self.e = TernausNet(256, FaceType.toString(FaceType.FULL) )
                 self.e.__enter__()
                     
             elif self.type == 'final':
@@ -279,7 +279,6 @@ class ExtractSubprocessor(Subprocessor):
                     fanseg_mask = self.e.extract( image / 255.0 )
                     src_dflimg.embed_and_set( filename_path_str, 
                                               fanseg_mask=fanseg_mask,
-                                              #fanseg_mask_ver=FANSegmentator.VERSION,
                                               )
         
         #overridable
@@ -889,4 +888,3 @@ def extract_umd_csv(input_file_csv,
     io.log_info ('Images found:        %d' % (images_found) )
     io.log_info ('Faces detected:      %d' % (faces_detected) )
     io.log_info ('-------------------------')
-

mainscripts/dev_misc.py

@@ -1,9 +1,20 @@
-from . import Extractor
+import multiprocessing
-from . import Sorter
-from pathlib import Path
-from utils import Path_utils
 import shutil
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+from facelib import LandmarksProcessor
 from interact import interact as io
+from joblib import Subprocessor
+from utils import Path_utils
+from utils.cv2_utils import *
+from utils.DFLJPG import DFLJPG
+from utils.DFLPNG import DFLPNG
+
+from . import Extractor, Sorter
+
 
 def extract_vggface2_dataset(input_dir, device_args={} ):
     multi_gpu = device_args.get('multi_gpu', False)
@@ -47,4 +58,169 @@ def extract_vggface2_dataset(input_dir, device_args={} ):
             io.log_info (f"Removing: {str(cur_input_path)} ")
             shutil.rmtree(cur_input_path)
         except:
-            io.log_info (f"unable to remove: {str(cur_input_path)} ")
+            io.log_info (f"unable to remove: {str(cur_input_path)} ")
+            
+            
+
+class CelebAMASKHQSubprocessor(Subprocessor):
+    class Cli(Subprocessor.Cli):
+        #override
+        def on_initialize(self, client_dict):
+            self.masks_files_paths = client_dict['masks_files_paths']
+            return None
+
+        #override
+        def process_data(self, data):
+            filename = data[0]
+            filepath = Path(filename)
+            if filepath.suffix == '.png':
+                dflimg = DFLPNG.load( str(filepath) )
+            elif filepath.suffix == '.jpg':
+                dflimg = DFLJPG.load ( str(filepath) )
+            else:
+                dflimg = None
+                
+            image_to_face_mat = dflimg.get_image_to_face_mat()    
+            src_filename = dflimg.get_source_filename()
+
+            img = cv2_imread(filename)
+            h,w,c = img.shape
+            
+            fanseg_mask = LandmarksProcessor.get_image_hull_mask(img.shape, dflimg.get_landmarks() )
+            
+            idx_name = '%.5d' % int(src_filename.split('.')[0])
+            idx_files = [ x for x in self.masks_files_paths if idx_name in x ]  
+                
+            skin_files = [ x for x in idx_files if 'skin' in x ]
+            eye_glass_files = [ x for x in idx_files if 'eye_g' in x ]
+            
+            for files, is_invert in [ (skin_files,False), 
+                                      (eye_glass_files,True) ]:
+                if len(files) > 0:
+                    mask = cv2_imread(files[0])            
+                    mask = mask[...,0]
+                    mask[mask == 255] = 1
+                    mask = mask.astype(np.float32)
+                    mask = cv2.resize(mask, (1024,1024) )                
+                    mask = cv2.warpAffine(mask, image_to_face_mat, (w, h), cv2.INTER_LANCZOS4)
+                    
+                    if not is_invert:
+                        fanseg_mask *= mask[...,None]
+                    else:
+                        fanseg_mask *= (1-mask[...,None])
+
+            dflimg.embed_and_set (filename, fanseg_mask=fanseg_mask)
+            return 1
+
+        #override
+        def get_data_name (self, data):
+            #return string identificator of your data
+            return data[0]
+
+    #override
+    def __init__(self, image_paths, masks_files_paths ):
+        self.image_paths = image_paths
+        self.masks_files_paths = masks_files_paths
+        
+        self.result = []
+        super().__init__('CelebAMASKHQSubprocessor', CelebAMASKHQSubprocessor.Cli, 60)
+
+    #override
+    def process_info_generator(self):
+        for i in range(min(multiprocessing.cpu_count(), 8)):
+            yield 'CPU%d' % (i), {}, {'masks_files_paths' : self.masks_files_paths }
+
+    #override
+    def on_clients_initialized(self):
+        io.progress_bar ("Processing", len (self.image_paths))
+
+    #override
+    def on_clients_finalized(self):
+        io.progress_bar_close()
+
+    #override
+    def get_data(self, host_dict):
+        if len (self.image_paths) > 0:
+            return [self.image_paths.pop(0)]
+        return None
+
+    #override
+    def on_data_return (self, host_dict, data):
+        self.image_paths.insert(0, data[0])
+
+    #override
+    def on_result (self, host_dict, data, result):
+        io.progress_bar_inc(1)
+
+    #override
+    def get_result(self):
+        return self.result
+        
+#unused in end user workflow
+def apply_celebamaskhq(input_dir ):
+    
+    input_path = Path(input_dir)    
+    
+    img_path = input_path / 'aligned'
+    mask_path = input_path / 'mask'
+
+    if not img_path.exists():
+        raise ValueError(f'{str(img_path)} directory not found. Please ensure it exists.')
+
+    CelebAMASKHQSubprocessor(Path_utils.get_image_paths(img_path), 
+                             Path_utils.get_image_paths(mask_path, subdirs=True) ).run()
+    
+    return
+    
+    paths_to_extract = []
+    for filename in io.progress_bar_generator(Path_utils.get_image_paths(img_path), desc="Processing"):
+        filepath = Path(filename)
+        if filepath.suffix == '.png':
+            dflimg = DFLPNG.load( str(filepath) )
+        elif filepath.suffix == '.jpg':
+            dflimg = DFLJPG.load ( str(filepath) )
+        else:
+            dflimg = None
+
+        if dflimg is not None:
+            paths_to_extract.append (filepath)
+            
+        image_to_face_mat = dflimg.get_image_to_face_mat()    
+        src_filename = dflimg.get_source_filename()
+
+        #img = cv2_imread(filename)
+        h,w,c = dflimg.get_shape()
+        
+        fanseg_mask = LandmarksProcessor.get_image_hull_mask( (h,w,c), dflimg.get_landmarks() )
+        
+        idx_name = '%.5d' % int(src_filename.split('.')[0])
+        idx_files = [ x for x in masks_files if idx_name in x ]  
+              
+        skin_files = [ x for x in idx_files if 'skin' in x ]
+        eye_glass_files = [ x for x in idx_files if 'eye_g' in x ]
+        
+        for files, is_invert in [ (skin_files,False), 
+                                  (eye_glass_files,True) ]:
+                       
+            if len(files) > 0:
+                mask = cv2_imread(files[0])            
+                mask = mask[...,0]
+                mask[mask == 255] = 1
+                mask = mask.astype(np.float32)
+                mask = cv2.resize(mask, (1024,1024) )                
+                mask = cv2.warpAffine(mask, image_to_face_mat, (w, h), cv2.INTER_LANCZOS4)
+                
+                if not is_invert:
+                    fanseg_mask *= mask[...,None]
+                else:
+                    fanseg_mask *= (1-mask[...,None])
+            
+        #cv2.imshow("", (fanseg_mask*255).astype(np.uint8) )
+        #cv2.waitKey(0)    
+        
+        
+        dflimg.embed_and_set (filename, fanseg_mask=fanseg_mask)
+        
+        
+        #import code
+        #code.interact(local=dict(globals(), **locals()))

models/Model_DEV_FANSEG/Model.py

@@ -1,9 +1,8 @@
 import numpy as np
 
-from nnlib import nnlib
+from nnlib import nnlib, TernausNet
 from models import ModelBase
 from facelib import FaceType
-from facelib import FANSegmentator
 from samplelib import *
 from interact import interact as io
 
@@ -29,17 +28,18 @@ class Model(ModelBase):
     #override
     def onInitialize(self):
         exec(nnlib.import_all(), locals(), globals())
-        self.set_vram_batch_requirements( {1.5:4} )
+        self.set_vram_batch_requirements( {1.5:4, 11:48} )
 
         self.resolution = 256
         self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
 
-        
+        model_name = 'FANSeg'
-        self.fan_seg = FANSegmentator(self.resolution, 
+        model_name = 'FANCHQ'
-                                      FaceType.toString(self.face_type), 
+        self.fan_seg = TernausNet(model_name, self.resolution, 
-                                      load_weights=not self.is_first_run(),
+                                            FaceType.toString(self.face_type), 
-                                      weights_file_root=self.get_model_root_path(),
+                                            load_weights=not self.is_first_run(),
-                                      training=True)
+                                            weights_file_root=self.get_model_root_path(),
+                                            training=True)
 
         if self.is_training_mode:
             t = SampleProcessor.Types
@@ -48,13 +48,13 @@ class Model(ModelBase):
             self.set_training_data_generators ([    
                     SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, 
                             sample_process_options=SampleProcessor.Options(random_flip=True), 
-                            output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'border_replicate':False },
+                            output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution' : self.resolution, 'motion_blur':(25, 5),  'gaussian_blur':(25,5), 'border_replicate':False, 'random_hsv_shift' : True },
                                                   { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M), 'resolution': self.resolution },
                                                 ]),
                                                 
                     SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, 
                             sample_process_options=SampleProcessor.Options(random_flip=True ), 
-                            output_sample_types=[ { 'types': (t.IMG_TRANSFORMED , face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution},
+                            output_sample_types=[ { 'types': (t.IMG_TRANSFORMED , face_type, t.MODE_BGR), 'resolution' : self.resolution, 'random_hsv_shift' : True},
                                                 ])
                                                ])
                 
@@ -73,11 +73,14 @@ class Model(ModelBase):
     #override
     def onGetPreview(self, sample):
         test_A      = sample[0][0][0:4] #first 4 samples
+        test_Am     = sample[0][1][0:4] #first 4 samples
         test_B      = sample[1][0][0:4] #first 4 samples
         
+        
         mAA = self.fan_seg.extract(test_A)
         mBB = self.fan_seg.extract(test_B)
 
+        test_Am = np.repeat ( test_Am, (3,), -1)
         mAA = np.repeat ( mAA, (3,), -1)
         mBB = np.repeat ( mBB, (3,), -1)
         
@@ -85,6 +88,7 @@ class Model(ModelBase):
         for i in range(0, len(test_A)):
             st.append ( np.concatenate ( (
                 test_A[i,:,:,0:3],
+                test_Am[i],
                 mAA[i],
                 test_A[i,:,:,0:3]*mAA[i],
                 ), axis=1) )
@@ -92,7 +96,7 @@ class Model(ModelBase):
         st2 = []
         for i in range(0, len(test_B)):
             st2.append ( np.concatenate ( (
-                test_B[i,:,:,0:3],
+                test_B[i,:,:,0:3],                
                 mBB[i],
                 test_B[i,:,:,0:3]*mBB[i],
                 ), axis=1) )

nnlib/TernausNet.py

@@ -10,8 +10,6 @@ from interact import interact as io
 from nnlib import nnlib
 
 """
-FANSegmentator is designed to exclude obstructions from faces such as hair, fingers, etc.
-
 Dataset used to train located in official DFL mega.nz folder
 https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg
 
@@ -19,19 +17,19 @@ using https://github.com/ternaus/TernausNet
 TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation
 """
 
-class FANSegmentator(object):
+class TernausNet(object):
     VERSION = 1
-    def __init__ (self, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False):
+    def __init__ (self, name, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False):
         exec( nnlib.import_all(), locals(), globals() )
 
-        self.model = FANSegmentator.BuildModel(resolution, ngf=64)
+        self.model = TernausNet.BuildModel(resolution, ngf=64)
 
         if weights_file_root is not None:
             weights_file_root = Path(weights_file_root)
         else:
             weights_file_root = Path(__file__).parent
 
-        self.weights_path = weights_file_root / ('FANSeg_%d_%s.h5' % (resolution, face_type_str) )
+        self.weights_path = weights_file_root / ('%s_%d_%s.h5' % (name, resolution, face_type_str) )
 
         if load_weights:
             self.model.load_weights (str(self.weights_path))
@@ -59,7 +57,6 @@ class FANSegmentator(object):
             real_t = Input ( (resolution, resolution, 1) )
             out_t = self.model(inp_t)
 
-            #loss = K.mean(10*K.square(out_t-real_t))
             loss = K.mean(10*K.binary_crossentropy(real_t,out_t) )
 
             out_t_diff1 = out_t[:, 1:, :, :] - out_t[:, :-1, :, :]
@@ -69,7 +66,7 @@ class FANSegmentator(object):
 
             opt = Adam(lr=0.0001, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)
 
-            self.train_func = K.function  ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss,total_var_loss], self.model.trainable_weights) )
+            self.train_func = K.function  ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss], self.model.trainable_weights) )
 
 
     def __enter__(self):
@@ -103,7 +100,7 @@ class FANSegmentator(object):
         exec( nnlib.import_all(), locals(), globals() )
         inp = Input ( (resolution,resolution,3) )
         x = inp
-        x = FANSegmentator.Flow(ngf=ngf)(x)
+        x = TernausNet.Flow(ngf=ngf)(x)
         model = Model(inp,x)
         return model
 
@@ -115,7 +112,7 @@ class FANSegmentator(object):
             x = input
 
             x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x)
-            x = BlurPool(filt_size=3)(x) #x = MaxPooling2D()(x)
+            x = BlurPool(filt_size=3)(x)
 
             x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x)
             x = BlurPool(filt_size=3)(x)

nnlib/__init__.py

@@ -1,2 +1,3 @@
 from .nnlib import nnlib
-from .FUNIT import FUNIT
+from .FUNIT import FUNIT
+from .TernausNet import TernausNet

samplelib/SampleProcessor.py

@@ -72,7 +72,7 @@ class SampleProcessor(object):
         MODE_GGG                   = 42  #3xGrayscale
         MODE_M                     = 43  #mask only
         MODE_BGR_SHUFFLE           = 44  #BGR shuffle
-        MODE_BGR_RANDOM_HUE_SHIFT  = 45
+        MODE_BGR_RANDOM_HSV_SHIFT  = 45
         MODE_END = 50
 
     class Options(object):
@@ -111,8 +111,6 @@ class SampleProcessor(object):
 
         sample_rnd_seed = np.random.randint(0x80000000)
 
-
-
         outputs = []
         for opts in output_sample_types:
 
@@ -124,6 +122,9 @@ class SampleProcessor(object):
             normalize_std_dev = opts.get('normalize_std_dev', False)
             normalize_vgg = opts.get('normalize_vgg', False)
             motion_blur = opts.get('motion_blur', None)
+            gaussian_blur = opts.get('gaussian_blur', None)
+            
+            random_hsv_shift = opts.get('random_hsv_shift', None)
             apply_ct = opts.get('apply_ct', False)
             normalize_tanh = opts.get('normalize_tanh', False)
 
@@ -205,6 +206,13 @@ class SampleProcessor(object):
                     if np.random.randint(100) < chance:
                         img = imagelib.LinearMotionBlur (img, np.random.randint( mb_max_size )+1, np.random.randint(360) )
 
+                if gaussian_blur is not None:
+                    chance, kernel_max_size = gaussian_blur
+                    chance = np.clip(chance, 0, 100)
+                    
+                    if np.random.randint(100) < chance:
+                        img = cv2.GaussianBlur(img, ( np.random.randint( kernel_max_size )*2+1 ,) *2 , 0)
+
                 if is_face_sample and target_face_type != SPTF.NONE:
                     target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[target_face_type]
                     if target_ft > sample.face_type:
@@ -232,7 +240,7 @@ class SampleProcessor(object):
                     start_y = rnd_state.randint(sub_size+1)
                     img = img[start_y:start_y+sub_size,start_x:start_x+sub_size,:]
 
-                img = np.clip(img, 0, 1)
+                img = np.clip(img, 0, 1).astype(np.float32)
                 img_bgr  = img[...,0:3]
                 img_mask = img[...,3:4]
 
@@ -244,6 +252,18 @@ class SampleProcessor(object):
 
                     img_bgr = imagelib.linear_color_transfer (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)
+                    hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
+                    h, s, v = cv2.split(hsv)
+                    
+                    h = (h + rnd_state.randint(360) ) % 360
+                    s = np.clip ( s + rnd_state.random()-0.5, 0, 1 )
+                    v = np.clip ( v + rnd_state.random()-0.5, 0, 1 )
+                    hsv = cv2.merge([h, s, v])
+                    
+                    img_bgr = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) , 0, 1 )
 
                 if normalize_std_dev:
                     img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) )
@@ -252,21 +272,12 @@ class SampleProcessor(object):
                     img_bgr[:,:,0] -= 103.939
                     img_bgr[:,:,1] -= 116.779
                     img_bgr[:,:,2] -= 123.68
-
+                    
                 if mode_type == SPTF.MODE_BGR:
                     img = img_bgr
                 elif mode_type == SPTF.MODE_BGR_SHUFFLE:
                     rnd_state = np.random.RandomState (sample_rnd_seed)
                     img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
-                elif mode_type == SPTF.MODE_BGR_RANDOM_HUE_SHIFT:
-                    rnd_state = np.random.RandomState (sample_rnd_seed)
-                    hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
-                    h, s, v = cv2.split(hsv)
-                    
-                    h = (h + rnd_state.randint(360) ) % 360
-                    hsv = cv2.merge([h, s, v])
-                    
-                    img = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) , 0, 1 )
                 elif mode_type == SPTF.MODE_G:
                     img = np.concatenate ( (np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1),img_mask) , -1 )
                 elif mode_type == SPTF.MODE_GGG:

utils/Path_utils.py

@@ -3,12 +3,26 @@ from os import scandir
 
 image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
 
-def get_image_paths(dir_path, image_extensions=image_extensions):
+def scantree(path):
+    """Recursively yield DirEntry objects for given directory."""
+    for entry in scandir(path):
+        if entry.is_dir(follow_symlinks=False):
+            yield from scantree(entry.path)  # see below for Python 2.x
+        else:
+            yield entry
+            
+def get_image_paths(dir_path, image_extensions=image_extensions, subdirs=False):
     dir_path = Path (dir_path)
 
     result = []
     if dir_path.exists():
-        for x in list(scandir(str(dir_path))):
+        
+        if subdirs:
+            gen = scantree(str(dir_path))
+        else:
+            gen = scandir(str(dir_path))
+        
+        for x in list(gen):
             if any([x.name.lower().endswith(ext) for ext in image_extensions]):
                 result.append(x.path)
     return sorted(result)