mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 13:33:24 -07:00
refactoring
This commit is contained in:
iperov
parent
44798c2b85
commit
8a223845fb
19 changed files with 963 additions and 468 deletions
|
|
@ -1,50 +0,0 @@
|
|||
from enum import IntEnum
|
||||
import cv2
|
||||
import numpy as np
|
||||
from random import randint
|
||||
from facelib import FaceType
|
||||
|
||||
|
||||
class TrainingDataType(IntEnum):
|
||||
IMAGE = 0 #raw image
|
||||
|
||||
FACE_BEGIN = 1
|
||||
FACE = 1 #aligned face unsorted
|
||||
FACE_YAW_SORTED = 2 #sorted by yaw
|
||||
FACE_YAW_SORTED_AS_TARGET = 3 #sorted by yaw and included only yaws which exist in TARGET also automatic mirrored
|
||||
FACE_END = 3
|
||||
|
||||
QTY = 4
|
||||
|
||||
|
||||
class TrainingDataSample(object):
|
||||
|
||||
def __init__(self, filename=None, face_type=None, shape=None, landmarks=None, yaw=None, mirror=None, nearest_target_list=None):
|
||||
self.filename = filename
|
||||
self.face_type = face_type
|
||||
self.shape = shape
|
||||
self.landmarks = np.array(landmarks) if landmarks is not None else None
|
||||
self.yaw = yaw
|
||||
self.mirror = mirror
|
||||
self.nearest_target_list = nearest_target_list
|
||||
|
||||
def copy_and_set(self, filename=None, face_type=None, shape=None, landmarks=None, yaw=None, mirror=None, nearest_target_list=None):
|
||||
return TrainingDataSample(
|
||||
filename=filename if filename is not None else self.filename,
|
||||
face_type=face_type if face_type is not None else self.face_type,
|
||||
shape=shape if shape is not None else self.shape,
|
||||
landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
|
||||
yaw=yaw if yaw is not None else self.yaw,
|
||||
mirror=mirror if mirror is not None else self.mirror,
|
||||
nearest_target_list=nearest_target_list if nearest_target_list is not None else self.nearest_target_list)
|
||||
|
||||
def load_bgr(self):
|
||||
img = cv2.imread (self.filename).astype(np.float32) / 255.0
|
||||
if self.mirror:
|
||||
img = img[:,::-1].copy()
|
||||
return img
|
||||
|
||||
def get_random_nearest_target_sample(self):
|
||||
if self.nearest_target_list is None:
|
||||
return None
|
||||
return self.nearest_target_list[randint (0, len(self.nearest_target_list)-1)]
|
||||
|
|
@ -10,7 +10,7 @@ from utils import image_utils
|
|||
import numpy as np
|
||||
import cv2
|
||||
import gpufmkmgr
|
||||
from .TrainingDataGeneratorBase import TrainingDataGeneratorBase
|
||||
from samples import SampleGeneratorBase
|
||||
|
||||
'''
|
||||
You can implement your own model. Check examples.
|
||||
|
|
@ -47,13 +47,11 @@ class ModelBase(object):
|
|||
self.epoch = model_data['epoch']
|
||||
self.options = model_data['options']
|
||||
self.loss_history = model_data['loss_history'] if 'loss_history' in model_data.keys() else []
|
||||
self.generator_dict_states = model_data['generator_dict_states'] if 'generator_dict_states' in model_data.keys() else None
|
||||
self.sample_for_preview = model_data['sample_for_preview'] if 'sample_for_preview' in model_data.keys() else None
|
||||
else:
|
||||
self.epoch = 0
|
||||
self.options = {}
|
||||
self.loss_history = []
|
||||
self.generator_dict_states = None
|
||||
self.sample_for_preview = None
|
||||
|
||||
if self.write_preview_history:
|
||||
|
|
@ -97,11 +95,8 @@ class ModelBase(object):
|
|||
raise Exception( 'You didnt set_training_data_generators()')
|
||||
else:
|
||||
for i, generator in enumerate(self.generator_list):
|
||||
if not isinstance(generator, TrainingDataGeneratorBase):
|
||||
raise Exception('training data generator is not subclass of TrainingDataGeneratorBase')
|
||||
|
||||
if self.generator_dict_states is not None and i < len(self.generator_dict_states):
|
||||
generator.set_dict_state ( self.generator_dict_states[i] )
|
||||
if not isinstance(generator, SampleGeneratorBase):
|
||||
raise Exception('training data generator is not subclass of SampleGeneratorBase')
|
||||
|
||||
if self.sample_for_preview is None:
|
||||
self.sample_for_preview = self.generate_next_sample()
|
||||
|
|
@ -212,7 +207,6 @@ class ModelBase(object):
|
|||
'epoch': self.epoch,
|
||||
'options': self.options,
|
||||
'loss_history': self.loss_history,
|
||||
'generator_dict_states' : [generator.get_dict_state() for generator in self.generator_list],
|
||||
'sample_for_preview' : self.sample_for_preview
|
||||
}
|
||||
self.model_data_path.write_bytes( pickle.dumps(model_data) )
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
from models import ModelBase
|
||||
from samples import *
|
||||
from nnlib import tf_dssim
|
||||
from nnlib import DSSIMLossClass
|
||||
from nnlib import conv
|
||||
|
|
@ -72,21 +72,20 @@ class Model(ModelBase):
|
|||
self.BA256_view = K.function ([input_B_warped64], [BA_rec256])
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 256],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_BGR, 256] ] ),
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 256],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 256] ] ),
|
||||
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_BGR, 256] ] )
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 256] ] )
|
||||
])
|
||||
#override
|
||||
def onSave(self):
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
|
@ -7,7 +6,7 @@ from nnlib import DSSIMMaskLossClass
|
|||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
from samples import *
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
|
|
@ -42,11 +41,17 @@ class Model(ModelBase):
|
|||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
|
||||
])
|
||||
#override
|
||||
def onSave(self):
|
||||
|
|
|
|||
|
|
@ -1,13 +1,12 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
|
||||
import cv2
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -48,11 +47,17 @@ class Model(ModelBase):
|
|||
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
|
||||
])
|
||||
|
||||
#override
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
from samples import *
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
|
|
@ -46,11 +46,17 @@ class Model(ModelBase):
|
|||
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 64] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 64] ], random_flip=True )
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] )
|
||||
])
|
||||
|
||||
#override
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
|
@ -7,6 +6,7 @@ from nnlib import DSSIMMaskLossClass
|
|||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -48,11 +48,17 @@ class Model(ModelBase):
|
|||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
|
||||
])
|
||||
|
||||
#override
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
|
@ -7,7 +6,8 @@ from nnlib import DSSIMMaskLossClass
|
|||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
|
|
@ -47,12 +47,18 @@ class Model(ModelBase):
|
|||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE_YAW_SORTED_AS_TARGET, self.training_data_src_path, target_training_data_path=self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
|
||||
])
|
||||
|
||||
#override
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from models import ModelBase
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -82,11 +82,18 @@ class Model(ModelBase):
|
|||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_GGG, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_G , 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_GGG, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128]], random_flip=True )
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_GGG, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_G , 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_GGG, 128] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128]] )
|
||||
])
|
||||
#override
|
||||
def onSave(self):
|
||||
|
|
|
|||
|
|
@ -1,149 +0,0 @@
|
|||
from facelib import FaceType
|
||||
from facelib import LandmarksProcessor
|
||||
import cv2
|
||||
import numpy as np
|
||||
from models import TrainingDataGeneratorBase
|
||||
from utils import image_utils
|
||||
from utils import random_utils
|
||||
from enum import IntEnum
|
||||
from models import TrainingDataType
|
||||
|
||||
class TrainingDataGenerator(TrainingDataGeneratorBase):
|
||||
class SampleTypeFlags(IntEnum):
|
||||
SOURCE = 0x000001,
|
||||
WARPED = 0x000002,
|
||||
WARPED_TRANSFORMED = 0x000004,
|
||||
TRANSFORMED = 0x000008,
|
||||
|
||||
HALF_FACE = 0x000010,
|
||||
FULL_FACE = 0x000020,
|
||||
HEAD_FACE = 0x000040,
|
||||
AVATAR_FACE = 0x000080,
|
||||
MARK_ONLY_FACE = 0x000100,
|
||||
|
||||
MODE_BGR = 0x001000, #BGR
|
||||
MODE_G = 0x002000, #Grayscale
|
||||
MODE_GGG = 0x004000, #3xGrayscale
|
||||
MODE_M = 0x008000, #mask only
|
||||
MODE_BGR_SHUFFLE = 0x010000, #BGR shuffle
|
||||
|
||||
MASK_FULL = 0x100000,
|
||||
MASK_EYES = 0x200000,
|
||||
|
||||
#overrided
|
||||
def onInitialize(self, random_flip=False, normalize_tanh=False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], output_sample_types=[], **kwargs):
|
||||
self.random_flip = random_flip
|
||||
self.normalize_tanh = normalize_tanh
|
||||
self.output_sample_types = output_sample_types
|
||||
self.rotation_range = rotation_range
|
||||
self.scale_range = scale_range
|
||||
self.tx_range = tx_range
|
||||
self.ty_range = ty_range
|
||||
|
||||
#overrided
|
||||
def onProcessSample(self, sample, debug):
|
||||
source = sample.load_bgr()
|
||||
h,w,c = source.shape
|
||||
|
||||
is_face_sample = self.trainingdatatype >= TrainingDataType.FACE_BEGIN and self.trainingdatatype <= TrainingDataType.FACE_END
|
||||
|
||||
if debug and is_face_sample:
|
||||
LandmarksProcessor.draw_landmarks (source, sample.landmarks, (0, 1, 0))
|
||||
|
||||
params = image_utils.gen_warp_params(source, self.random_flip, rotation_range=self.rotation_range, scale_range=self.scale_range, tx_range=self.tx_range, ty_range=self.ty_range )
|
||||
|
||||
images = [[None]*3 for _ in range(4)]
|
||||
|
||||
outputs = []
|
||||
for t,size in self.output_sample_types:
|
||||
if t & self.SampleTypeFlags.SOURCE != 0:
|
||||
img_type = 0
|
||||
elif t & self.SampleTypeFlags.WARPED != 0:
|
||||
img_type = 1
|
||||
elif t & self.SampleTypeFlags.WARPED_TRANSFORMED != 0:
|
||||
img_type = 2
|
||||
elif t & self.SampleTypeFlags.TRANSFORMED != 0:
|
||||
img_type = 3
|
||||
else:
|
||||
raise ValueError ('expected SampleTypeFlags type')
|
||||
|
||||
mask_type = 0
|
||||
if t & self.SampleTypeFlags.MASK_FULL != 0:
|
||||
mask_type = 1
|
||||
elif t & self.SampleTypeFlags.MASK_EYES != 0:
|
||||
mask_type = 2
|
||||
|
||||
if images[img_type][mask_type] is None:
|
||||
img = source
|
||||
if is_face_sample:
|
||||
if mask_type == 1:
|
||||
img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (source, sample.landmarks) ), -1 )
|
||||
elif mask_type == 2:
|
||||
mask = LandmarksProcessor.get_image_eye_mask (source, sample.landmarks)
|
||||
mask = np.expand_dims (cv2.blur (mask, ( w // 32, w // 32 ) ), -1)
|
||||
mask[mask > 0.0] = 1.0
|
||||
img = np.concatenate( (img, mask ), -1 )
|
||||
|
||||
images[img_type][mask_type] = image_utils.warp_by_params (params, img, (img_type==1 or img_type==2), (img_type==2 or img_type==3), img_type != 0)
|
||||
|
||||
img = images[img_type][mask_type]
|
||||
|
||||
target_face_type = -1
|
||||
if t & self.SampleTypeFlags.HALF_FACE != 0:
|
||||
target_face_type = FaceType.HALF
|
||||
elif t & self.SampleTypeFlags.FULL_FACE != 0:
|
||||
target_face_type = FaceType.FULL
|
||||
elif t & self.SampleTypeFlags.HEAD_FACE != 0:
|
||||
target_face_type = FaceType.HEAD
|
||||
elif t & self.SampleTypeFlags.AVATAR_FACE != 0:
|
||||
target_face_type = FaceType.AVATAR
|
||||
elif t & self.SampleTypeFlags.MARK_ONLY_FACE != 0:
|
||||
target_face_type = FaceType.MARK_ONLY
|
||||
|
||||
if is_face_sample and target_face_type != -1 and target_face_type != FaceType.MARK_ONLY:
|
||||
if target_face_type > sample.face_type:
|
||||
raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_face_type) )
|
||||
|
||||
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, size, target_face_type), (size,size), flags=cv2.INTER_LANCZOS4 )
|
||||
else:
|
||||
img = cv2.resize( img, (size,size), cv2.INTER_LANCZOS4 )
|
||||
|
||||
img_bgr = img[...,0:3]
|
||||
img_mask = img[...,3:4]
|
||||
|
||||
if t & self.SampleTypeFlags.MODE_BGR != 0:
|
||||
img = img
|
||||
elif t & self.SampleTypeFlags.MODE_BGR_SHUFFLE != 0:
|
||||
img_bgr = np.take (img_bgr, np.random.permutation(img_bgr.shape[-1]), axis=-1)
|
||||
img = np.concatenate ( (img_bgr,img_mask) , -1 )
|
||||
elif t & self.SampleTypeFlags.MODE_G != 0:
|
||||
img = np.concatenate ( (np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1),img_mask) , -1 )
|
||||
elif t & self.SampleTypeFlags.MODE_GGG != 0:
|
||||
img = np.concatenate ( ( np.repeat ( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1), (3,), -1), img_mask), -1)
|
||||
elif is_face_sample and t & self.SampleTypeFlags.MODE_M != 0:
|
||||
if mask_type== 0:
|
||||
raise ValueError ('no mask mode defined')
|
||||
img = img_mask
|
||||
else:
|
||||
raise ValueError ('expected SampleTypeFlags mode')
|
||||
|
||||
if not debug and self.normalize_tanh:
|
||||
img = img * 2.0 - 1.0
|
||||
|
||||
outputs.append ( img )
|
||||
|
||||
if debug:
|
||||
result = ()
|
||||
|
||||
for output in outputs:
|
||||
if output.shape[2] < 4:
|
||||
result += (output,)
|
||||
elif output.shape[2] == 4:
|
||||
result += (output[...,0:3]*output[...,3:4],)
|
||||
|
||||
return result
|
||||
else:
|
||||
return outputs
|
||||
|
||||
|
||||
|
||||
|
|
@ -1,240 +0,0 @@
|
|||
import traceback
|
||||
import random
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
import cv2
|
||||
from utils.DFLPNG import DFLPNG
|
||||
from utils import iter_utils
|
||||
from utils import Path_utils
|
||||
from .BaseTypes import TrainingDataType
|
||||
from .BaseTypes import TrainingDataSample
|
||||
from facelib import FaceType
|
||||
from facelib import LandmarksProcessor
|
||||
|
||||
'''
|
||||
You can implement your own TrainingDataGenerator
|
||||
'''
|
||||
class TrainingDataGeneratorBase(object):
|
||||
cache = dict()
|
||||
|
||||
#DONT OVERRIDE
|
||||
#use YourOwnTrainingDataGenerator (..., your_opt=1)
|
||||
#and then this opt will be passed in YourOwnTrainingDataGenerator.onInitialize ( your_opt )
|
||||
def __init__ (self, trainingdatatype, training_data_path, target_training_data_path=None, debug=False, batch_size=1, **kwargs):
|
||||
if not isinstance(trainingdatatype, TrainingDataType):
|
||||
raise Exception('TrainingDataGeneratorBase() trainingdatatype is not TrainingDataType')
|
||||
|
||||
if training_data_path is None:
|
||||
raise Exception('training_data_path is None')
|
||||
|
||||
self.training_data_path = Path(training_data_path)
|
||||
self.target_training_data_path = Path(target_training_data_path) if target_training_data_path is not None else None
|
||||
|
||||
self.debug = debug
|
||||
self.batch_size = 1 if self.debug else batch_size
|
||||
self.trainingdatatype = trainingdatatype
|
||||
self.data = TrainingDataGeneratorBase.load (trainingdatatype, self.training_data_path, self.target_training_data_path)
|
||||
|
||||
if self.debug:
|
||||
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, self.data)]
|
||||
else:
|
||||
if len(self.data) > 1:
|
||||
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, self.data[0::2] ),
|
||||
iter_utils.SubprocessGenerator ( self.batch_func, self.data[1::2] )]
|
||||
else:
|
||||
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, self.data )]
|
||||
|
||||
self.generator_counter = -1
|
||||
self.onInitialize(**kwargs)
|
||||
|
||||
#overridable
|
||||
def onInitialize(self, **kwargs):
|
||||
#your TrainingDataGenerator initialization here
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onProcessSample(self, sample, debug):
|
||||
#process sample and return tuple of images for your model in onTrainOneEpoch
|
||||
return ( np.zeros( (64,64,4), dtype=np.float32 ), )
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
self.generator_counter += 1
|
||||
generator = self.generators[self.generator_counter % len(self.generators) ]
|
||||
x = next(generator)
|
||||
return x
|
||||
|
||||
def batch_func(self, data):
|
||||
data_len = len(data)
|
||||
if data_len == 0:
|
||||
raise ValueError('No training data provided.')
|
||||
|
||||
if self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED or self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if all ( [ x == None for x in data] ):
|
||||
raise ValueError('Not enough training data. Gather more faces!')
|
||||
|
||||
if self.trainingdatatype == TrainingDataType.IMAGE or self.trainingdatatype == TrainingDataType.FACE:
|
||||
shuffle_idxs = []
|
||||
elif self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED or self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
shuffle_idxs = []
|
||||
shuffle_idxs_2D = [[]]*data_len
|
||||
|
||||
while True:
|
||||
|
||||
batches = None
|
||||
for n_batch in range(0, self.batch_size):
|
||||
while True:
|
||||
sample = None
|
||||
|
||||
if self.trainingdatatype == TrainingDataType.IMAGE or self.trainingdatatype == TrainingDataType.FACE:
|
||||
if len(shuffle_idxs) == 0:
|
||||
shuffle_idxs = [ i for i in range(0, data_len) ]
|
||||
random.shuffle(shuffle_idxs)
|
||||
idx = shuffle_idxs.pop()
|
||||
sample = data[ idx ]
|
||||
elif self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED or self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if len(shuffle_idxs) == 0:
|
||||
shuffle_idxs = [ i for i in range(0, data_len) ]
|
||||
random.shuffle(shuffle_idxs)
|
||||
|
||||
idx = shuffle_idxs.pop()
|
||||
if data[idx] != None:
|
||||
if len(shuffle_idxs_2D[idx]) == 0:
|
||||
shuffle_idxs_2D[idx] = [ i for i in range(0, len(data[idx])) ]
|
||||
random.shuffle(shuffle_idxs_2D[idx])
|
||||
|
||||
idx2 = shuffle_idxs_2D[idx].pop()
|
||||
sample = data[idx][idx2]
|
||||
|
||||
if sample is not None:
|
||||
try:
|
||||
x = self.onProcessSample (sample, self.debug)
|
||||
except:
|
||||
raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )
|
||||
|
||||
if type(x) != tuple and type(x) != list:
|
||||
raise Exception('TrainingDataGenerator.onProcessSample() returns NOT tuple/list')
|
||||
|
||||
x_len = len(x)
|
||||
if batches is None:
|
||||
batches = [ [] for _ in range(0,x_len) ]
|
||||
|
||||
for i in range(0,x_len):
|
||||
batches[i].append ( x[i] )
|
||||
|
||||
break
|
||||
|
||||
yield [ np.array(batch) for batch in batches]
|
||||
|
||||
def get_dict_state(self):
|
||||
return {}
|
||||
|
||||
def set_dict_state(self, state):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def load(trainingdatatype, training_data_path, target_training_data_path=None):
|
||||
cache = TrainingDataGeneratorBase.cache
|
||||
|
||||
if str(training_data_path) not in cache.keys():
|
||||
cache[str(training_data_path)] = [None]*TrainingDataType.QTY
|
||||
|
||||
if target_training_data_path is not None and str(target_training_data_path) not in cache.keys():
|
||||
cache[str(target_training_data_path)] = [None]*TrainingDataType.QTY
|
||||
|
||||
datas = cache[str(training_data_path)]
|
||||
|
||||
if trainingdatatype == TrainingDataType.IMAGE:
|
||||
if datas[trainingdatatype] is None:
|
||||
datas[trainingdatatype] = [ TrainingDataSample(filename=filename) for filename in tqdm( Path_utils.get_image_paths(training_data_path), desc="Loading" ) ]
|
||||
|
||||
elif trainingdatatype == TrainingDataType.FACE:
|
||||
if datas[trainingdatatype] is None:
|
||||
datas[trainingdatatype] = X_LOAD( [ TrainingDataSample(filename=filename) for filename in Path_utils.get_image_paths(training_data_path) ] )
|
||||
|
||||
elif trainingdatatype == TrainingDataType.FACE_YAW_SORTED:
|
||||
if datas[trainingdatatype] is None:
|
||||
datas[trainingdatatype] = X_YAW_SORTED( TrainingDataGeneratorBase.load(TrainingDataType.FACE, training_data_path) )
|
||||
|
||||
elif trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if datas[trainingdatatype] is None:
|
||||
if target_training_data_path is None:
|
||||
raise Exception('target_training_data_path is None for FACE_YAW_SORTED_AS_TARGET')
|
||||
datas[trainingdatatype] = X_YAW_AS_Y_SORTED( TrainingDataGeneratorBase.load(TrainingDataType.FACE_YAW_SORTED, training_data_path), TrainingDataGeneratorBase.load(TrainingDataType.FACE_YAW_SORTED, target_training_data_path) )
|
||||
|
||||
return datas[trainingdatatype]
|
||||
|
||||
def X_LOAD ( RAWS ):
|
||||
sample_list = []
|
||||
|
||||
for s in tqdm( RAWS, desc="Loading" ):
|
||||
|
||||
s_filename_path = Path(s.filename)
|
||||
if s_filename_path.suffix != '.png':
|
||||
print ("%s is not a png file required for training" % (s_filename_path.name) )
|
||||
continue
|
||||
|
||||
dflpng = DFLPNG.load ( str(s_filename_path), print_on_no_embedded_data=True )
|
||||
if dflpng is None:
|
||||
continue
|
||||
|
||||
sample_list.append( s.copy_and_set(face_type=FaceType.fromString (dflpng.get_face_type()),
|
||||
shape=dflpng.get_shape(),
|
||||
landmarks=dflpng.get_landmarks(),
|
||||
yaw=dflpng.get_yaw_value()) )
|
||||
|
||||
return sample_list
|
||||
|
||||
def X_YAW_SORTED( YAW_RAWS ):
|
||||
|
||||
lowest_yaw, highest_yaw = -32, +32
|
||||
gradations = 64
|
||||
diff_rot_per_grad = abs(highest_yaw-lowest_yaw) / gradations
|
||||
|
||||
yaws_sample_list = [None]*gradations
|
||||
|
||||
for i in tqdm( range(0, gradations), desc="Sorting" ):
|
||||
yaw = lowest_yaw + i*diff_rot_per_grad
|
||||
next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad
|
||||
|
||||
yaw_samples = []
|
||||
for s in YAW_RAWS:
|
||||
s_yaw = s.yaw
|
||||
if (i == 0 and s_yaw < next_yaw) or \
|
||||
(i < gradations-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
|
||||
(i == gradations-1 and s_yaw >= yaw):
|
||||
yaw_samples.append ( s )
|
||||
|
||||
if len(yaw_samples) > 0:
|
||||
yaws_sample_list[i] = yaw_samples
|
||||
|
||||
return yaws_sample_list
|
||||
|
||||
def X_YAW_AS_Y_SORTED (s, t):
|
||||
l = len(s)
|
||||
if l != len(t):
|
||||
raise Exception('X_YAW_AS_Y_SORTED() s_len != t_len')
|
||||
b = l // 2
|
||||
|
||||
s_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in s] ) == 1 )[:,0]
|
||||
t_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in t] ) == 1 )[:,0]
|
||||
|
||||
new_s = [None]*l
|
||||
|
||||
for t_idx in t_idxs:
|
||||
search_idxs = []
|
||||
for i in range(0,l):
|
||||
search_idxs += [t_idx - i, (l-t_idx-1) - i, t_idx + i, (l-t_idx-1) + i]
|
||||
|
||||
for search_idx in search_idxs:
|
||||
if search_idx in s_idxs:
|
||||
mirrored = ( t_idx != search_idx and ((t_idx < b and search_idx >= b) or (search_idx < b and t_idx >= b)) )
|
||||
new_s[t_idx] = [ sample.copy_and_set(mirror=True, yaw=-sample.yaw, landmarks=LandmarksProcessor.mirror_landmarks (sample.landmarks, sample.shape[1] ))
|
||||
for sample in s[search_idx]
|
||||
] if mirrored else s[search_idx]
|
||||
break
|
||||
|
||||
return new_s
|
||||
|
|
@ -1,12 +1,7 @@
|
|||
from .BaseTypes import TrainingDataType
|
||||
from .BaseTypes import TrainingDataSample
|
||||
|
||||
from .ModelBase import ModelBase
|
||||
from .ConverterBase import ConverterBase
|
||||
from .ConverterMasked import ConverterMasked
|
||||
from .ConverterImage import ConverterImage
|
||||
from .TrainingDataGeneratorBase import TrainingDataGeneratorBase
|
||||
from .TrainingDataGenerator import TrainingDataGenerator
|
||||
|
||||
def import_model(name):
|
||||
module = __import__('Model_'+name, globals(), locals(), [], 1)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue