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)
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 = TernausNet("FANSeg", self.fanseg_input_size , FaceType.toString( face_type ) )
self.fanseg_by_face_type[face_type] = fanseg
return fanseg.extract(*args, **kwargs)
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
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'
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
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
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), '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), 'resolution' : self.resolution, 'random_hsv_shift' : True},
])
])
def on_initialize(self, client_dict):
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.max_faces_from_image = client_dict['max_faces_from_image']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = Path(client_dict['final_output_dir']) if 'final_output_dir' in client_dict.keys() else None
self.debug_dir = client_dict['debug_dir']
#transfer and set stdin in order to work code.interact in debug subprocess
stdin_fd = client_dict['stdin_fd']
if stdin_fd is not None and DEBUG:
sys.stdin = os.fdopen(stdin_fd)
self.cached_image = (None, None)
self.e = None
device_config = nnlib.DeviceConfig ( cpu_only=self.cpu_only, force_gpu_idx=self.device_idx, allow_growth=True)
self.device_vram = device_config.gpu_vram_gb[0]
intro_str = 'Running on %s.' % (client_dict['device_name'])
if not self.cpu_only and self.device_vram <= 2:
intro_str += " Recommended to close all programs using this device."
self.log_info (intro_str)
if 'rects' in self.type:
if self.type == 'rects-mt':
nnlib.import_all (device_config)
self.e = facelib.MTCExtractor()
elif self.type == 'rects-dlib':
nnlib.import_dlib (device_config)
self.e = facelib.DLIBExtractor(nnlib.dlib)
elif self.type == 'rects-s3fd':
nnlib.import_all (device_config)
self.e = facelib.S3FDExtractor()
else:
raise ValueError ("Wrong type.")
if self.e is not None:
self.e.__enter__()
elif self.type == 'landmarks':
nnlib.import_all (device_config)
self.e = facelib.FANExtractor()
self.e.__enter__()
if self.device_vram >= 2:
self.second_pass_e = facelib.S3FDExtractor()
self.second_pass_e.__enter__()
else:
self.second_pass_e = None
elif self.type == 'fanseg':
nnlib.import_all (device_config)
self.e = TernausNet(256, FaceType.toString(FaceType.FULL) )
self.e.__enter__()
elif self.type == 'final':
pass
class Cli(Subprocessor.Cli):
#override
def on_initialize(self, client_dict):
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.max_faces_from_image = client_dict['max_faces_from_image']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = Path(
client_dict['final_output_dir']
) if 'final_output_dir' in client_dict.keys() else None
self.debug_dir = client_dict['debug_dir']
#transfer and set stdin in order to work code.interact in debug subprocess
stdin_fd = client_dict['stdin_fd']
if stdin_fd is not None and DEBUG:
sys.stdin = os.fdopen(stdin_fd)
self.cached_image = (None, None)
self.e = None
device_config = nnlib.DeviceConfig(cpu_only=self.cpu_only,
force_gpu_idx=self.device_idx,
allow_growth=True)
self.device_vram = device_config.gpu_vram_gb[0]
intro_str = 'Running on %s.' % (client_dict['device_name'])
if not self.cpu_only and self.device_vram <= 2:
intro_str += " Recommended to close all programs using this device."
self.log_info(intro_str)
if 'rects' in self.type:
if self.type == 'rects-mt':
nnlib.import_all(device_config)
self.e = facelib.MTCExtractor()
elif self.type == 'rects-dlib':
nnlib.import_dlib(device_config)
self.e = facelib.DLIBExtractor(nnlib.dlib)
elif self.type == 'rects-s3fd':
nnlib.import_all(device_config)
self.e = facelib.S3FDExtractor(do_dummy_predict=True)
else:
raise ValueError("Wrong type.")
if self.e is not None:
self.e.__enter__()
elif self.type == 'landmarks':
nnlib.import_all(device_config)
self.e = facelib.FANExtractor()
self.e.__enter__()
if self.device_vram >= 2:
self.second_pass_e = facelib.S3FDExtractor(
do_dummy_predict=False)
self.second_pass_e.__enter__()
else:
self.second_pass_e = None
elif self.type == 'fanseg':
nnlib.import_all(device_config)
self.e = TernausNet(256, FaceType.toString(FaceType.FULL))
self.e.__enter__()
elif self.type == 'final':
pass
#override
def on_finalize(self):
if self.e is not None:
self.e.__exit__()
#override
def process_data(self, data):
filename_path = Path(data.filename)
filename_path_str = str(filename_path)
if self.cached_image[0] == filename_path_str:
image = self.cached_image[
1] #cached image for manual extractor
else:
image = cv2_imread(filename_path_str)
if image is None:
self.log_err(
'Failed to extract %s, reason: cv2_imread() fail.' %
(str(filename_path)))
return data
image = imagelib.normalize_channels(image, 3)
h, w, ch = image.shape
wm, hm = w % 2, h % 2
if wm + hm != 0: #fix odd image
image = image[0:h - hm, 0:w - wm, :]
self.cached_image = (filename_path_str, image)
src_dflimg = None
h, w, ch = image.shape
if h == w:
#extracting from already extracted jpg image?
if filename_path.suffix == '.png':
src_dflimg = DFLPNG.load(str(filename_path))
if filename_path.suffix == '.jpg':
src_dflimg = DFLJPG.load(str(filename_path))
if 'rects' in self.type:
if min(w, h) < 128:
self.log_err('Image is too small %s : [%d, %d]' %
(str(filename_path), w, h))
data.rects = []
else:
for rot in ([0, 90, 270, 180]):
data.rects_rotation = rot
if rot == 0:
rotated_image = image
elif rot == 90:
rotated_image = image.swapaxes(0, 1)[:, ::-1, :]
elif rot == 180:
rotated_image = image[::-1, ::-1, :]
elif rot == 270:
rotated_image = image.swapaxes(0, 1)[::-1, :, :]
rects = data.rects = self.e.extract(rotated_image,
is_bgr=True)
if len(rects) != 0:
break
if self.max_faces_from_image != 0 and len(data.rects) > 1:
data.rects = data.rects[0:self.max_faces_from_image]
return data
elif self.type == 'landmarks':
if data.rects_rotation == 0:
rotated_image = image
elif data.rects_rotation == 90:
rotated_image = image.swapaxes(0, 1)[:, ::-1, :]
elif data.rects_rotation == 180:
rotated_image = image[::-1, ::-1, :]
elif data.rects_rotation == 270:
rotated_image = image.swapaxes(0, 1)[::-1, :, :]
data.landmarks = self.e.extract(
rotated_image,
data.rects,
self.second_pass_e if
(src_dflimg is None and data.landmarks_accurate) else None,
is_bgr=True)
if data.rects_rotation != 0:
for i, (rect,
lmrks) in enumerate(zip(data.rects,
data.landmarks)):
new_rect, new_lmrks = rect, lmrks
(l, t, r, b) = rect
if data.rects_rotation == 90:
new_rect = (t, h - l, b, h - r)
if lmrks is not None:
new_lmrks = lmrks[:, ::-1].copy()
new_lmrks[:, 1] = h - new_lmrks[:, 1]
elif data.rects_rotation == 180:
if lmrks is not None:
new_rect = (w - l, h - t, w - r, h - b)
new_lmrks = lmrks.copy()
new_lmrks[:, 0] = w - new_lmrks[:, 0]
new_lmrks[:, 1] = h - new_lmrks[:, 1]
elif data.rects_rotation == 270:
new_rect = (w - b, l, w - t, r)
if lmrks is not None:
new_lmrks = lmrks[:, ::-1].copy()
new_lmrks[:, 0] = w - new_lmrks[:, 0]
data.rects[i], data.landmarks[i] = new_rect, new_lmrks
return data
elif self.type == 'final':
data.final_output_files = []
rects = data.rects
landmarks = data.landmarks
if self.debug_dir is not None:
debug_output_file = str(
Path(self.debug_dir) / (filename_path.stem + '.jpg'))
debug_image = image.copy()
if src_dflimg is not None and len(rects) != 1:
#if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it
print("src_dflimg is not None and len(rects) != 1",
str(filename_path))
output_file = str(self.final_output_path /
filename_path.name)
if str(filename_path) != str(output_file):
shutil.copy(str(filename_path), str(output_file))
data.final_output_files.append(output_file)
else:
face_idx = 0
for rect, image_landmarks in zip(rects, landmarks):
if src_dflimg is not None and face_idx > 1:
#cannot extract more than 1 face from dflimg
break
if image_landmarks is None:
continue
rect = np.array(rect)
if self.face_type == FaceType.MARK_ONLY:
image_to_face_mat = None
face_image = image
face_image_landmarks = image_landmarks
else:
image_to_face_mat = LandmarksProcessor.get_transform_mat(
image_landmarks, self.image_size,
self.face_type)
face_image = cv2.warpAffine(
image, image_to_face_mat,
(self.image_size, self.image_size),
cv2.INTER_LANCZOS4)
face_image_landmarks = LandmarksProcessor.transform_points(
image_landmarks, image_to_face_mat)
landmarks_bbox = LandmarksProcessor.transform_points(
[(0, 0), (0, self.image_size - 1),
(self.image_size - 1, self.image_size - 1),
(self.image_size - 1, 0)], image_to_face_mat,
True)
rect_area = mathlib.polygon_area(
np.array(rect[[0, 2, 2, 0]]),
np.array(rect[[1, 1, 3, 3]]))
landmarks_area = mathlib.polygon_area(
landmarks_bbox[:, 0], landmarks_bbox[:, 1])
if landmarks_area > 4 * rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
continue
if self.debug_dir is not None:
LandmarksProcessor.draw_rect_landmarks(
debug_image,
rect,
image_landmarks,
self.image_size,
self.face_type,
transparent_mask=True)
if src_dflimg is not None and filename_path.suffix == '.jpg':
#if extracting from dflimg and jpg copy it in order not to lose quality
output_file = str(self.final_output_path /
filename_path.name)
if str(filename_path) != str(output_file):
shutil.copy(str(filename_path),
str(output_file))
else:
output_file = '{}_{}{}'.format(
str(self.final_output_path /
filename_path.stem), str(face_idx), '.jpg')
cv2_imwrite(output_file, face_image,
[int(cv2.IMWRITE_JPEG_QUALITY), 85])
DFLJPG.embed_data(
output_file,
face_type=FaceType.toString(self.face_type),
landmarks=face_image_landmarks.tolist(),
source_filename=filename_path.name,
source_rect=rect,
source_landmarks=image_landmarks.tolist(),
image_to_face_mat=image_to_face_mat,
pitch_yaw_roll=data.pitch_yaw_roll)
data.final_output_files.append(output_file)
face_idx += 1
data.faces_detected = face_idx
if self.debug_dir is not None:
cv2_imwrite(debug_output_file, debug_image,
[int(cv2.IMWRITE_JPEG_QUALITY), 50])
return data
elif self.type == 'fanseg':
if src_dflimg is not None:
fanseg_mask = self.e.extract(image / 255.0)
src_dflimg.embed_and_set(
filename_path_str,
fanseg_mask=fanseg_mask,
)
#overridable
def get_data_name(self, data):
#return string identificator of your data
return data.filename
class Model(ModelBase):
def __init__(self, *args, **kwargs):
super().__init__(*args,
**kwargs,
ask_enable_autobackup=False,
ask_write_preview_history=False,
ask_target_iter=False,
ask_sort_by_yaw=False,
ask_random_flip=False,
ask_src_scale_mod=False)
#override
def onInitializeOptions(self, is_first_run, ask_override):
default_face_type = 'f'
if is_first_run:
self.options['face_type'] = io.input_str(
"Half or Full face? (h/f, ?:help skip:f) : ",
default_face_type, ['h', 'f'],
help_message="").lower()
else:
self.options['face_type'] = self.options.get(
'face_type', default_face_type)
#override
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
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 = 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
face_type = t.FACE_TYPE_FULL if self.options[
'face_type'] == 'f' else t.FACE_TYPE_HALF
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),
'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),
'resolution': self.resolution,
'random_hsv_shift': True
},
])
])
#override
def onSave(self):
self.fan_seg.save_weights()
#override
def onTrainOneIter(self, generators_samples, generators_list):
target_src, target_src_mask = generators_samples[0]
loss = self.fan_seg.train(target_src, target_src_mask)
return (('loss', loss), )
#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)
st = []
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))
st2 = []
for i in range(0, len(test_B)):
st2.append(
np.concatenate((
test_B[i, :, :, 0:3],
mBB[i],
test_B[i, :, :, 0:3] * mBB[i],
),
axis=1))
return [
('training data', np.concatenate(st, axis=0)),
('evaluating data', np.concatenate(st2, axis=0)),
]