diff --git a/converters/ConvertMasked.py b/converters/ConvertMasked.py index 6a5b274..e06a46b 100644 --- a/converters/ConvertMasked.py +++ b/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] diff --git a/converters/ConverterConfig.py b/converters/ConverterConfig.py index 9f1e174..0bbe556 100644 --- a/converters/ConverterConfig.py +++ b/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'} diff --git a/facelib/__init__.py b/facelib/__init__.py index 47a0de9..8ca071b 100644 --- a/facelib/__init__.py +++ b/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 \ No newline at end of file diff --git a/main.py b/main.py index 8401aee..a7bf13c 100644 --- a/main.py +++ b/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() diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 1968567..77405b8 100644 --- a/mainscripts/Converter.py +++ b/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) diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 0935596..e452c8d 100644 --- a/mainscripts/Extractor.py +++ b/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 ('-------------------------') - diff --git a/mainscripts/dev_misc.py b/mainscripts/dev_misc.py index fc3f18b..72be53f 100644 --- a/mainscripts/dev_misc.py +++ b/mainscripts/dev_misc.py @@ -1,9 +1,20 @@ -from . import Extractor -from . import Sorter -from pathlib import Path -from utils import Path_utils +import multiprocessing 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)} ") \ No newline at end of file + 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())) diff --git a/models/Model_DEV_FANSEG/Model.py b/models/Model_DEV_FANSEG/Model.py index ea24c9a..7fca4b5 100644 --- a/models/Model_DEV_FANSEG/Model.py +++ b/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 - - self.fan_seg = FANSegmentator(self.resolution, - FaceType.toString(self.face_type), - load_weights=not self.is_first_run(), - weights_file_root=self.get_model_root_path(), - training=True) + model_name = 'FANSeg' + model_name = 'FANCHQ' + self.fan_seg = TernausNet(model_name, self.resolution, + FaceType.toString(self.face_type), + load_weights=not self.is_first_run(), + 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) ) diff --git a/facelib/FANSeg_256_full_face.h5 b/nnlib/FANSeg_256_full_face.h5 similarity index 100% rename from facelib/FANSeg_256_full_face.h5 rename to nnlib/FANSeg_256_full_face.h5 diff --git a/facelib/FANSegmentator.py b/nnlib/TernausNet.py similarity index 90% rename from facelib/FANSegmentator.py rename to nnlib/TernausNet.py index 21b4f07..9016ead 100644 --- a/facelib/FANSegmentator.py +++ b/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) diff --git a/nnlib/__init__.py b/nnlib/__init__.py index 60fc709..d0d6a59 100644 --- a/nnlib/__init__.py +++ b/nnlib/__init__.py @@ -1,2 +1,3 @@ from .nnlib import nnlib -from .FUNIT import FUNIT \ No newline at end of file +from .FUNIT import FUNIT +from .TernausNet import TernausNet \ No newline at end of file diff --git a/facelib/vgg11_enc_weights.npy b/nnlib/vgg11_enc_weights.npy similarity index 100% rename from facelib/vgg11_enc_weights.npy rename to nnlib/vgg11_enc_weights.npy diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index a8cc851..7166b1f 100644 --- a/samplelib/SampleProcessor.py +++ b/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: diff --git a/utils/Path_utils.py b/utils/Path_utils.py index e753834..c609572 100644 --- a/utils/Path_utils.py +++ b/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)