def __init__(self, predictor_func,
predictor_input_size=0,
predictor_masked=True,
face_type=FaceType.FULL,
default_mode = 4,
base_erode_mask_modifier = 0,
base_blur_mask_modifier = 0,
default_erode_mask_modifier = 0,
default_blur_mask_modifier = 0,
clip_hborder_mask_per = 0,
force_mask_mode=-1):
super().__init__(predictor_func, Converter.TYPE_FACE)
#dummy predict and sleep, tensorflow caching kernels. If remove it, conversion speed will be x2 slower
predictor_func ( np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ) )
time.sleep(2)
predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func)
self.predictor_func_host = AntiPickler(predictor_func_host)
self.predictor_func = predictor_func
self.predictor_masked = predictor_masked
self.predictor_input_size = predictor_input_size
self.face_type = face_type
self.clip_hborder_mask_per = clip_hborder_mask_per
mode = io.input_int ("Choose mode: (1) overlay, (2) hist match, (3) hist match bw, (4) seamless, (5) raw. Default - %d : " % (default_mode) , default_mode)
mode_dict = {1:'overlay',
2:'hist-match',
3:'hist-match-bw',
4:'seamless',
5:'raw'}
self.mode = mode_dict.get (mode, mode_dict[default_mode] )
if self.mode == 'raw':
mode = io.input_int ("Choose raw mode: (1) rgb, (2) rgb+mask (default), (3) mask only, (4) predicted only : ", 2)
self.raw_mode = {1:'rgb',
2:'rgb-mask',
3:'mask-only',
4:'predicted-only'}.get (mode, 'rgb-mask')
if self.mode != 'raw':
if self.mode == 'seamless':
if io.input_bool("Seamless hist match? (y/n skip:n) : ", False):
self.mode = 'seamless-hist-match'
if self.mode == 'hist-match' or self.mode == 'hist-match-bw':
self.masked_hist_match = io.input_bool("Masked hist match? (y/n skip:y) : ", True)
if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match':
self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255)
if force_mask_mode != -1:
self.mask_mode = force_mask_mode
else:
if face_type == FaceType.FULL:
self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks."), 1, 6 )
else:
self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst . Default - %d : " % (1) , 1), 1, 2 )
if self.mode != 'raw':
self.erode_mask_modifier = base_erode_mask_modifier + np.clip ( io.input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (default_erode_mask_modifier), default_erode_mask_modifier), -200, 200)
self.blur_mask_modifier = base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (default_blur_mask_modifier), default_blur_mask_modifier), -200, 200)
self.output_face_scale = np.clip ( 1.0 + io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5)
if self.mode != 'raw':
self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct'])
self.super_resolution = io.input_bool("Apply super resolution? (y/n ?:help skip:n) : ", False, help_message="Enhance details by applying DCSCN network.")
if self.mode != 'raw':
self.final_image_color_degrade_power = np.clip ( io.input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100)
self.alpha = io.input_bool("Export png with alpha channel? (y/n skip:n) : ", False)
io.log_info ("")
if self.super_resolution:
host_proc, dc_upscale = SubprocessFunctionCaller.make_pair( imagelib.DCSCN().upscale )
self.dc_host = AntiPickler(host_proc)
self.dc_upscale = dc_upscale
else:
self.dc_host = None
if (self.mask_mode >= 3 and self.mask_mode <= 6):
host_proc, fan_seg_extract = SubprocessFunctionCaller.make_pair( FANSegmentator(256, FaceType.toString( self.face_type ) ).extract )
self.fan_seg_host = AntiPickler(host_proc)
self.fan_seg_extract = fan_seg_extract
else:
self.fan_seg_host = None
def process_data(self, data):
filename_path = Path(data[0])
filename_path_str = str(filename_path)
if self.cached_image[0] == filename_path_str:
image = self.cached_image[1]
else:
image = cv2_imread(filename_path_str)
self.cached_image = (filename_path_str, image)
if image is None:
self.log_err(
'Failed to extract %s, reason: cv2_imread() fail.' %
(str(filename_path)))
else:
if self.type == 'rects':
rects = self.e.extract_from_bgr(image)
return [str(filename_path), rects]
elif self.type == 'landmarks':
rects = data[1]
landmarks = self.e.extract_from_bgr(image, rects)
return [str(filename_path), landmarks]
elif self.type == 'final':
src_dflimg = None
(h, w, c) = image.shape
if h == w:
#extracting from already extracted jpg image?
if filename_path.suffix == '.jpg':
src_dflimg = DFLJPG.load(str(filename_path))
result = []
faces = data[1]
if self.debug:
debug_output_file = '{}{}'.format(
str(
Path(str(self.output_path) + '_debug') /
filename_path.stem), '.jpg')
debug_image = image.copy()
for (face_idx, face) in enumerate(faces):
output_file = '{}_{}{}'.format(
str(self.output_path / filename_path.stem),
str(face_idx), '.jpg')
rect = face[0]
image_landmarks = np.array(face[1])
if self.debug:
LandmarksProcessor.draw_rect_landmarks(
debug_image, rect, image_landmarks,
self.image_size, self.face_type)
if self.face_type == FaceType.MARK_ONLY:
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)
if src_dflimg is not None:
#if extracting from dflimg just copy it in order not to lose quality
shutil.copy(str(filename_path), str(output_file))
else:
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())
result.append(output_file)
if self.debug:
cv2_imwrite(debug_output_file, debug_image,
[int(cv2.IMWRITE_JPEG_QUALITY), 50])
return result
return None
def main(
detector=None,
input_path=None,
output_path=None,
output_debug=None,
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
face_type='full_face',
max_faces_from_image=None,
image_size=None,
jpeg_quality=None,
cpu_only=False,
force_gpu_idxs=None,
):
if not input_path.exists():
io.log_err('Input directory not found. Please ensure it exists.')
return
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
if face_type is not None:
face_type = FaceType.fromString(face_type)
if face_type is None:
if manual_output_debug_fix:
files = pathex.get_image_paths(output_path)
if len(files) != 0:
dflimg = DFLIMG.load(Path(files[0]))
if dflimg is not None and dflimg.has_data():
face_type = FaceType.fromString(dflimg.get_face_type())
input_image_paths = pathex.get_image_unique_filestem_paths(
input_path, verbose_print_func=io.log_info)
output_images_paths = pathex.get_image_paths(output_path)
output_debug_path = output_path.parent / (output_path.name + '_debug')
continue_extraction = False
if not manual_output_debug_fix and len(output_images_paths) > 0:
if len(output_images_paths) > 128:
continue_extraction = io.input_bool(
"Continue extraction?",
True,
help_message=
"Extraction can be continued, but you must specify the same options again."
)
if len(output_images_paths) > 128 and continue_extraction:
try:
input_image_paths = input_image_paths[
[Path(x).stem for x in input_image_paths].
index(Path(output_images_paths[-128]).stem.split('_')[0]):]
except:
io.log_err(
"Error in fetching the last index. Extraction cannot be continued."
)
return
elif input_path != output_path:
io.input(
f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n"
)
for filename in output_images_paths:
Path(filename).unlink()
device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=detector=='manual', suggest_all_gpu=True) ) \
if not cpu_only else nn.DeviceConfig.CPU()
if face_type is None:
face_type = io.input_str(
"Face type",
'wf', ['f', 'wf', 'head'],
help_message=
"Full face / whole face / head. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset."
).lower()
face_type = {
'f': FaceType.FULL,
'wf': FaceType.WHOLE_FACE,
'head': FaceType.HEAD
}[face_type]
if max_faces_from_image is None:
max_faces_from_image = io.input_int(
f"Max number of faces from image",
0,
help_message=
"If you extract a src faceset that has frames with a large number of faces, it is advisable to set max faces to 3 to speed up extraction. 0 - unlimited"
)
if image_size is None:
image_size = io.input_int(
f"Image size",
512 if face_type < FaceType.HEAD else 768,
valid_range=[256, 2048],
help_message=
"Output image size. The higher image size, the worse face-enhancer works. Use higher than 512 value only if the source image is sharp enough and the face does not need to be enhanced."
)
if jpeg_quality is None:
jpeg_quality = io.input_int(
f"Jpeg quality",
90,
valid_range=[1, 100],
help_message=
"Jpeg quality. The higher jpeg quality the larger the output file size."
)
if detector is None:
io.log_info("Choose detector type.")
io.log_info("[0] S3FD")
io.log_info("[1] manual")
detector = {0: 's3fd', 1: 'manual'}[io.input_int("", 0, [0, 1])]
if output_debug is None:
output_debug = io.input_bool(
f"Write debug images to {output_debug_path.name}?", False)
if output_debug:
output_debug_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
if not output_debug_path.exists():
io.log_err(
f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.'
)
return
else:
detector = 'manual'
io.log_info(
'Performing re-extract frames which were deleted from _debug directory.'
)
input_image_paths = DeletedFilesSearcherSubprocessor(
input_image_paths,
pathex.get_image_paths(output_debug_path)).run()
input_image_paths = sorted(input_image_paths)
io.log_info('Found %d images.' % (len(input_image_paths)))
else:
if not continue_extraction and output_debug_path.exists():
for filename in pathex.get_image_paths(output_debug_path):
Path(filename).unlink()
images_found = len(input_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info('Performing manual extract...')
data = ExtractSubprocessor(
[
ExtractSubprocessor.Data(Path(filename))
for filename in input_image_paths
],
'landmarks-manual',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
manual_window_size=manual_window_size,
device_config=device_config).run()
io.log_info('Performing 3rd pass...')
data = ExtractSubprocessor(
data,
'final',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
final_output_path=output_path,
device_config=device_config).run()
else:
io.log_info('Extracting faces...')
data = ExtractSubprocessor(
[
ExtractSubprocessor.Data(Path(filename))
for filename in input_image_paths
],
'all',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
max_faces_from_image=max_faces_from_image,
final_output_path=output_path,
device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all(np.array([d.faces_detected > 0 for d in data]) == True):
io.log_info('All faces are detected, manual fix not needed.')
else:
fix_data = [
ExtractSubprocessor.Data(d.filepath) for d in data
if d.faces_detected == 0
]
io.log_info('Performing manual fix for %d images...' %
(len(fix_data)))
fix_data = ExtractSubprocessor(
fix_data,
'landmarks-manual',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
manual_window_size=manual_window_size,
device_config=device_config).run()
fix_data = ExtractSubprocessor(
fix_data,
'final',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
final_output_path=output_path,
device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info('-------------------------')
io.log_info('Images found: %d' % (images_found))
io.log_info('Faces detected: %d' % (faces_detected))
io.log_info('-------------------------')
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
def apply_xseg(input_path, model_path):
if not input_path.exists():
raise ValueError(f'{input_path} not found. Please ensure it exists.')
if not model_path.exists():
raise ValueError(f'{model_path} not found. Please ensure it exists.')
face_type = None
model_dat = model_path / 'XSeg_data.dat'
if model_dat.exists():
dat = pickle.loads( model_dat.read_bytes() )
dat_options = dat.get('options', None)
if dat_options is not None:
face_type = dat_options.get('face_type', None)
if face_type is None:
face_type = io.input_str ("XSeg model face type", 'same', ['h','mf','f','wf','head','same'], help_message="Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'").lower()
if face_type == 'same':
face_type = None
if face_type is not None:
face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[face_type]
io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.')
device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
nn.initialize(device_config)
xseg = XSegNet(name='XSeg',
load_weights=True,
weights_file_root=model_path,
data_format=nn.data_format,
raise_on_no_model_files=True)
xseg_res = xseg.get_resolution()
images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
for filepath in io.progress_bar_generator(images_paths, "Processing"):
dflimg = DFLIMG.load(filepath)
if dflimg is None or not dflimg.has_data():
io.log_info(f'{filepath} is not a DFLIMG')
continue
img = cv2_imread(filepath).astype(np.float32) / 255.0
h,w,c = img.shape
img_face_type = FaceType.fromString( dflimg.get_face_type() )
if face_type is not None and img_face_type != face_type:
lmrks = dflimg.get_source_landmarks()
fmat = LandmarksProcessor.get_transform_mat(lmrks, w, face_type)
imat = LandmarksProcessor.get_transform_mat(lmrks, w, img_face_type)
g_p = LandmarksProcessor.transform_points (np.float32([(0,0),(w,0),(0,w) ]), fmat, True)
g_p2 = LandmarksProcessor.transform_points (g_p, imat)
mat = cv2.getAffineTransform( g_p2, np.float32([(0,0),(w,0),(0,w) ]) )
img = cv2.warpAffine(img, mat, (w, w), cv2.INTER_LANCZOS4)
img = cv2.resize(img, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
else:
if w != xseg_res:
img = cv2.resize( img, (xseg_res,xseg_res), interpolation=cv2.INTER_LANCZOS4 )
if len(img.shape) == 2:
img = img[...,None]
mask = xseg.extract(img)
if face_type is not None and img_face_type != face_type:
mask = cv2.resize(mask, (w, w), interpolation=cv2.INTER_LANCZOS4)
mask = cv2.warpAffine( mask, mat, (w,w), np.zeros( (h,w,c), dtype=np.float), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4)
mask = cv2.resize(mask, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
mask[mask < 0.5]=0
mask[mask >= 0.5]=1
dflimg.set_xseg_mask(mask)
dflimg.save()
def dev_test_68(input_dir ):
# process 68 landmarks dataset with .pts files
input_path = Path(input_dir)
if not input_path.exists():
raise ValueError('input_dir not found. Please ensure it exists.')
output_path = input_path.parent / (input_path.name+'_aligned')
io.log_info(f'Output dir is % {output_path}')
if output_path.exists():
output_images_paths = pathex.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool("WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False )
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
images_paths = pathex.get_image_paths(input_path)
for filepath in io.progress_bar_generator(images_paths, "Processing"):
filepath = Path(filepath)
pts_filepath = filepath.parent / (filepath.stem+'.pts')
if pts_filepath.exists():
pts = pts_filepath.read_text()
pts_lines = pts.split('\n')
lmrk_lines = None
for pts_line in pts_lines:
if pts_line == '{':
lmrk_lines = []
elif pts_line == '}':
break
else:
if lmrk_lines is not None:
lmrk_lines.append (pts_line)
if lmrk_lines is not None and len(lmrk_lines) == 68:
try:
lmrks = [ np.array ( lmrk_line.strip().split(' ') ).astype(np.float32).tolist() for lmrk_line in lmrk_lines]
except Exception as e:
print(e)
print(filepath)
continue
rect = LandmarksProcessor.get_rect_from_landmarks(lmrks)
output_filepath = output_path / (filepath.stem+'.jpg')
img = cv2_imread(filepath)
img = imagelib.normalize_channels(img, 3)
cv2_imwrite(output_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 95] )
DFLJPG.embed_data(output_filepath, face_type=FaceType.toString(FaceType.MARK_ONLY),
landmarks=lmrks,
source_filename=filepath.name,
source_rect=rect,
source_landmarks=lmrks
)
io.log_info("Done.")
def extract_umd_csv(input_file_csv,
image_size=256,
face_type='full_face',
device_args={}):
#extract faces from umdfaces.io dataset csv file with pitch,yaw,roll info.
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
face_type = FaceType.fromString(face_type)
input_file_csv_path = Path(input_file_csv)
if not input_file_csv_path.exists():
raise ValueError('input_file_csv not found. Please ensure it exists.')
input_file_csv_root_path = input_file_csv_path.parent
output_path = input_file_csv_path.parent / ('aligned_' +
input_file_csv_path.name)
io.log_info("Output dir is %s." % (str(output_path)))
if output_path.exists():
output_images_paths = Path_utils.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool(
"WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue."
% (str(output_path)), False)
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
try:
with open(str(input_file_csv_path), 'r') as f:
csv_file = f.read()
except Exception as e:
io.log_err("Unable to open or read file " + str(input_file_csv_path) +
": " + str(e))
return
strings = csv_file.split('\n')
keys = strings[0].split(',')
keys_len = len(keys)
csv_data = []
for i in range(1, len(strings)):
values = strings[i].split(',')
if keys_len != len(values):
io.log_err("Wrong string in csv file, skipping.")
continue
csv_data += [{keys[n]: values[n] for n in range(keys_len)}]
data = []
for d in csv_data:
filename = input_file_csv_root_path / d['FILE']
pitch, yaw, roll = float(d['PITCH']), float(d['YAW']), float(d['ROLL'])
if pitch < -90 or pitch > 90 or yaw < -90 or yaw > 90 or roll < -90 or roll > 90:
continue
pitch_yaw_roll = pitch / 90.0, yaw / 90.0, roll / 90.0
x, y, w, h = float(d['FACE_X']), float(d['FACE_Y']), float(
d['FACE_WIDTH']), float(d['FACE_HEIGHT'])
data += [
ExtractSubprocessor.Data(filename=filename,
rects=[[x, y, x + w, y + h]],
pitch_yaw_roll=pitch_yaw_roll)
]
images_found = len(data)
faces_detected = 0
if len(data) > 0:
io.log_info("Performing 2nd pass from csv file...")
data = ExtractSubprocessor(data,
'landmarks',
multi_gpu=multi_gpu,
cpu_only=cpu_only).run()
io.log_info('Performing 3rd pass...')
data = ExtractSubprocessor(data,
'final',
image_size,
face_type,
None,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in data])
io.log_info('-------------------------')
io.log_info('Images found: %d' % (images_found))
io.log_info('Faces detected: %d' % (faces_detected))
io.log_info('-------------------------')
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
},
])
])
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements({4: 32})
self.resolution = 128
self.face_type = FaceType.FULL if self.options[
'face_type'] == 'f' else FaceType.HALF
self.pose_est = PoseEstimator(
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,
generators_count=4,
sample_process_options=SampleProcessor.Options(
rotation_range=[0, 0]), #random_flip=True,
output_sample_types=[{
'types': (t.IMG_WARPED_TRANSFORMED, face_type,
t.MODE_BGR_SHUFFLE),
'resolution':
self.resolution,
'motion_blur': (25, 1)
}, {
'types':
(t.IMG_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE),
'resolution':
self.resolution
}, {
'types': (t.IMG_TRANSFORMED, face_type, t.MODE_M,
t.FACE_MASK_FULL),
'resolution':
self.resolution
}, {
'types': (t.IMG_PITCH_YAW_ROLL, )
}]),
SampleGeneratorFace(
self.training_data_dst_path,
debug=self.is_debug(),
batch_size=self.batch_size,
generators_count=4,
sample_process_options=SampleProcessor.Options(
rotation_range=[0, 0]), #random_flip=True,
output_sample_types=[{
'types':
(t.IMG_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE),
'resolution':
self.resolution
}, {
'types': (t.IMG_PITCH_YAW_ROLL, )
}])
])
def onClientProcessData(self, data):
filename_path = Path(data[0])
image = cv2.imread(str(filename_path))
if image is None:
print('Failed to extract %s, reason: cv2.imread() fail.' %
(str(filename_path)))
else:
if self.type == 'rects':
rects = self.e.extract_from_bgr(image)
return [str(filename_path), rects]
elif self.type == 'landmarks':
rects = data[1]
landmarks = self.e.extract_from_bgr(image, rects)
return [str(filename_path), landmarks]
elif self.type == 'final':
result = []
faces = data[1]
if self.debug:
debug_output_file = '{}_{}'.format(
str(
Path(str(self.output_path) + '_debug') /
filename_path.stem), 'debug.png')
debug_image = image.copy()
for (face_idx, face) in enumerate(faces):
output_file = '{}_{}{}'.format(
str(self.output_path / filename_path.stem),
str(face_idx), '.png')
rect = face[0]
image_landmarks = np.array(face[1])
if self.debug:
facelib.LandmarksProcessor.draw_rect_landmarks(
debug_image, rect, image_landmarks,
self.image_size, self.face_type)
if self.face_type == FaceType.MARK_ONLY:
face_image = image
face_image_landmarks = image_landmarks
else:
image_to_face_mat = facelib.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 = facelib.LandmarksProcessor.transform_points(
image_landmarks, image_to_face_mat)
cv2.imwrite(output_file, face_image)
a_png = AlignedPNG.load(output_file)
d = {
'face_type':
FaceType.toString(self.face_type),
'landmarks':
face_image_landmarks.tolist(),
'yaw_value':
facelib.LandmarksProcessor.calc_face_yaw(
face_image_landmarks),
'pitch_value':
facelib.LandmarksProcessor.calc_face_pitch(
face_image_landmarks),
'source_filename':
filename_path.name,
'source_rect':
rect,
'source_landmarks':
image_landmarks.tolist()
}
a_png.setFaceswapDictData(d)
a_png.save(output_file)
result.append(output_file)
if self.debug:
cv2.imwrite(debug_output_file, debug_image)
return result
return None
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({2: 1, 3: 1, 4: 4, 5: 8, 6: 16})
resolution = self.options['resolution']
face_type = self.face_type = FaceType.FULL if self.options[
'face_type'] == 'f' else FaceType.HALF
self.model = FUNIT(face_type_str=FaceType.toString(face_type),
batch_size=self.batch_size,
encoder_nf=64,
encoder_downs=2,
encoder_res_blk=2,
class_downs=4,
class_nf=64,
class_latent=64,
mlp_nf=256,
mlp_blks=2,
dis_nf=64,
dis_res_blks=10,
num_classes=2,
subpixel_decoder=True,
initialize_weights=self.is_first_run(),
is_training=self.is_training_mode,
tf_cpu_mode=self.options['optimizer_mode'] - 1)
if not self.is_first_run():
self.load_weights_safe(self.model.get_model_filename_list())
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL if self.options[
'face_type'] == 'f' else t.FACE_TYPE_HALF
if self.is_training_mode:
output_sample_types = [
{
'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'resolution': resolution,
'normalize_tanh': True
},
]
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=output_sample_types),
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=output_sample_types)
])
else:
generator = SampleGeneratorFace(
self.training_data_src_path,
batch_size=1,
sample_process_options=SampleProcessor.Options(),
output_sample_types=[{
'types': (t.IMG_SOURCE, face_type, t.MODE_BGR),
'resolution':
resolution,
'normalize_tanh':
True
}])
io.log_info("Calculating average src face style...")
codes = []
for i in io.progress_bar_generator(
range(generator.get_total_sample_count())):
codes += self.model.get_average_class_code(
generator.generate_next())
self.average_class_code = np.mean(np.array(codes), axis=0)[None,
...]
def main(input_dir,
output_dir,
debug_dir=None,
detector='mt',
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
image_size=256,
face_type='full_face',
device_args={}):
input_path = Path(input_dir)
output_path = Path(output_dir)
face_type = FaceType.fromString(face_type)
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
if not input_path.exists():
raise ValueError('输入路径不存在，请先确认workspace目录下是否存在data_dst和data_src文件夹')
if output_path.exists():
if not manual_output_debug_fix and input_path != output_path:
output_images_paths = Path_utils.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool("警告 !!! \n %s 已经有图片了! \n 继续执行会删除之前的图片. \n 按回车[Enter]继续" % (str(output_path)), False )
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
if debug_dir is None:
raise ValueError('debug-dir must be specified')
detector = 'manual'
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
if debug_dir is not None:
debug_output_path = Path(debug_dir)
if manual_output_debug_fix:
if not debug_output_path.exists():
raise ValueError("%s not found " % ( str(debug_output_path) ))
input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, Path_utils.get_image_paths(debug_output_path) ).run()
input_path_image_paths = sorted (input_path_image_paths)
io.log_info('Found %d images.' % (len(input_path_image_paths)))
else:
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info ('正在启动手动提取...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug_dir, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).run()
else:
io.log_info ('第一阶段...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in input_path_image_paths ], 'rects-'+detector, image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
io.log_info ('第二阶段...')
data = ExtractSubprocessor (data, 'landmarks', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
io.log_info ('第三阶段...')
data = ExtractSubprocessor (data, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all ( np.array ( [ d.faces_detected > 0 for d in data] ) == True ):
io.log_info ('All faces are detected, manual fix not needed.')
else:
fix_data = [ ExtractSubprocessor.Data(d.filename) for d in data if d.faces_detected == 0 ]
io.log_info ('Performing manual fix for %d images...' % (len(fix_data)) )
fix_data = ExtractSubprocessor (fix_data, 'landmarks', image_size, face_type, debug_dir, manual=True, manual_window_size=manual_window_size).run()
fix_data = ExtractSubprocessor (fix_data, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info ('-------------------------')
io.log_info ('图片数量: %d' % (images_found) )
io.log_info ('人脸数量: %d' % (faces_detected) )
io.log_info ('-------------------------')
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({4: 4})
resolution = self.options['resolution']
bgr_shape = (resolution, resolution, 3)
mask_shape = (resolution, resolution, 1)
bgrm_shape = (resolution, resolution, 4)
ngf = 64
ndf = 64
lambda_A = 100
lambda_B = 100
use_batch_norm = True #created_batch_size > 1
poseest = self.poseest = PoseEstimator(
resolution, FaceType.toString(FaceType.FULL))
self.enc = modelify(AVATARModel.DFEncFlow())([
Input(bgr_shape),
Input((poseest.class_nums[0], )),
Input((poseest.class_nums[0], )),
Input((poseest.class_nums[0], ))
])
dec_Inputs = [Input(K.int_shape(x)[1:]) for x in self.enc.outputs]
self.decA = modelify(AVATARModel.DFDecFlow(bgr_shape[2]))(dec_Inputs)
self.decB = modelify(AVATARModel.DFDecFlow(bgr_shape[2]))(dec_Inputs)
def GA(x):
return self.decA(self.enc(x))
self.GA = GA
def GB(x):
return self.decB(self.enc(x))
self.GB = GB
#self.GA = modelify(AVATARModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))( Input(bgr_shape) )
#self.GB = modelify(AVATARModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))( Input(bgr_shape) )
#self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
#self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
self.DA = modelify(
AVATARModel.NLayerDiscriminator(use_batch_norm,
ndf=ndf))(Input(bgr_shape))
self.DB = modelify(
AVATARModel.NLayerDiscriminator(use_batch_norm,
ndf=ndf))(Input(bgr_shape))
if not self.is_first_run():
weights_to_load = [
# (self.GA, 'GA.h5'),
# (self.GB, 'GB.h5'),
(self.enc, 'enc.h5'),
(self.decA, 'decA.h5'),
(self.decB, 'decB.h5'),
(self.DA, 'DA.h5'),
(self.DB, 'DB.h5'),
]
self.load_weights_safe(weights_to_load)
real_A0 = Input(bgr_shape)
real_A0m = Input(mask_shape)
real_B0 = Input(bgr_shape)
real_B0m = Input(mask_shape)
real_A0p, real_A0y, real_A0r = poseest.flow(real_A0)
real_B0p, real_B0y, real_B0r = poseest.flow(real_B0)
DA_ones = K.ones_like(K.shape(self.DA.outputs[0]))
DA_zeros = K.zeros_like(K.shape(self.DA.outputs[0]))
DB_ones = K.ones_like(K.shape(self.DB.outputs[0]))
DB_zeros = K.zeros_like(K.shape(self.DB.outputs[0]))
def DLoss(labels, logits):
return K.mean(K.binary_crossentropy(labels, logits))
def CycleLOSS(t1, t2):
return dssim(kernel_size=int(resolution / 11.6),
max_value=2.0)(t1 + 1, t2 + 1)
return K.mean(K.abs(t1 - t2))
fake_B0 = self.GA([real_A0, real_B0p, real_B0y, real_B0r])
fake_A0 = self.GB([real_B0, real_A0p, real_A0y, real_A0r])
fake_B0p, fake_B0y, fake_B0r = poseest.flow(fake_B0)
fake_A0p, fake_A0y, fake_A0r = poseest.flow(fake_A0)
real_A0_d = self.DA(real_A0)
real_A0_d_ones = K.ones_like(real_A0_d)
fake_A0_d = self.DA(fake_A0)
fake_A0_d_ones = K.ones_like(fake_A0_d)
fake_A0_d_zeros = K.zeros_like(fake_A0_d)
real_B0_d = self.DB(real_B0)
real_B0_d_ones = K.ones_like(real_B0_d)
fake_B0_d = self.DB(fake_B0)
fake_B0_d_ones = K.ones_like(fake_B0_d)
fake_B0_d_zeros = K.zeros_like(fake_B0_d)
rec_A0 = self.GB([fake_B0, real_A0p, real_A0y, real_A0r])
rec_B0 = self.GA([fake_A0, real_B0p, real_B0y, real_B0r])
#import code
#code.interact(local=dict(globals(), **locals()))
loss_GA = DLoss(fake_B0_d_ones, fake_B0_d ) + \
lambda_A * 0.1 * K.mean( K.square(fake_B0p-real_A0p) + K.square(fake_B0y-real_A0y) + K.square(fake_B0r-real_A0r) ) + \
lambda_A * (CycleLOSS(rec_B0, real_B0) )
weights_GA = self.enc.trainable_weights + self.decA.trainable_weights # + #self.GA.trainable_weights
loss_GB = DLoss(fake_A0_d_ones, fake_A0_d ) + \
lambda_B * 0.1 * K.mean( K.square(fake_A0p-real_B0p) + K.square(fake_A0y-real_B0y) + K.square(fake_A0r-real_B0r) ) + \
lambda_B * (CycleLOSS(rec_A0, real_A0) )
weights_GB = self.enc.trainable_weights + self.decB.trainable_weights # + #self.GB.trainable_weights
def opt():
return Adam(lr=2e-5, beta_1=0.5, beta_2=0.999,
tf_cpu_mode=2) #, clipnorm=1)
self.GA_train = K.function([real_A0, real_A0m, real_B0, real_B0m],
[loss_GA],
opt().get_updates(loss_GA, weights_GA))
self.GB_train = K.function([real_A0, real_A0m, real_B0, real_B0m],
[loss_GB],
opt().get_updates(loss_GB, weights_GB))
###########
loss_D_A = ( DLoss(real_A0_d_ones, real_A0_d ) + \
DLoss(fake_A0_d_zeros, fake_A0_d ) ) * 0.5
self.DA_train = K.function(
[real_A0, real_A0m, real_B0, real_B0m], [loss_D_A],
opt().get_updates(loss_D_A, self.DA.trainable_weights))
############
loss_D_B = ( DLoss(real_B0_d_ones, real_B0_d ) + \
DLoss(fake_B0_d_zeros, fake_B0_d ) ) * 0.5
self.DB_train = K.function(
[real_A0, real_A0m, real_B0, real_B0m], [loss_D_B],
opt().get_updates(loss_D_B, self.DB.trainable_weights))
############
self.G_view = K.function([real_A0, real_A0m, real_B0, real_B0m],
[fake_A0, rec_A0, fake_B0, rec_B0])
if self.is_training_mode:
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL
output_sample_types = [
{
'types': (t.IMG_SOURCE, face_type, t.MODE_BGR),
'resolution': resolution
},
{
'types':
(t.IMG_SOURCE, face_type, t.MODE_M, t.FACE_MASK_FULL),
'resolution': resolution
},
]
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=self.random_flip, normalize_tanh=True),
output_sample_types=output_sample_types),
SampleGeneratorFace(
self.training_data_dst_path,
debug=self.is_debug(),
batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(
random_flip=self.random_flip, normalize_tanh=True),
output_sample_types=output_sample_types)
])
else:
self.G_convert = K.function([real_A0, real_B0m], [fake_B0])
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.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.LandmarksExtractor(nnlib.keras)
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 = facelib.FANSegmentator(
256, FaceType.toString(FaceType.FULL))
self.e.__enter__()
elif self.type == "final":
pass
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({4: 16})
resolution = self.options['resolution']
face_type = FaceType.FULL if self.options[
'face_type'] == 'f' else FaceType.HALF
person_id_max_count = SampleGeneratorFace.get_person_id_max_count(
self.training_data_src_path)
self.model = FUNIT(face_type_str=FaceType.toString(face_type),
batch_size=self.batch_size,
encoder_nf=64,
encoder_downs=2,
encoder_res_blk=2,
class_downs=4,
class_nf=64,
class_latent=64,
mlp_nf=256,
mlp_blks=2,
dis_nf=64,
dis_res_blks=10,
num_classes=person_id_max_count,
subpixel_decoder=True,
initialize_weights=self.is_first_run(),
is_training=self.is_training_mode,
tf_cpu_mode=self.options['optimizer_mode'] - 1)
if not self.is_first_run():
self.load_weights_safe(self.model.get_model_filename_list())
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
output_sample_types = [{
'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'resolution':
128,
'normalize_tanh':
True
}]
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=output_sample_types,
person_id_mode=True),
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=output_sample_types,
person_id_mode=True),
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=output_sample_types,
person_id_mode=True),
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=output_sample_types,
person_id_mode=True),
])
def main(input_dir,
output_dir,
debug,
detector='mt',
multi_gpu=True,
cpu_only=False,
manual_fix=False,
manual_window_size=0,
image_size=256,
face_type='full_face'):
print("Running extractor.\r\n")
input_path = Path(input_dir)
output_path = Path(output_dir)
face_type = FaceType.fromString(face_type)
if not input_path.exists():
print('Input directory not found. Please ensure it exists.')
return
if output_path.exists():
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if debug:
debug_output_path = Path(str(output_path) + '_debug')
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(
input_path, verbose=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
print('Performing manual extract...')
extracted_faces = ExtractSubprocessor(
[(filename, []) for filename in input_path_image_paths],
'landmarks',
image_size,
face_type,
debug,
cpu_only=cpu_only,
manual=True,
manual_window_size=manual_window_size).process()
else:
print('Performing 1st pass...')
extracted_rects = ExtractSubprocessor(
[(x, ) for x in input_path_image_paths],
'rects',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
detector=detector).process()
print('Performing 2nd pass...')
extracted_faces = ExtractSubprocessor(extracted_rects,
'landmarks',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False).process()
if manual_fix:
print('Performing manual fix...')
if all(
np.array(
[len(data[1]) > 0
for data in extracted_faces]) == True):
print('All faces are detected, manual fix not needed.')
else:
extracted_faces = ExtractSubprocessor(
extracted_faces,
'landmarks',
image_size,
face_type,
debug,
manual=True,
manual_window_size=manual_window_size).process()
if len(extracted_faces) > 0:
print('Performing 3rd pass...')
final_imgs_paths = ExtractSubprocessor(
extracted_faces,
'final',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
output_path=output_path).process()
faces_detected = len(final_imgs_paths)
print('-------------------------')
print('Images found: %d' % (images_found))
print('Faces detected: %d' % (faces_detected))
print('-------------------------')
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
nn.initialize(data_format="NHWC")
tf = nn.tf
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.resolution = resolution = 256 #self.options['resolution']
#self.face_type = {'h' : FaceType.HALF,
# 'mf' : FaceType.MID_FULL,
# 'f' : FaceType.FULL,
# 'wf' : FaceType.WHOLE_FACE}[ self.options['face_type'] ]
self.face_type = FaceType.FULL
place_model_on_cpu = len(devices) == 0
models_opt_device = '/CPU:0' if place_model_on_cpu else '/GPU:0'
bgr_shape = nn.get4Dshape(resolution, resolution, 3)
mask_shape = nn.get4Dshape(resolution, resolution, 1)
# Initializing model classes
self.model = TernausNet(f'{self.model_name}_FANSeg',
resolution,
FaceType.toString(self.face_type),
load_weights=not self.is_first_run(),
weights_file_root=self.get_model_root_path(),
training=True,
place_model_on_cpu=place_model_on_cpu)
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_list = []
gpu_losses = []
gpu_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(
f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0'):
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
batch_slice = slice(gpu_id * bs_per_gpu,
(gpu_id + 1) * bs_per_gpu)
gpu_input_t = self.model.input_t[batch_slice, :, :, :]
gpu_target_t = self.model.target_t[
batch_slice, :, :, :]
# process model tensors
gpu_pred_logits_t, gpu_pred_t = self.model.net(
[gpu_input_t])
gpu_pred_list.append(gpu_pred_t)
gpu_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=gpu_target_t, logits=gpu_pred_logits_t),
axis=[1, 2, 3])
gpu_losses += [gpu_loss]
gpu_loss_gvs += [
nn.tf_gradients(gpu_loss, self.model.net_weights)
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
pred = nn.tf_concat(gpu_pred_list, 0)
loss = tf.reduce_mean(gpu_losses)
loss_gv_op = self.model.opt.get_update_op(
nn.tf_average_gv_list(gpu_loss_gvs))
# Initializing training and view functions
def train(input_np, target_np):
l, _ = nn.tf_sess.run([loss, loss_gv_op],
feed_dict={
self.model.input_t: input_np,
self.model.target_t: target_np
})
return l
self.train = train
def view(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.model.input_t: input_np})
self.view = view
# initializing sample generators
training_data_src_path = self.training_data_src_path
training_data_dst_path = self.training_data_dst_path
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
src_generators_count = int(src_generators_count * 1.5)
src_generator = SampleGeneratorFace(
training_data_src_path,
random_ct_samples_path=training_data_src_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True),
output_sample_types=[
{
'sample_type': SampleProcessor.SampleType.FACE_IMAGE,
'ct_mode': 'lct',
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'motion_blur': (25, 5),
'gaussian_blur': (25, 5),
'data_format': nn.data_format,
'resolution': resolution
},
{
'sample_type': SampleProcessor.SampleType.FACE_MASK,
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.G,
'face_mask_type':
SampleProcessor.FaceMaskType.FULL_FACE,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution
},
],
generators_count=src_generators_count)
dst_generator = SampleGeneratorFace(
training_data_dst_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True),
output_sample_types=[
{
'sample_type': SampleProcessor.SampleType.FACE_IMAGE,
'warp': False,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'motion_blur': (25, 5),
'gaussian_blur': (25, 5),
'data_format': nn.data_format,
'resolution': resolution
},
],
generators_count=dst_generators_count,
raise_on_no_data=False)
if not dst_generator.is_initialized():
io.log_info(
f"\nTo view the model on unseen faces, place any aligned faces in {training_data_dst_path}.\n"
)
self.set_training_data_generators([src_generator, dst_generator])
def on_cli_initialize(self):
if (self.mask_mode >= 3
and self.mask_mode <= 6) and self.fan_seg == None:
self.fan_seg = FANSegmentator(256,
FaceType.toString(FaceType.FULL))
def process_data(self, filepath):
try:
dflimg = DFLIMG.load (filepath)
if dflimg is None or not dflimg.has_data():
self.log_err (f"{filepath.name} is not a dfl image file")
else:
img = cv2_imread(filepath)
h,w = img.shape[:2]
if h != w:
raise Exception(f'w != h in {filepath}')
image_size = self.image_size
face_type = self.face_type
output_filepath = self.output_dirpath / filepath.name
if face_type is not None:
lmrks = dflimg.get_landmarks()
mat = LandmarksProcessor.get_transform_mat(lmrks, image_size, face_type)
img = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 )
img = np.clip(img, 0, 255).astype(np.uint8)
cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
dfl_dict = dflimg.get_dict()
dflimg = DFLIMG.load (output_filepath)
dflimg.set_dict(dfl_dict)
xseg_mask = dflimg.get_xseg_mask()
if xseg_mask is not None:
xseg_res = 256
xseg_lmrks = lmrks.copy()
xseg_lmrks *= (xseg_res / w)
xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type)
xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 )
xseg_mask[xseg_mask < 0.5] = 0
xseg_mask[xseg_mask >= 0.5] = 1
dflimg.set_xseg_mask(xseg_mask)
seg_ie_polys = dflimg.get_seg_ie_polys()
for poly in seg_ie_polys.get_polys():
poly_pts = poly.get_pts()
poly_pts = LandmarksProcessor.transform_points(poly_pts, mat)
poly.set_points(poly_pts)
dflimg.set_seg_ie_polys(seg_ie_polys)
lmrks = LandmarksProcessor.transform_points(lmrks, mat)
dflimg.set_landmarks(lmrks)
image_to_face_mat = dflimg.get_image_to_face_mat()
if image_to_face_mat is not None:
image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
dflimg.set_image_to_face_mat(image_to_face_mat)
dflimg.set_face_type( FaceType.toString(face_type) )
dflimg.save()
else:
dfl_dict = dflimg.get_dict()
scale = w / image_size
img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4)
cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
dflimg = DFLIMG.load (output_filepath)
dflimg.set_dict(dfl_dict)
lmrks = dflimg.get_landmarks()
lmrks /= scale
dflimg.set_landmarks(lmrks)
seg_ie_polys = dflimg.get_seg_ie_polys()
seg_ie_polys.mult_points( 1.0 / scale)
dflimg.set_seg_ie_polys(seg_ie_polys)
image_to_face_mat = dflimg.get_image_to_face_mat()
if image_to_face_mat is not None:
face_type = FaceType.fromString ( dflimg.get_face_type() )
image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
dflimg.set_image_to_face_mat(image_to_face_mat)
dflimg.save()
return (1, filepath, output_filepath)
except:
self.log_err (f"Exception occured while processing file {filepath}. Error: {traceback.format_exc()}")
return (0, filepath, None)
def process_data(self, data):
filename_path = Path(data[0])
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 None
image_shape = image.shape
if len(image_shape) == 2:
h, w = image.shape
ch = 1
else:
h, w, ch = image.shape
if ch == 1:
image = np.repeat(image[:, :, np.newaxis], 3, -1)
elif ch == 4:
image = image[:, :, 0:3]
wm = w % 2
hm = 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 == '.jpg':
src_dflimg = DFLJPG.load(str(filename_path))
if self.type == 'rects':
if min(w, h) < 128:
self.log_err('Image is too small %s : [%d, %d]' %
(str(filename_path), w, h))
rects = []
else:
rects = self.e.extract_from_bgr(image)
return [str(filename_path), rects]
elif self.type == 'landmarks':
rects = data[1]
if rects is None:
landmarks = None
else:
landmarks = self.e.extract_from_bgr(
image, rects,
self.second_pass_e if src_dflimg is None else None)
return [str(filename_path), landmarks]
elif self.type == 'final':
result = []
faces = data[1]
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(faces) != 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(faces) != 1",
str(filename_path))
output_file = str(self.output_path / filename_path.name)
if str(filename_path) != str(output_file):
shutil.copy(str(filename_path), str(output_file))
result.append(output_file)
else:
face_idx = 0
for face in faces:
rect = np.array(face[0])
image_landmarks = face[1]
if image_landmarks is None:
continue
image_landmarks = np.array(image_landmarks)
if self.face_type == FaceType.MARK_ONLY:
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:
#if extracting from dflimg copy it in order not to lose quality
output_file = str(self.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.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)
result.append(output_file)
face_idx += 1
if self.debug_dir is not None:
cv2_imwrite(debug_output_file, debug_image,
[int(cv2.IMWRITE_JPEG_QUALITY), 50])
return result
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,
)
def convert_face(self, img_bgr, img_face_landmarks, debug):
if (self.mask_mode == 3
or self.mask_mode == 4) and self.fan_seg == None:
self.fan_seg = FANSegmentator(256,
FaceType.toString(FaceType.FULL))
if self.over_res != 1:
img_bgr = cv2.resize(img_bgr, (img_bgr.shape[1] * self.over_res,
img_bgr.shape[0] * self.over_res))
img_face_landmarks = img_face_landmarks * self.over_res
if debug:
debugs = [img_bgr.copy()]
img_size = img_bgr.shape[1], img_bgr.shape[0]
img_face_mask_a = LandmarksProcessor.get_image_hull_mask(
img_bgr.shape, img_face_landmarks)
face_mat = LandmarksProcessor.get_transform_mat(
img_face_landmarks, self.output_size, face_type=self.face_type)
face_output_mat = LandmarksProcessor.get_transform_mat(
img_face_landmarks,
self.output_size,
face_type=self.face_type,
scale=self.output_face_scale)
dst_face_bgr = cv2.warpAffine(img_bgr,
face_mat,
(self.output_size, self.output_size),
flags=cv2.INTER_LANCZOS4)
dst_face_mask_a_0 = cv2.warpAffine(
img_face_mask_a,
face_mat, (self.output_size, self.output_size),
flags=cv2.INTER_LANCZOS4)
predictor_input_bgr = cv2.resize(
dst_face_bgr,
(self.predictor_input_size, self.predictor_input_size))
predictor_input_mask_a_0 = cv2.resize(
dst_face_mask_a_0,
(self.predictor_input_size, self.predictor_input_size))
predictor_input_mask_a = np.expand_dims(predictor_input_mask_a_0, -1)
predicted_bgra = self.predictor_func(
np.concatenate((predictor_input_bgr, predictor_input_mask_a), -1))
prd_face_bgr = np.clip(predicted_bgra[:, :, 0:3], 0, 1.0)
prd_face_mask_a_0 = np.clip(predicted_bgra[:, :, 3], 0.0, 1.0)
if self.mask_mode == 2: #dst
prd_face_mask_a_0 = predictor_input_mask_a_0
elif self.mask_mode == 3: #FAN-prd
prd_face_bgr_256 = cv2.resize(prd_face_bgr, (256, 256))
prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr(
np.expand_dims(prd_face_bgr_256, 0))[0]
prd_face_mask_a_0 = cv2.resize(
prd_face_bgr_256_mask,
(self.predictor_input_size, self.predictor_input_size))
elif self.mask_mode == 4: #FAN-dst
face_256_mat = LandmarksProcessor.get_transform_mat(
img_face_landmarks, 256, face_type=FaceType.FULL)
dst_face_256_bgr = cv2.warpAffine(img_bgr,
face_256_mat, (256, 256),
flags=cv2.INTER_LANCZOS4)
dst_face_256_mask = self.fan_seg.extract_from_bgr(
np.expand_dims(dst_face_256_bgr, 0))[0]
prd_face_mask_a_0 = cv2.resize(
dst_face_256_mask,
(self.predictor_input_size, self.predictor_input_size))
prd_face_mask_a_0[prd_face_mask_a_0 < 0.001] = 0.0
prd_face_mask_a = np.expand_dims(prd_face_mask_a_0, axis=-1)
prd_face_mask_aaa = np.repeat(prd_face_mask_a, (3, ), axis=-1)
img_face_mask_aaa = cv2.warpAffine(prd_face_mask_aaa,
face_output_mat,
img_size,
np.zeros(img_bgr.shape,
dtype=np.float32),
flags=cv2.WARP_INVERSE_MAP
| cv2.INTER_LANCZOS4)
img_face_mask_aaa = np.clip(img_face_mask_aaa, 0.0, 1.0)
img_face_mask_aaa[img_face_mask_aaa <= 0.1] = 0.0 #get rid of noise
if debug:
debugs += [img_face_mask_aaa.copy()]
if 'seamless' in self.mode:
#mask used for cv2.seamlessClone
img_face_seamless_mask_aaa = None
for i in range(9, 0, -1):
a = img_face_mask_aaa > i / 10.0
if len(np.argwhere(a)) == 0:
continue
img_face_seamless_mask_aaa = img_face_mask_aaa.copy()
img_face_seamless_mask_aaa[a] = 1.0
img_face_seamless_mask_aaa[img_face_seamless_mask_aaa <= i /
10.0] = 0.0
out_img = img_bgr.copy()
if self.mode == 'raw':
if self.raw_mode == 'rgb' or self.raw_mode == 'rgb-mask':
out_img = cv2.warpAffine(
prd_face_bgr, face_output_mat, img_size, out_img,
cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT)
if self.raw_mode == 'rgb-mask':
out_img = np.concatenate(
[out_img,
np.expand_dims(img_face_mask_aaa[:, :, 0], -1)], -1)
if self.raw_mode == 'mask-only':
out_img = img_face_mask_aaa
if self.raw_mode == 'predicted-only':
out_img = cv2.warpAffine(
prd_face_bgr, face_output_mat, img_size,
np.zeros(out_img.shape, dtype=np.float32),
cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT)
elif ('seamless' not in self.mode) or (img_face_seamless_mask_aaa
is not None):
#averaging [lenx, leny, maskx, masky] by grayscale gradients of upscaled mask
ar = []
for i in range(1, 10):
maxregion = np.argwhere(img_face_mask_aaa > i / 10.0)
if maxregion.size != 0:
miny, minx = maxregion.min(axis=0)[:2]
maxy, maxx = maxregion.max(axis=0)[:2]
lenx = maxx - minx
leny = maxy - miny
maskx = (minx + (lenx / 2))
masky = (miny + (leny / 2))
if lenx >= 4 and leny >= 4:
ar += [[lenx, leny, maskx, masky]]
if len(ar) > 0:
lenx, leny, maskx, masky = np.mean(ar, axis=0)
if debug:
io.log_info("lenx/leny:(%d/%d) maskx/masky:(%f/%f)" %
(lenx, leny, maskx, masky))
maskx = int(maskx)
masky = int(masky)
lowest_len = min(lenx, leny)
if debug:
io.log_info("lowest_len = %f" % (lowest_len))
img_mask_blurry_aaa = img_face_mask_aaa
if self.erode_mask_modifier != 0:
ero = int(lowest_len *
(0.126 - lowest_len * 0.00004551365) * 0.01 *
self.erode_mask_modifier)
if debug:
io.log_info("erode_size = %d" % (ero))
if ero > 0:
img_mask_blurry_aaa = cv2.erode(
img_mask_blurry_aaa,
cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(ero, ero)),
iterations=1)
elif ero < 0:
img_mask_blurry_aaa = cv2.dilate(
img_mask_blurry_aaa,
cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(-ero, -ero)),
iterations=1)
if self.seamless_erode_mask_modifier != 0:
ero = int(lowest_len *
(0.126 - lowest_len * 0.00004551365) * 0.01 *
self.seamless_erode_mask_modifier)
if debug:
io.log_info("seamless_erode_size = %d" % (ero))
if ero > 0:
img_face_seamless_mask_aaa = cv2.erode(
img_face_seamless_mask_aaa,
cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(ero, ero)),
iterations=1)
elif ero < 0:
img_face_seamless_mask_aaa = cv2.dilate(
img_face_seamless_mask_aaa,
cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(-ero, -ero)),
iterations=1)
img_face_seamless_mask_aaa = np.clip(
img_face_seamless_mask_aaa, 0, 1)
if self.clip_hborder_mask_per > 0: #clip hborder before blur
prd_hborder_rect_mask_a = np.ones(prd_face_mask_a.shape,
dtype=np.float32)
prd_border_size = int(prd_hborder_rect_mask_a.shape[1] *
self.clip_hborder_mask_per)
prd_hborder_rect_mask_a[:, 0:prd_border_size, :] = 0
prd_hborder_rect_mask_a[:, -prd_border_size:, :] = 0
prd_hborder_rect_mask_a = np.expand_dims(
cv2.blur(prd_hborder_rect_mask_a,
(prd_border_size, prd_border_size)), -1)
img_prd_hborder_rect_mask_a = cv2.warpAffine(
prd_hborder_rect_mask_a, face_output_mat, img_size,
np.zeros(img_bgr.shape, dtype=np.float32),
cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4)
img_prd_hborder_rect_mask_a = np.expand_dims(
img_prd_hborder_rect_mask_a, -1)
img_mask_blurry_aaa *= img_prd_hborder_rect_mask_a
img_mask_blurry_aaa = np.clip(img_mask_blurry_aaa, 0, 1.0)
if debug:
debugs += [img_mask_blurry_aaa.copy()]
if self.blur_mask_modifier > 0:
blur = int(lowest_len * 0.10 * 0.01 *
self.blur_mask_modifier)
if debug:
io.log_info("blur_size = %d" % (blur))
if blur > 0:
img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa,
(blur, blur))
img_mask_blurry_aaa = np.clip(img_mask_blurry_aaa, 0, 1.0)
if debug:
debugs += [img_mask_blurry_aaa.copy()]
if self.color_transfer_mode is not None:
if self.color_transfer_mode == 'rct':
if debug:
debugs += [
np.clip(
cv2.warpAffine(
prd_face_bgr, face_output_mat,
img_size,
np.zeros(img_bgr.shape,
dtype=np.float32),
cv2.WARP_INVERSE_MAP
| cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT), 0, 1.0)
]
prd_face_bgr = image_utils.reinhard_color_transfer(
np.clip((prd_face_bgr * 255).astype(np.uint8), 0,
255),
np.clip((dst_face_bgr * 255).astype(np.uint8), 0,
255),
source_mask=prd_face_mask_a,
target_mask=prd_face_mask_a)
prd_face_bgr = np.clip(
prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
if debug:
debugs += [
np.clip(
cv2.warpAffine(
prd_face_bgr, face_output_mat,
img_size,
np.zeros(img_bgr.shape,
dtype=np.float32),
cv2.WARP_INVERSE_MAP
| cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT), 0, 1.0)
]
elif self.color_transfer_mode == 'lct':
if debug:
debugs += [
np.clip(
cv2.warpAffine(
prd_face_bgr, face_output_mat,
img_size,
np.zeros(img_bgr.shape,
dtype=np.float32),
cv2.WARP_INVERSE_MAP
| cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT), 0, 1.0)
]
prd_face_bgr = image_utils.linear_color_transfer(
prd_face_bgr, dst_face_bgr)
prd_face_bgr = np.clip(prd_face_bgr, 0.0, 1.0)
if debug:
debugs += [
np.clip(
cv2.warpAffine(
prd_face_bgr, face_output_mat,
img_size,
np.zeros(img_bgr.shape,
dtype=np.float32),
cv2.WARP_INVERSE_MAP
| cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT), 0, 1.0)
]
if self.mode == 'hist-match-bw':
prd_face_bgr = cv2.cvtColor(prd_face_bgr,
cv2.COLOR_BGR2GRAY)
prd_face_bgr = np.repeat(np.expand_dims(prd_face_bgr, -1),
(3, ), -1)
if self.mode == 'hist-match' or self.mode == 'hist-match-bw':
if debug:
debugs += [
cv2.warpAffine(
prd_face_bgr, face_output_mat, img_size,
np.zeros(img_bgr.shape, dtype=np.float32),
cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT)
]
hist_mask_a = np.ones(prd_face_bgr.shape[:2] + (1, ),
dtype=np.float32)
if self.masked_hist_match:
hist_mask_a *= prd_face_mask_a
hist_match_1 = prd_face_bgr * hist_mask_a + (
1.0 - hist_mask_a) * np.ones(
prd_face_bgr.shape[:2] + (1, ), dtype=np.float32)
hist_match_1[hist_match_1 > 1.0] = 1.0
hist_match_2 = dst_face_bgr * hist_mask_a + (
1.0 - hist_mask_a) * np.ones(
prd_face_bgr.shape[:2] + (1, ), dtype=np.float32)
hist_match_2[hist_match_1 > 1.0] = 1.0
prd_face_bgr = image_utils.color_hist_match(
hist_match_1, hist_match_2, self.hist_match_threshold)
if self.mode == 'hist-match-bw':
prd_face_bgr = prd_face_bgr.astype(dtype=np.float32)
out_img = cv2.warpAffine(
prd_face_bgr, face_output_mat, img_size, out_img,
cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT)
out_img = np.clip(out_img, 0.0, 1.0)
if debug:
debugs += [out_img.copy()]
if self.mode == 'overlay':
pass
if 'seamless' in self.mode:
try:
out_img = cv2.seamlessClone(
(out_img * 255).astype(np.uint8),
(img_bgr * 255).astype(np.uint8),
(img_face_seamless_mask_aaa * 255).astype(
np.uint8), (maskx, masky), cv2.NORMAL_CLONE)
out_img = out_img.astype(dtype=np.float32) / 255.0
except Exception as e:
#seamlessClone may fail in some cases
e_str = traceback.format_exc()
if 'MemoryError' in e_str:
raise Exception(
"Seamless fail: " + e_str
) #reraise MemoryError in order to reprocess this data by other processes
else:
print("Seamless fail: " + e_str)
if debug:
debugs += [out_img.copy()]
out_img = np.clip(
img_bgr * (1 - img_mask_blurry_aaa) +
(out_img * img_mask_blurry_aaa), 0, 1.0)
if self.mode == 'seamless-hist-match':
out_face_bgr = cv2.warpAffine(
out_img, face_mat,
(self.output_size, self.output_size))
new_out_face_bgr = image_utils.color_hist_match(
out_face_bgr, dst_face_bgr, self.hist_match_threshold)
new_out = cv2.warpAffine(
new_out_face_bgr, face_mat, img_size, img_bgr.copy(),
cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4,
cv2.BORDER_TRANSPARENT)
out_img = np.clip(
img_bgr * (1 - img_mask_blurry_aaa) +
(new_out * img_mask_blurry_aaa), 0, 1.0)
if self.final_image_color_degrade_power != 0:
if debug:
debugs += [out_img.copy()]
out_img_reduced = image_utils.reduce_colors(out_img, 256)
if self.final_image_color_degrade_power == 100:
out_img = out_img_reduced
else:
alpha = self.final_image_color_degrade_power / 100.0
out_img = (out_img * (1.0 - alpha) +
out_img_reduced * alpha)
if self.alpha:
out_img = np.concatenate([
out_img,
np.expand_dims(img_mask_blurry_aaa[:, :, 0], -1)
], -1)
if self.over_res != 1:
out_img = cv2.resize(out_img, (img_bgr.shape[1] // self.over_res,
img_bgr.shape[0] // self.over_res))
out_img = np.clip(out_img, 0.0, 1.0)
if debug:
debugs += [out_img.copy()]
return debugs if debug else out_img
def main(input_dir,
output_dir,
debug_dir=None,
detector='mt',
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
image_size=256,
face_type='full_face',
max_faces_from_image=0,
device_args={}):
input_path = Path(input_dir)
output_path = Path(output_dir)
face_type = FaceType.fromString(face_type)
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
if output_path.exists():
if not manual_output_debug_fix and input_path != output_path:
output_images_paths = Path_utils.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool(
"WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue."
% (str(output_path)), False)
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
if debug_dir is None:
raise ValueError('debug-dir must be specified')
detector = 'manual'
io.log_info(
'Performing re-extract frames which were deleted from _debug directory.'
)
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(
input_path, verbose_print_func=io.log_info)
if debug_dir is not None:
debug_output_path = Path(debug_dir)
if manual_output_debug_fix:
if not debug_output_path.exists():
raise ValueError("%s not found " % (str(debug_output_path)))
input_path_image_paths = DeletedFilesSearcherSubprocessor(
input_path_image_paths,
Path_utils.get_image_paths(debug_output_path)).run()
input_path_image_paths = sorted(input_path_image_paths)
io.log_info('Found %d images.' % (len(input_path_image_paths)))
else:
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info('Performing manual extract...')
data = ExtractSubprocessor(
[
ExtractSubprocessor.Data(filename)
for filename in input_path_image_paths
],
'landmarks',
image_size,
face_type,
debug_dir,
cpu_only=cpu_only,
manual=True,
manual_window_size=manual_window_size).run()
else:
io.log_info('Performing 1st pass...')
data = ExtractSubprocessor(
[
ExtractSubprocessor.Data(filename)
for filename in input_path_image_paths
],
'rects-' + detector,
image_size,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
max_faces_from_image=max_faces_from_image).run()
io.log_info('Performing 2nd pass...')
data = ExtractSubprocessor(data,
'landmarks',
image_size,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False).run()
io.log_info('Performing 3rd pass...')
data = ExtractSubprocessor(data,
'final',
image_size,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all(np.array([d.faces_detected > 0 for d in data]) == True):
io.log_info('All faces are detected, manual fix not needed.')
else:
fix_data = [
ExtractSubprocessor.Data(d.filename) for d in data
if d.faces_detected == 0
]
io.log_info('Performing manual fix for %d images...' %
(len(fix_data)))
fix_data = ExtractSubprocessor(
fix_data,
'landmarks',
image_size,
face_type,
debug_dir,
manual=True,
manual_window_size=manual_window_size).run()
fix_data = ExtractSubprocessor(
fix_data,
'final',
image_size,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
final_output_path=output_path).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info('-------------------------')
io.log_info('Images found: %d' % (images_found))
io.log_info('Faces detected: %d' % (faces_detected))
io.log_info('-------------------------')
def onClientProcessData(self, data):
filename_path = Path(data[0])
image = cv2_imread(str(filename_path))
if image is None:
print('Failed to extract %s, reason: cv2_imread() fail.' %
(str(filename_path)))
else:
if self.type == 'rects':
rects = self.e.extract_from_bgr(image)
return [str(filename_path), rects]
elif self.type == 'landmarks':
rects = data[1]
landmarks = self.e.extract_from_bgr(image, rects)
return [str(filename_path), landmarks]
elif self.type == 'final':
result = []
faces = data[1]
if self.debug:
debug_output_file = '{}{}'.format(
str(
Path(str(self.output_path) + '_debug') /
filename_path.stem), '.jpg')
debug_image = image.copy()
for (face_idx, face) in enumerate(faces):
output_file = '{}_{}{}'.format(
str(self.output_path / filename_path.stem),
str(face_idx), '.jpg')
rect = face[0]
image_landmarks = np.array(face[1])
if self.debug:
facelib.LandmarksProcessor.draw_rect_landmarks(
debug_image, rect, image_landmarks,
self.image_size, self.face_type)
if self.face_type == FaceType.MARK_ONLY:
face_image = image
face_image_landmarks = image_landmarks
else:
image_to_face_mat = facelib.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 = facelib.LandmarksProcessor.transform_points(
image_landmarks, image_to_face_mat)
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(),
yaw_value=facelib.LandmarksProcessor.calc_face_yaw(
face_image_landmarks),
pitch_value=facelib.LandmarksProcessor.calc_face_pitch(
face_image_landmarks),
source_filename=filename_path.name,
source_rect=rect,
source_landmarks=image_landmarks.tolist())
result.append(output_file)
if self.debug:
cv2_imwrite(debug_output_file, debug_image,
[int(cv2.IMWRITE_JPEG_QUALITY), 50])
return result
return None
def main(input_dir,
output_dir,
debug=False,
detector='mt',
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
image_size=256,
face_type='full_face',
device_args={}):
input_path = Path(input_dir)
output_path = Path(output_dir)
face_type = FaceType.fromString(face_type)
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
if output_path.exists():
if not manual_output_debug_fix:
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
debug = True
detector = 'manual'
io.log_info(
'Performing re-extract frames which were deleted from _debug directory.'
)
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(
input_path, verbose_print_func=io.log_info)
if debug:
debug_output_path = Path(str(output_path) + '_debug')
if manual_output_debug_fix:
if not debug_output_path.exists():
raise ValueError("%s not found " % (str(debug_output_path)))
input_path_image_paths = DeletedFilesSearcherSubprocessor(
input_path_image_paths,
Path_utils.get_image_paths(debug_output_path)).run()
input_path_image_paths = sorted(input_path_image_paths)
else:
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info('Performing manual extract...')
extracted_faces = ExtractSubprocessor(
[(filename, []) for filename in input_path_image_paths],
'landmarks',
image_size,
face_type,
debug,
cpu_only=cpu_only,
manual=True,
manual_window_size=manual_window_size).run()
else:
io.log_info('Performing 1st pass...')
extracted_rects = ExtractSubprocessor(
[(x, ) for x in input_path_image_paths],
'rects',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
detector=detector).run()
io.log_info('Performing 2nd pass...')
extracted_faces = ExtractSubprocessor(extracted_rects,
'landmarks',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False).run()
if manual_fix:
io.log_info('Performing manual fix...')
if all(
np.array(
[len(data[1]) > 0
for data in extracted_faces]) == True):
io.log_info(
'All faces are detected, manual fix not needed.')
else:
extracted_faces = ExtractSubprocessor(
extracted_faces,
'landmarks',
image_size,
face_type,
debug,
manual=True,
manual_window_size=manual_window_size).run()
if len(extracted_faces) > 0:
io.log_info('Performing 3rd pass...')
final_imgs_paths = ExtractSubprocessor(
extracted_faces,
'final',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
output_path=output_path).run()
faces_detected = len(final_imgs_paths)
io.log_info('-------------------------')
io.log_info('Images found: %d' % (images_found))
io.log_info('Faces detected: %d' % (faces_detected))
io.log_info('-------------------------')
def main(detector=None,
input_path=None,
output_path=None,
output_debug=None,
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
face_type='full_face',
max_faces_from_image=0,
cpu_only = False,
force_gpu_idxs = None,
):
face_type = FaceType.fromString(face_type)
image_size = 512
if not input_path.exists():
io.log_err ('Input directory not found. Please ensure it exists.')
return
if detector is None:
io.log_info ("Choose detector type.")
io.log_info ("[0] S3FD")
io.log_info ("[1] manual")
detector = {0:'s3fd', 1:'manual'}[ io.input_int("", 0, [0,1]) ]
device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=detector=='manual', suggest_all_gpu=True) ) \
if not cpu_only else nn.DeviceConfig.CPU()
output_debug_path = output_path.parent / (output_path.name + '_debug')
if output_debug is None:
output_debug = io.input_bool (f"Write debug images to {output_debug_path.name}?", False)
if output_path.exists():
if not manual_output_debug_fix and input_path != output_path:
output_images_paths = pathex.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
input_path_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
if manual_output_debug_fix:
if not output_debug_path.exists():
io.log_err(f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.')
return
else:
detector = 'manual'
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, pathex.get_image_paths(output_debug_path) ).run()
input_path_image_paths = sorted (input_path_image_paths)
io.log_info('Found %d images.' % (len(input_path_image_paths)))
else:
if output_debug_path.exists():
for filename in pathex.get_image_paths(output_debug_path):
Path(filename).unlink()
else:
output_debug_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info ('Performing manual extract...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ], 'landmarks-manual', image_size, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
io.log_info ('Performing 3rd pass...')
data = ExtractSubprocessor (data, 'final', image_size, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run()
else:
io.log_info ('Extracting faces...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ],
'all',
image_size,
face_type,
output_debug_path if output_debug else None,
max_faces_from_image=max_faces_from_image,
final_output_path=output_path,
device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all ( np.array ( [ d.faces_detected > 0 for d in data] ) == True ):
io.log_info ('All faces are detected, manual fix not needed.')
else:
fix_data = [ ExtractSubprocessor.Data(d.filepath) for d in data if d.faces_detected == 0 ]
io.log_info ('Performing manual fix for %d images...' % (len(fix_data)) )
fix_data = ExtractSubprocessor (fix_data, 'landmarks-manual', image_size, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
fix_data = ExtractSubprocessor (fix_data, 'final', image_size, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info ('-------------------------')
io.log_info ('Images found: %d' % (images_found) )
io.log_info ('Faces detected: %d' % (faces_detected) )
io.log_info ('-------------------------')
def final_stage(
data,
image,
face_type,
image_size,
jpeg_quality,
extract_from_dflimg=False,
output_debug_path=None,
final_output_path=None,
):
data.final_output_files = []
filepath = data.filepath
rects = data.rects
landmarks = data.landmarks
if output_debug_path is not None:
debug_image = image.copy()
if extract_from_dflimg and len(rects) != 1:
#if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it
print("extract_from_dflimg and len(rects) != 1", filepath)
output_filepath = final_output_path / filepath.name
if filepath != str(output_file):
shutil.copy(str(filepath), str(output_filepath))
data.final_output_files.append(output_filepath)
else:
face_idx = 0
for rect, image_landmarks in zip(rects, landmarks):
if extract_from_dflimg and face_idx > 1:
#cannot extract more than 1 face from dflimg
break
if image_landmarks is None:
continue
rect = np.array(rect)
if 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, image_size, face_type)
face_image = cv2.warpAffine(image, image_to_face_mat,
(image_size, 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, image_size - 1),
(image_size - 1, image_size - 1),
(image_size - 1, 0)], image_to_face_mat, True)
rect_area = mathlib.polygon_area(
np.array(rect[[0, 2, 2, 0]]).astype(np.float32),
np.array(rect[[1, 1, 3, 3]]).astype(np.float32))
landmarks_area = mathlib.polygon_area(
landmarks_bbox[:, 0].astype(np.float32),
landmarks_bbox[:, 1].astype(np.float32))
if not data.manual and face_type <= FaceType.FULL_NO_ALIGN and landmarks_area > 4 * rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
continue
if output_debug_path is not None:
LandmarksProcessor.draw_rect_landmarks(
debug_image,
rect,
image_landmarks,
face_type,
image_size,
transparent_mask=True)
output_path = final_output_path
if data.force_output_path is not None:
output_path = data.force_output_path
if extract_from_dflimg and filepath.suffix == '.jpg':
#if extracting from dflimg and jpg copy it in order not to lose quality
output_filepath = output_path / filepath.name
if filepath != output_filepath:
shutil.copy(str(filepath), str(output_filepath))
else:
output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg"
cv2_imwrite(
output_filepath, face_image,
[int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
dflimg = DFLJPG.load(output_filepath)
dflimg.set_face_type(FaceType.toString(face_type))
dflimg.set_landmarks(face_image_landmarks.tolist())
dflimg.set_source_filename(filepath.name)
dflimg.set_source_rect(rect)
dflimg.set_source_landmarks(image_landmarks.tolist())
dflimg.set_image_to_face_mat(image_to_face_mat)
dflimg.save()
data.final_output_files.append(output_filepath)
face_idx += 1
data.faces_detected = face_idx
if output_debug_path is not None:
cv2_imwrite(output_debug_path / (filepath.stem + '.jpg'),
debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50])
return data
def get_face_type(self):
return self.dfl_dict.get('face_type', FaceType.toString(FaceType.FULL))
def extract_vggface2_dataset(input_dir, device_args={}):
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
input_path = Path(input_dir)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
bb_csv = input_path / 'loose_bb_train.csv'
if not bb_csv.exists():
raise ValueError('loose_bb_train.csv found. Please ensure it exists.')
bb_lines = bb_csv.read_text().split('\n')
bb_lines.pop(0)
bb_dict = {}
for line in bb_lines:
name, l, t, w, h = line.split(',')
name = name[1:-1]
l, t, w, h = [int(x) for x in (l, t, w, h)]
bb_dict[name] = (l, t, w, h)
output_path = input_path.parent / (input_path.name + '_out')
dir_names = pathex.get_all_dir_names(input_path)
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
data = []
for dir_name in io.progress_bar_generator(dir_names, "Collecting"):
cur_input_path = input_path / dir_name
cur_output_path = output_path / dir_name
if not cur_output_path.exists():
cur_output_path.mkdir(parents=True, exist_ok=True)
input_path_image_paths = pathex.get_image_paths(cur_input_path)
for filename in input_path_image_paths:
filename_path = Path(filename)
name = filename_path.parent.name + '/' + filename_path.stem
if name not in bb_dict:
continue
l, t, w, h = bb_dict[name]
if min(w, h) < 128:
continue
data += [
ExtractSubprocessor.Data(filename=filename,
rects=[(l, t, l + w, t + h)],
landmarks_accurate=False,
force_output_path=cur_output_path)
]
face_type = FaceType.fromString('full_face')
io.log_info('Performing 2nd pass...')
data = ExtractSubprocessor(data,
'landmarks',
256,
face_type,
debug_dir=None,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False).run()
io.log_info('Performing 3rd pass...')
ExtractSubprocessor(data,
'final',
256,
face_type,
debug_dir=None,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
final_output_path=None).run()
def process_data(self, data):
filename_path = Path(data[0])
filename_path_str = str(filename_path)
if self.cached_image[0] == filename_path_str:
image = self.cached_image[1]
else:
image = cv2_imread(filename_path_str)
self.cached_image = (filename_path_str, image)
if image is None:
self.log_err(
'Failed to extract %s, reason: cv2_imread() fail.' %
(str(filename_path)))
else:
if self.type == 'rects':
rects = self.e.extract_from_bgr(image)
return [str(filename_path), rects]
elif self.type == 'landmarks':
rects = data[1]
landmarks = self.e.extract_from_bgr(image, rects)
return [str(filename_path), landmarks]
elif self.type == 'final':
src_dflimg = None
(h, w, c) = image.shape
if h == w:
#extracting from already extracted jpg image?
if filename_path.suffix == '.jpg':
src_dflimg = DFLJPG.load(str(filename_path))
result = []
faces = data[1]
if self.debug:
debug_output_file = '{}{}'.format(
str(
Path(str(self.output_path) + '_debug') /
filename_path.stem), '.jpg')
debug_image = image.copy()
face_idx = 0
for face in faces:
rect = np.array(face[0])
image_landmarks = np.array(face[1])
if self.face_type == FaceType.MARK_ONLY:
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:
LandmarksProcessor.draw_rect_landmarks(
debug_image, rect, image_landmarks,
self.image_size, self.face_type)
output_file = '{}_{}{}'.format(
str(self.output_path / filename_path.stem),
str(face_idx), '.jpg')
face_idx += 1
if src_dflimg is not None:
#if extracting from dflimg just copy it in order not to lose quality
shutil.copy(str(filename_path), str(output_file))
else:
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())
result.append(output_file)
if self.debug:
cv2_imwrite(debug_output_file, debug_image,
[int(cv2.IMWRITE_JPEG_QUALITY), 50])
return result
return None