def onInitializeOptions(self, is_first_run, ask_override):
if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
def ask_settings(self):
self.add_source_image = io.input_bool(
"Add source image? (y/n ?:help skip:n) : ",
False,
help_message="Add source image for comparison.")
super().ask_settings()
def main(args, device_args):
io.log_info("Running converter.\r\n")
aligned_dir = args.get('aligned_dir', None)
avaperator_aligned_dir = args.get('avaperator_aligned_dir', None)
try:
input_path = Path(args['input_dir'])
output_path = Path(args['output_dir'])
model_path = Path(args['model_dir'])
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 not model_path.exists():
io.log_err('Model directory not found. Please ensure it exists.')
return
is_interactive = io.input_bool(
"Use interactive converter? (y/n skip:y) : ",
True) if not io.is_colab() else False
import models
model = models.import_model(args['model_name'])(
model_path, device_args=device_args)
converter_session_filepath = model.get_strpath_storage_for_file(
'converter_session.dat')
predictor_func, predictor_input_shape, cfg = model.get_ConverterConfig(
)
if not is_interactive:
cfg.ask_settings()
input_path_image_paths = Path_utils.get_image_paths(input_path)
if cfg.type == ConverterConfig.TYPE_MASKED:
if aligned_dir is None:
io.log_err(
'Aligned directory not found. Please ensure it exists.')
return
aligned_path = Path(aligned_dir)
if not aligned_path.exists():
io.log_err(
'Aligned directory not found. Please ensure it exists.')
return
alignments = {}
multiple_faces_detected = False
aligned_path_image_paths = Path_utils.get_image_paths(aligned_path)
for filepath in io.progress_bar_generator(aligned_path_image_paths,
"Collecting alignments"):
filepath = Path(filepath)
if filepath.suffix == '.png':
dflimg = DFLPNG.load(str(filepath))
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load(str(filepath))
else:
dflimg = None
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
continue
source_filename_stem = Path(dflimg.get_source_filename()).stem
if source_filename_stem not in alignments.keys():
alignments[source_filename_stem] = []
alignments_ar = alignments[source_filename_stem]
alignments_ar.append(dflimg.get_source_landmarks())
if len(alignments_ar) > 1:
multiple_faces_detected = True
if multiple_faces_detected:
io.log_info(
"Warning: multiple faces detected. Strongly recommended to process them separately."
)
frames = [
ConvertSubprocessor.Frame(frame_info=FrameInfo(
filename=p,
landmarks_list=alignments.get(Path(p).stem, None)))
for p in input_path_image_paths
]
if multiple_faces_detected:
io.log_info(
"Warning: multiple faces detected. Motion blur will not be used."
)
else:
s = 256
local_pts = [(s // 2 - 1, s // 2 - 1),
(s // 2 - 1, 0)] #center+up
frames_len = len(frames)
for i in io.progress_bar_generator(range(len(frames)),
"Computing motion vectors"):
fi_prev = frames[max(0, i - 1)].frame_info
fi = frames[i].frame_info
fi_next = frames[min(i + 1, frames_len - 1)].frame_info
if len(fi_prev.landmarks_list) == 0 or \
len(fi.landmarks_list) == 0 or \
len(fi_next.landmarks_list) == 0:
continue
mat_prev = LandmarksProcessor.get_transform_mat(
fi_prev.landmarks_list[0], s, face_type=FaceType.FULL)
mat = LandmarksProcessor.get_transform_mat(
fi.landmarks_list[0], s, face_type=FaceType.FULL)
mat_next = LandmarksProcessor.get_transform_mat(
fi_next.landmarks_list[0], s, face_type=FaceType.FULL)
pts_prev = LandmarksProcessor.transform_points(
local_pts, mat_prev, True)
pts = LandmarksProcessor.transform_points(
local_pts, mat, True)
pts_next = LandmarksProcessor.transform_points(
local_pts, mat_next, True)
prev_vector = pts[0] - pts_prev[0]
next_vector = pts_next[0] - pts[0]
motion_vector = pts_next[0] - pts_prev[0]
fi.motion_power = npla.norm(motion_vector)
motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array(
[0, 0], dtype=np.float32)
fi.motion_deg = -math.atan2(
motion_vector[1], motion_vector[0]) * 180 / math.pi
elif cfg.type == ConverterConfig.TYPE_FACE_AVATAR:
filesdata = []
for filepath in io.progress_bar_generator(input_path_image_paths,
"Collecting info"):
filepath = Path(filepath)
if filepath.suffix == '.png':
dflimg = DFLPNG.load(str(filepath))
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load(str(filepath))
else:
dflimg = None
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
continue
filesdata += [
(FrameInfo(filename=str(filepath),
landmarks_list=[dflimg.get_landmarks()]),
dflimg.get_source_filename())
]
filesdata = sorted(filesdata,
key=operator.itemgetter(1)) #sort by filename
frames = []
filesdata_len = len(filesdata)
for i in range(len(filesdata)):
frame_info = filesdata[i][0]
prev_temporal_frame_infos = []
next_temporal_frame_infos = []
for t in range(cfg.temporal_face_count):
prev_frame_info = filesdata[max(i - t, 0)][0]
next_frame_info = filesdata[min(i + t,
filesdata_len - 1)][0]
prev_temporal_frame_infos.insert(0, prev_frame_info)
next_temporal_frame_infos.append(next_frame_info)
frames.append(
ConvertSubprocessor.Frame(
prev_temporal_frame_infos=prev_temporal_frame_infos,
frame_info=frame_info,
next_temporal_frame_infos=next_temporal_frame_infos))
if len(frames) == 0:
io.log_info("No frames to convert in input_dir.")
else:
ConvertSubprocessor(
is_interactive=is_interactive,
converter_session_filepath=converter_session_filepath,
predictor_func=predictor_func,
predictor_input_shape=predictor_input_shape,
converter_config=cfg,
frames=frames,
output_path=output_path,
model_iter=model.get_iter()).run()
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
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(
"选择模式: (1)覆盖，(2)直方图匹配，(3)直方图匹配白平衡，(4)无缝，(5)raw. 默认 - %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(
"选择raw模式: (1) rgb， (2) rgb+掩码(默认)，(3)仅掩码，(4)仅预测 : ", 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("无缝直方图匹配? (y/n 默认: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(
"面部遮罩直方图匹配? (y/n 默认: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("直方图匹配阈值 [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(
"面部遮罩模式: (1) 学习, (2) dst原始视频, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) 帮助. 默认 - %d : "
% (1),
1,
help_message=
"如果你学过蒙版,那么选择选项1.“dst”遮罩是原始的抖动遮罩从dst对齐的图像.“扇-prd”-使用超光滑的面具,通过预先训练的扇模型从预测的脸.“风扇-dst”-使用超光滑的面具,由预先训练的风扇模型从dst的脸.“FAN-prd*FAN-dst”或“learned*FAN-prd*FAN-dst”——使用多个口罩."
), 1, 6)
else:
self.mask_mode = np.clip(
io.input_int("面部遮罩模式: (1) 学习, (2) dst . 默认 - %d : " % (1),
1), 1, 2)
if self.mask_mode >= 3 and self.mask_mode <= 6:
self.fan_seg = None
if self.mode != 'raw':
self.erode_mask_modifier = base_erode_mask_modifier + np.clip(
io.input_int(
"侵蚀遮罩 [-200..200] (默认:%d) : " %
(default_erode_mask_modifier),
default_erode_mask_modifier), -200, 200)
self.blur_mask_modifier = base_blur_mask_modifier + np.clip(
io.input_int(
"选择模糊遮罩边缘 [-200..200] (默认:%d) : " %
(default_blur_mask_modifier), default_blur_mask_modifier),
-200, 200)
self.output_face_scale = np.clip(
1.0 + io.input_int("选择输出脸部比例调整器 [-50..50] (默认:0) : ", 0) * 0.01,
0.5, 1.5)
if self.mode != 'raw':
self.color_transfer_mode = io.input_str(
"应用颜色转移到预测的脸? 选择模式 ( rct/lct 默认:None ) : ", None,
['rct', 'lct'])
self.super_resolution = io.input_bool("应用超分辨率? (y/n ?:帮助 默认:n) : ",
False,
help_message="通过应用DCSCN网络增强细节.")
if self.mode != 'raw':
self.final_image_color_degrade_power = np.clip(
io.input_int("降低最终图像色权 [0..100] (默认:0) : ", 0), 0, 100)
self.alpha = io.input_bool("使用alpha通道导出png? (y/n 默认: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
def ask_settings(self):
s = """Choose mode: \n"""
for key in mode_dict.keys():
s += f"""({key}) {mode_dict[key]}\n"""
s += f"""Default: {self.default_mode} : """
mode = io.input_int(s, self.default_mode)
self.mode = mode_dict.get(mode, mode_dict[self.default_mode])
if 'raw' not in self.mode:
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 self.face_type == FaceType.FULL:
s = """Choose mask mode: \n"""
for key in full_face_mask_mode_dict.keys():
s += f"""({key}) {full_face_mask_mode_dict[key]}\n"""
s += f"""?:help Default: 1 : """
self.mask_mode = io.input_int(
s,
1,
valid_list=full_face_mask_mode_dict.keys(),
help_message=
"If you learned the 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."
)
else:
s = """Choose mask mode: \n"""
for key in half_face_mask_mode_dict.keys():
s += f"""({key}) {half_face_mask_mode_dict[key]}\n"""
s += f"""?:help , Default: 1 : """
self.mask_mode = io.input_int(
s,
1,
valid_list=half_face_mask_mode_dict.keys(),
help_message=
"If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images."
)
if 'raw' not in self.mode:
self.erode_mask_modifier = np.clip(
io.input_int(
"Choose erode mask modifier [-400..400] (skip:%d) : " % 0,
0), -400, 400)
self.blur_mask_modifier = np.clip(
io.input_int(
"Choose blur mask modifier [-400..400] (skip:%d) : " % 0,
0), -400, 400)
self.motion_blur_power = np.clip(
io.input_int(
"Choose motion blur power [0..100] (skip:%d) : " % (0), 0),
0, 100)
self.output_face_scale = np.clip(
io.input_int(
"Choose output face scale modifier [-50..50] (skip:0) : ", 0),
-50, 50)
if 'raw' not in self.mode:
self.color_transfer_mode = io.input_str(
"Apply color transfer to predicted face? Choose mode ( rct/lct/ebs skip:None ) : ",
None, ctm_str_dict.keys())
self.color_transfer_mode = ctm_str_dict[self.color_transfer_mode]
super().ask_settings()
if 'raw' not in self.mode:
self.color_degrade_power = np.clip(
io.input_int(
"Degrade color power of final image [0..100] (skip:0) : ",
0), 0, 100)
self.export_mask_alpha = io.input_bool(
"Export png with alpha channel of the mask? (y/n skip:n) : ",
False)
io.log_info("")
def pack(samples_path):
samples_dat_path = samples_path / packed_faceset_filename
if samples_dat_path.exists():
io.log_info(f"{samples_dat_path} : file already exists !")
io.input_bool("Press enter to continue and overwrite.", False)
as_person_faceset = False
dir_names = Path_utils.get_all_dir_names(samples_path)
if len(dir_names) != 0:
as_person_faceset = io.input_bool(
f"{len(dir_names)} subdirectories found, process as person faceset? (y/n) skip:y : ",
True)
if as_person_faceset:
image_paths = []
for dir_name in dir_names:
image_paths += Path_utils.get_image_paths(samples_path /
dir_name)
else:
image_paths = Path_utils.get_image_paths(samples_path)
samples = samplelib.SampleHost.load_face_samples(image_paths)
samples_len = len(samples)
samples_configs = []
for sample in io.progress_bar_generator(samples, "Processing"):
sample_filepath = Path(sample.filename)
sample.filename = sample_filepath.name
if as_person_faceset:
sample.person_name = sample_filepath.parent.name
samples_configs.append(sample.get_config())
samples_bytes = pickle.dumps(samples_configs, 4)
of = open(samples_dat_path, "wb")
of.write(struct.pack("Q", PackedFaceset.VERSION))
of.write(struct.pack("Q", len(samples_bytes)))
of.write(samples_bytes)
del samples_bytes #just free mem
del samples_configs
sample_data_table_offset = of.tell()
of.write(bytes(8 * (samples_len + 1))) #sample data offset table
data_start_offset = of.tell()
offsets = []
for sample in io.progress_bar_generator(samples, "Packing"):
try:
if sample.person_name is not None:
sample_path = samples_path / sample.person_name / sample.filename
else:
sample_path = samples_path / sample.filename
with open(sample_path, "rb") as f:
b = f.read()
offsets.append(of.tell() - data_start_offset)
of.write(b)
except:
raise Exception(f"error while processing sample {sample_path}")
offsets.append(of.tell())
of.seek(sample_data_table_offset, 0)
for offset in offsets:
of.write(struct.pack("Q", offset))
of.seek(0, 2)
of.close()
for filename in io.progress_bar_generator(image_paths,
"Deleting files"):
Path(filename).unlink()
if as_person_faceset:
for dir_name in io.progress_bar_generator(dir_names,
"Deleting dirs"):
dir_path = samples_path / dir_name
try:
shutil.rmtree(dir_path)
except:
io.log_info(f"unable to remove: {dir_path} ")
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 __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.mask_mode >= 3 and self.mask_mode <= 6:
self.fan_seg = None
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
def onInitializeOptions(self, is_first_run, ask_override):
yn_str = {True: 'y', False: 'n'}
default_resolution = 128
default_archi = 'df'
default_face_type = 'f'
if is_first_run:
resolution = io.input_int(
"Resolution ( 16-1024 ?:help skip:128) : ",
default_resolution,
help_message=
"More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16."
)
resolution = np.clip(resolution, 16, 1024)
while np.modf(resolution / 16)[0] != 0.0:
resolution -= 1
self.options['resolution'] = resolution
self.options['face_type'] = io.input_str(
"Half or Full face? (h/f, ?:help skip:f) : ",
default_face_type, ['h', 'f'],
help_message=
"Half face has better resolution, but covers less area of cheeks."
).lower()
else:
self.options['resolution'] = self.options.get(
'resolution', default_resolution)
self.options['face_type'] = self.options.get(
'face_type', default_face_type)
if is_first_run or ask_override:
default_learn_mask = self.options.get('learn_mask', True)
self.options['learn_mask'] = io.input_bool(
f'Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]}) : ',
default_learn_mask,
help_message=
"Learning mask can help model to recognize face directions. Learn without mask can reduce "
"model size, in this case converter forced to use 'not predicted mask' that is not smooth "
"as predicted. Model with style values can be learned without mask and produce same "
"quality result.")
if (is_first_run or
ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int(
"Optimizer mode? ( 1,2,3 ?:help skip:%d) : " %
(def_optimizer_mode),
def_optimizer_mode,
help_message=
"1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power."
)
else:
self.options['optimizer_mode'] = self.options.get(
'optimizer_mode', 1)
if is_first_run:
self.options['archi'] = io.input_str(
"AE architecture (df, liae ?:help skip:%s) : " %
(default_archi),
default_archi, ['df', 'liae'],
help_message=
"'df' keeps faces more natural. 'liae' can fix overly different face shapes."
).lower(
) # -s version is slower, but has decreased change to collapse.
else:
self.options['archi'] = self.options.get('archi', default_archi)
default_ae_dims = 256 if 'liae' in self.options['archi'] else 512
default_e_ch_dims = 42
default_d_ch_dims = default_e_ch_dims // 2
def_ca_weights = False
if is_first_run:
self.options['ae_dims'] = np.clip(
io.input_int(
"AutoEncoder dims (1-2048 ?:help skip:%d) : " %
(default_ae_dims),
default_ae_dims,
help_message=
"All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU."
), 1, 2048)
self.options['e_ch_dims'] = np.clip(
io.input_int(
"Encoder dims per channel (1-128 ?:help skip:%d) : " %
(default_e_ch_dims),
default_e_ch_dims,
help_message=
"More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU."
), 1, 128)
default_d_ch_dims = self.options['e_ch_dims'] // 2
self.options['d_ch_dims'] = np.clip(
io.input_int(
"Decoder dims per channel (1-128 ?:help skip:%d) : " %
(default_d_ch_dims),
default_d_ch_dims,
help_message=
"More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU."
), 1, 128)
self.options['multiscale_decoder'] = io.input_bool(
"Use multiscale decoder? (y/n, ?:help skip:n) : ",
False,
help_message="Multiscale decoder helps to get better details.")
self.options['ca_weights'] = io.input_bool(
"Use CA weights? (y/n, ?:help skip: %s ) : " %
(yn_str[def_ca_weights]),
def_ca_weights,
help_message=
"Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run."
)
else:
self.options['ae_dims'] = self.options.get('ae_dims',
default_ae_dims)
self.options['e_ch_dims'] = self.options.get(
'e_ch_dims', default_e_ch_dims)
self.options['d_ch_dims'] = self.options.get(
'd_ch_dims', default_d_ch_dims)
self.options['multiscale_decoder'] = self.options.get(
'multiscale_decoder', False)
self.options['ca_weights'] = self.options.get(
'ca_weights', def_ca_weights)
default_face_style_power = 0.0
default_bg_style_power = 0.0
if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = io.input_bool(
"Use pixel loss? (y/n, ?:help skip: %s ) : " %
(yn_str[def_pixel_loss]),
def_pixel_loss,
help_message=
"Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. Enabling this option too early increases the chance of model collapse."
)
default_face_style_power = default_face_style_power if is_first_run else self.options.get(
'face_style_power', default_face_style_power)
self.options['face_style_power'] = np.clip(
io.input_number(
"Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " %
(default_face_style_power),
default_face_style_power,
help_message=
"Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
default_bg_style_power = default_bg_style_power if is_first_run else self.options.get(
'bg_style_power', default_bg_style_power)
self.options['bg_style_power'] = np.clip(
io.input_number(
"Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : "
% (default_bg_style_power),
default_bg_style_power,
help_message=
"Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
default_apply_random_ct = ColorTransferMode.NONE if is_first_run else self.options.get(
'apply_random_ct', ColorTransferMode.NONE)
self.options['apply_random_ct'] = np.clip(
io.input_int(
"Apply random color transfer to src faceset? (0) None, (1) LCT, (2) RCT, (3) RCT-c, (4) RCT-p, "
"(5) RCT-pc, (6) mRTC, (7) mRTC-c, (8) mRTC-p, (9) mRTC-pc ?:help skip:%s) : "
% default_apply_random_ct,
default_apply_random_ct,
help_message=
"Increase variativity of src samples by apply LCT color transfer from random dst "
"samples. It is like 'face_style' learning, but more precise color transfer and without "
"risk of model collapse, also it does not require additional GPU resources, "
"but the training time may be longer, due to the src faceset is becoming more diverse."
), ColorTransferMode.NONE,
ColorTransferMode.MASKED_RCT_PAPER_CLIP)
default_random_color_change = False if is_first_run else self.options.get(
'random_color_change', False)
self.options['random_color_change'] = io.input_bool(
"Enable random color change? (y/n, ?:help skip:%s) : " %
(yn_str[default_random_color_change]),
default_random_color_change,
help_message="")
if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301
default_clipgrad = False if is_first_run else self.options.get(
'clipgrad', False)
self.options['clipgrad'] = io.input_bool(
"Enable gradient clipping? (y/n, ?:help skip:%s) : " %
(yn_str[default_clipgrad]),
default_clipgrad,
help_message=
"Gradient clipping reduces chance of model collapse, sacrificing speed of training."
)
else:
self.options['clipgrad'] = False
self.options['pretrain'] = io.input_bool(
"Pretrain the model? (y/n, ?:help skip:n) : ",
False,
help_message="Pretrain the model with large amount of various "
"faces. This technique may help to train the fake "
"with overly different face shapes and light "
"conditions of src/dst data. Face will be look more "
"like a morphed. To reduce the morph effect, "
"some model files will be initialized but not be "
"updated after pretrain: LIAE: inter_AB.h5 DF: "
"encoder.h5. The longer you pretrain the model the "
"more morphed face will look. After that, "
"save and run the training again.")
else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
self.options['face_style_power'] = self.options.get(
'face_style_power', default_face_style_power)
self.options['bg_style_power'] = self.options.get(
'bg_style_power', default_bg_style_power)
self.options['apply_random_ct'] = self.options.get(
'apply_random_ct', ColorTransferMode.NONE)
self.options['clipgrad'] = self.options.get('clipgrad', False)
self.options['random_color_change'] = self.options.get(
'random_color_change', False)
self.options['pretrain'] = self.options.get('pretrain', False)
def __init__(self,
predictor_func,
predictor_input_size=0,
output_size=0,
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):
super().__init__(predictor_func, Converter.TYPE_FACE)
self.predictor_input_size = predictor_input_size
self.output_size = output_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])
self.suppress_seamless_jitter = False
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':
io.input_bool(
"Suppress seamless jitter? [ y/n ] (?:help skip:n ) : ",
False,
help_message=
"Seamless clone produces face jitter. You can suppress it, but process can take a long time."
)
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 face_type == FaceType.FULL:
self.mask_mode = io.input_int(
"Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) 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."
)
else:
self.mask_mode = io.input_int(
"Mask mode: (1) learned, (2) dst . Default - %d : " % (1), 1)
if self.mask_mode == 3 or self.mask_mode == 4:
self.fan_seg = None
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.seamless_erode_mask_modifier = 0
if 'seamless' in self.mode:
self.seamless_erode_mask_modifier = np.clip(
io.input_int(
"Choose seamless erode mask modifier [-100..100] (skip:0) : ",
0), -100, 100)
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)
self.color_transfer_mode = io.input_str(
"Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ",
None, ['rct', 'lct'])
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("")
self.over_res = 4 if self.suppress_seamless_jitter else 1
def onInitializeOptions(self, is_first_run, ask_override):
yn_str = {True: 'y', False: 'n'}
default_resolution = 128
default_archi = 'df'
default_face_type = 'f'
if is_first_run:
resolution = io.input_int(
"像素[Resolution=]( 64-256 ?:help skip:128) : ",
default_resolution,
help_message="更高的分辨率需要更多的VRAM和训练时间。取值为16的倍数.")
resolution = np.clip(resolution, 64, 256)
while np.modf(resolution / 16)[0] != 0.0:
resolution -= 1
self.options['resolution'] = resolution
self.options['face_type'] = io.input_str(
"全脸还是半脸[Half or Full face]? (h/f, ?:help skip:f) : ",
default_face_type, ['h', 'f'],
help_message=
"Half face has better resolution, but covers less area of cheeks."
).lower()
self.options['learn_mask'] = io.input_bool(
"Learn mask? (y/n, ?:help skip:y) : ",
True,
help_message=
"Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result."
)
else:
self.options['resolution'] = self.options.get(
'resolution', default_resolution)
self.options['face_type'] = self.options.get(
'face_type', default_face_type)
self.options['learn_mask'] = self.options.get('learn_mask', True)
if (is_first_run or
ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int(
"优化模式[Optimizer mode]? ( 1,2,3 ?:help skip:%d) : " %
(def_optimizer_mode),
def_optimizer_mode,
help_message=
"1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power."
)
else:
self.options['optimizer_mode'] = self.options.get(
'optimizer_mode', 1)
if is_first_run:
self.options['archi'] = io.input_str(
"模型结构[AE architecture] (df, liae ?:help skip:%s) : " %
(default_archi),
default_archi, ['df', 'liae'],
help_message=
"'df' keeps faces more natural. 'liae' can fix overly different face shapes."
).lower(
) #-s version is slower, but has decreased change to collapse.
else:
self.options['archi'] = self.options.get('archi', default_archi)
default_ae_dims = 256 if 'liae' in self.options['archi'] else 512
default_e_ch_dims = 42
default_d_ch_dims = default_e_ch_dims // 2
def_ca_weights = False
if is_first_run:
self.options['ae_dims'] = np.clip(
io.input_int(
"自动编码器维度[AutoEncoder dims] (32-1024 ?:help skip:%d) : " %
(default_ae_dims),
default_ae_dims,
help_message=
"All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU."
), 32, 1024)
self.options['e_ch_dims'] = np.clip(
io.input_int(
"每个通道编码器维度 [Encoder dims per channel] (21-85 ?:help skip:%d) : "
% (default_e_ch_dims),
default_e_ch_dims,
help_message=
"More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU."
), 21, 85)
default_d_ch_dims = self.options['e_ch_dims'] // 2
self.options['d_ch_dims'] = np.clip(
io.input_int(
"每个通道解码器维度 [Decoder dims per channel] (10-85 ?:help skip:%d) : "
% (default_d_ch_dims),
default_d_ch_dims,
help_message=
"More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU."
), 10, 85)
self.options['multiscale_decoder'] = io.input_bool(
"使用多尺度解码器 [Use multiscale decoder]? (y/n, ?:help skip:n) : ",
False,
help_message="Multiscale decoder helps to get better details.")
self.options['ca_weights'] = io.input_bool(
"使用CA权重[Use CA weights]? (y/n, ?:help skip: %s ) : " %
(yn_str[def_ca_weights]),
def_ca_weights,
help_message=
"Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run."
)
else:
self.options['ae_dims'] = self.options.get('ae_dims',
default_ae_dims)
self.options['e_ch_dims'] = self.options.get(
'e_ch_dims', default_e_ch_dims)
self.options['d_ch_dims'] = self.options.get(
'd_ch_dims', default_d_ch_dims)
self.options['multiscale_decoder'] = self.options.get(
'multiscale_decoder', False)
self.options['ca_weights'] = self.options.get(
'ca_weights', def_ca_weights)
default_face_style_power = 0.0
default_bg_style_power = 0.0
if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = io.input_bool(
"使用像素丢失[pixel loss]? (y/n, ?:help skip: %s ) : " %
(yn_str[def_pixel_loss]),
def_pixel_loss,
help_message="像素丢失可能有助于增强细节和稳定面部颜色。 仅在质量不随时间改善时使用。")
default_face_style_power = default_face_style_power if is_first_run else self.options.get(
'face_style_power', default_face_style_power)
self.options['face_style_power'] = np.clip(
io.input_number(
"脸部风格强度[Face style power] ( 0.0 .. 100.0 ?:help skip:%.2f) : "
% (default_face_style_power),
default_face_style_power,
help_message=
"Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
default_bg_style_power = default_bg_style_power if is_first_run else self.options.get(
'bg_style_power', default_bg_style_power)
self.options['bg_style_power'] = np.clip(
io.input_number(
"背景风格强度[Background style power] ( 0.0 .. 100.0 ?:help skip:%.2f) : "
% (default_bg_style_power),
default_bg_style_power,
help_message=
"Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
default_apply_random_ct = False if is_first_run else self.options.get(
'apply_random_ct', False)
self.options['apply_random_ct'] = io.input_bool(
"随机颜色[Apply random color transfer to src faceset]? (y/n, ?:help skip:%s) : "
% (yn_str[default_apply_random_ct]),
default_apply_random_ct,
help_message=
"Increase variativity of src samples by apply LCT color transfer from random dst samples. It is like 'face_style' learning, but more precise color transfer and without risk of model collapse, also it does not require additional GPU resources, but the training time may be longer, due to the src faceset is becoming more diverse."
)
if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301
default_clipgrad = False if is_first_run else self.options.get(
'clipgrad', False)
self.options['clipgrad'] = io.input_bool(
"启用渐变剪裁[Enable gradient clipping]? (y/n, ?:help skip:%s) : "
% (yn_str[default_clipgrad]),
default_clipgrad,
help_message=
"Gradient clipping reduces chance of model collapse, sacrificing speed of training."
)
else:
self.options['clipgrad'] = False
else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
self.options['face_style_power'] = self.options.get(
'face_style_power', default_face_style_power)
self.options['bg_style_power'] = self.options.get(
'bg_style_power', default_bg_style_power)
self.options['apply_random_ct'] = self.options.get(
'apply_random_ct', False)
self.options['clipgrad'] = self.options.get('clipgrad', False)
if is_first_run:
self.options['pretrain'] = io.input_bool(
"模型预训练[Pretrain the model]? (y/n, ?:help skip:n) : ",
False,
help_message=
"Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again."
)
else:
self.options['pretrain'] = False
def video_from_sequence(input_dir,
output_file,
reference_file=None,
ext=None,
fps=None,
bitrate=None,
lossless=None):
input_path = Path(input_dir)
output_file_path = Path(output_file)
reference_file_path = Path(
reference_file) if reference_file is not None else None
if not input_path.exists():
io.log_err("input_dir not found.")
return
if not output_file_path.parent.exists():
output_file_path.parent.mkdir(parents=True, exist_ok=True)
return
out_ext = output_file_path.suffix
if ext is None:
ext = io.input_str(
"Input image format (extension)? ( default:png ) : ", "png")
if lossless is None:
lossless = io.input_bool("Use lossless codec ? ( default:no ) : ",
False)
video_id = None
audio_id = None
ref_in_a = None
if reference_file_path is not None:
if reference_file_path.suffix == '.*':
reference_file_path = Path_utils.get_first_file_by_stem(
reference_file_path.parent, reference_file_path.stem)
else:
if not reference_file_path.exists():
reference_file_path = None
if reference_file_path is None:
io.log_err("reference_file not found.")
return
#probing reference file
probe = ffmpeg.probe(str(reference_file_path))
#getting first video and audio streams id with fps
for stream in probe['streams']:
if video_id is None and stream['codec_type'] == 'video':
video_id = stream['index']
fps = stream['r_frame_rate']
if audio_id is None and stream['codec_type'] == 'audio':
audio_id = stream['index']
if audio_id is not None:
#has audio track
ref_in_a = ffmpeg.input(str(reference_file_path))[str(audio_id)]
if fps is None:
#if fps not specified and not overwritten by reference-file
fps = max(1, io.input_int("FPS ? (default:25) : ", 25))
if not lossless and bitrate is None:
bitrate = max(
1,
io.input_int("Bitrate of output file in MB/s ? (default:16) : ",
16))
i_in = ffmpeg.input(str(input_path / ('%5d.' + ext)), r=fps)
output_args = [i_in]
if ref_in_a is not None:
output_args += [ref_in_a]
output_args += [str(output_file_path)]
output_kwargs = {}
if lossless:
output_kwargs.update({"c:v": "png"})
else:
output_kwargs.update({
"c:v": "libx264",
"b:v": "%dM" % (bitrate),
"pix_fmt": "yuv420p",
})
output_kwargs.update({
"c:a": "aac",
"b:a": "192k",
"ar": "48000",
"strict": "experimental"
})
job = (ffmpeg.output(*output_args, **output_kwargs).overwrite_output())
try:
job = job.run()
except:
io.log_err("ffmpeg fail, job commandline:" + str(job.compile()))
def main(input_dir,
output_dir,
debug_dir=None,
detector='mt',
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
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)
toscale = io.input_int(
"Output image Size (?:help skip:0 ) : ",
0,
help_message=
"Select extracted image size. A size of 0 will leave the extracted images unscaled"
)
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',
face_type,
debug_dir,
cpu_only=cpu_only,
manual=True,
size=toscale,
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,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
size=toscale).run()
io.log_info('Performing 2nd pass...')
data = ExtractSubprocessor(data,
'landmarks',
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
size=toscale).run()
io.log_info('Performing 3rd pass...')
data = ExtractSubprocessor(data,
'final',
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
final_output_path=output_path,
size=toscale).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',
face_type,
debug_dir,
manual=True,
manual_window_size=manual_window_size,
size=toscale).run()
fix_data = ExtractSubprocessor(fix_data,
'final',
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
final_output_path=output_path,
size=toscale).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 __init__(self, is_interactive, converter_session_filepath,
predictor_func, predictor_input_shape, converter_config,
frames, output_path, model_iter):
if len(frames) == 0:
raise ValueError("len (frames) == 0")
super().__init__('Converter',
ConvertSubprocessor.Cli,
86400 if CONVERTER_DEBUG else 60,
io_loop_sleep_time=0.001,
initialize_subprocesses_in_serial=False)
self.is_interactive = is_interactive
self.converter_session_filepath = Path(converter_session_filepath)
self.converter_config = converter_config
#dummy predict and sleep, tensorflow caching kernels. If remove it, sometime conversion speed can be x2 slower
predictor_func(dummy_predict=True)
time.sleep(2)
self.predictor_func_host, self.predictor_func = SubprocessFunctionCaller.make_pair(
predictor_func)
self.predictor_input_shape = predictor_input_shape
self.dcscn = None
self.ranksrgan = None
def superres_func(mode, *args, **kwargs):
if mode == 1:
if self.ranksrgan is None:
self.ranksrgan = imagelib.RankSRGAN()
return self.ranksrgan.upscale(*args, **kwargs)
self.dcscn_host, self.superres_func = SubprocessFunctionCaller.make_pair(
superres_func)
self.output_path = output_path
self.model_iter = model_iter
self.prefetch_frame_count = self.process_count = min(
6, multiprocessing.cpu_count())
session_data = None
if self.is_interactive and self.converter_session_filepath.exists():
if io.input_bool("Use saved session? (y/n skip:y) : ", True):
try:
with open(str(self.converter_session_filepath), "rb") as f:
session_data = pickle.loads(f.read())
except Exception as e:
pass
self.frames = frames
self.frames_idxs = [*range(len(self.frames))]
self.frames_done_idxs = []
if self.is_interactive and session_data is not None:
s_frames = session_data.get('frames', None)
s_frames_idxs = session_data.get('frames_idxs', None)
s_frames_done_idxs = session_data.get('frames_done_idxs', None)
s_model_iter = session_data.get('model_iter', None)
frames_equal = (s_frames is not None) and \
(s_frames_idxs is not None) and \
(s_frames_done_idxs is not None) and \
(s_model_iter is not None) and \
(len(frames) == len(s_frames))
if frames_equal:
for i in range(len(frames)):
frame = frames[i]
s_frame = s_frames[i]
if frame.frame_info.filename != s_frame.frame_info.filename:
frames_equal = False
if not frames_equal:
break
if frames_equal:
io.log_info(
'Using saved session from ' +
'/'.join(self.converter_session_filepath.parts[-2:]))
for frame in s_frames:
if frame.cfg is not None:
#recreate ConverterConfig class using constructor with get_config() as dict params
#so if any new param will be added, old converter session will work properly
frame.cfg = frame.cfg.__class__(
**frame.cfg.get_config())
self.frames = s_frames
self.frames_idxs = s_frames_idxs
self.frames_done_idxs = s_frames_done_idxs
if self.model_iter != s_model_iter:
#model is more trained, recompute all frames
for frame in self.frames:
frame.is_done = False
if self.model_iter != s_model_iter or \
len(self.frames_idxs) == 0:
#rewind to begin if model is more trained or all frames are done
while len(self.frames_done_idxs) > 0:
prev_frame = self.frames[self.frames_done_idxs.pop()]
self.frames_idxs.insert(0, prev_frame.idx)
if len(self.frames_idxs) != 0:
cur_frame = self.frames[self.frames_idxs[0]]
cur_frame.is_shown = False
if not frames_equal:
session_data = None
if session_data is None:
for filename in Path_utils.get_image_paths(
self.output_path): #remove all images in output_path
Path(filename).unlink()
frames[0].cfg = self.converter_config.copy()
for i in range(len(self.frames)):
frame = self.frames[i]
frame.idx = i
frame.output_filename = self.output_path / (
Path(frame.frame_info.filename).stem + '.png')
def __init__(self,
model_path,
training_data_src_path=None,
training_data_dst_path=None,
debug=False,
device_args=None):
device_args['force_gpu_idx'] = device_args.get('force_gpu_idx', -1)
if device_args['force_gpu_idx'] == -1:
idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList()
if len(idxs_names_list) > 1:
io.log_info("You have multi GPUs in a system: ")
for idx, name in idxs_names_list:
io.log_info("[%d] : %s" % (idx, name))
device_args['force_gpu_idx'] = io.input_int(
"Which GPU idx to choose? ( skip: best GPU ) : ", -1,
[x[0] for x in idxs_names_list])
self.device_args = device_args
io.log_info("Loading model...")
self.model_path = model_path
self.model_data_path = Path(
self.get_strpath_storage_for_file('data.dat'))
self.training_data_src_path = training_data_src_path
self.training_data_dst_path = training_data_dst_path
self.src_images_paths = None
self.dst_images_paths = None
self.src_yaw_images_paths = None
self.dst_yaw_images_paths = None
self.src_data_generator = None
self.dst_data_generator = None
self.debug = debug
self.is_training_mode = (training_data_src_path is not None
and training_data_dst_path is not None)
self.epoch = 0
self.options = {}
self.loss_history = []
self.sample_for_preview = None
if self.model_data_path.exists():
model_data = pickle.loads(self.model_data_path.read_bytes())
self.epoch = model_data['epoch']
if self.epoch != 0:
self.options = model_data['options']
self.loss_history = model_data[
'loss_history'] if 'loss_history' in model_data.keys(
) else []
self.sample_for_preview = model_data[
'sample_for_preview'] if 'sample_for_preview' in model_data.keys(
) else None
ask_override = self.is_training_mode and self.epoch != 0 and io.input_in_time(
"Press enter in 2 seconds to override model settings.", 2)
yn_str = {True: 'y', False: 'n'}
if self.epoch == 0:
io.log_info(
"\nModel first run. Enter model options as default for each run."
)
if self.epoch == 0 or ask_override:
default_write_preview_history = False if self.epoch == 0 else self.options.get(
'write_preview_history', False)
self.options['write_preview_history'] = io.input_bool(
"Write preview history? (y/n ?:help skip:%s) : " %
(yn_str[default_write_preview_history]),
default_write_preview_history,
help_message=
"Preview history will be writed to <ModelName>_history folder."
)
else:
self.options['write_preview_history'] = self.options.get(
'write_preview_history', False)
if self.epoch == 0 or ask_override:
self.options['target_epoch'] = max(
0, io.input_int("Target epoch (skip:unlimited/default) : ", 0))
else:
self.options['target_epoch'] = self.options.get('target_epoch', 0)
if self.epoch == 0 or ask_override:
default_batch_size = 0 if self.epoch == 0 else self.options.get(
'batch_size', 0)
self.options['batch_size'] = max(
0,
io.input_int(
"Batch_size (?:help skip:0/default) : ",
default_batch_size,
help_message=
"Larger batch size is always better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."
))
else:
self.options['batch_size'] = self.options.get('batch_size', 0)
if self.epoch == 0:
self.options['sort_by_yaw'] = io.input_bool(
"Feed faces to network sorted by yaw? (y/n ?:help skip:n) : ",
False,
help_message=
"NN will not learn src face directions that don't match dst face directions."
)
else:
self.options['sort_by_yaw'] = self.options.get(
'sort_by_yaw', False)
if self.epoch == 0:
self.options['random_flip'] = io.input_bool(
"Flip faces randomly? (y/n ?:help skip:y) : ",
True,
help_message=
"Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset."
)
else:
self.options['random_flip'] = self.options.get('random_flip', True)
if self.epoch == 0:
self.options['src_scale_mod'] = np.clip(
io.input_int(
"Src face scale modifier % ( -30...30, ?:help skip:0) : ",
0,
help_message=
"If src face shape is wider than dst, try to decrease this value to get a better result."
), -30, 30)
else:
self.options['src_scale_mod'] = self.options.get(
'src_scale_mod', 0)
self.write_preview_history = self.options['write_preview_history']
if not self.options['write_preview_history']:
self.options.pop('write_preview_history')
self.target_epoch = self.options['target_epoch']
if self.options['target_epoch'] == 0:
self.options.pop('target_epoch')
self.batch_size = self.options['batch_size']
self.sort_by_yaw = self.options['sort_by_yaw']
self.random_flip = self.options['random_flip']
self.src_scale_mod = self.options['src_scale_mod']
if self.src_scale_mod == 0:
self.options.pop('src_scale_mod')
self.onInitializeOptions(self.epoch == 0, ask_override)
nnlib.import_all(
nnlib.DeviceConfig(allow_growth=False, **self.device_args))
self.device_config = nnlib.active_DeviceConfig
self.keras = nnlib.keras
self.K = nnlib.keras.backend
self.onInitialize()
self.options['batch_size'] = self.batch_size
if self.debug or self.batch_size == 0:
self.batch_size = 1
if self.is_training_mode:
if self.write_preview_history:
if self.device_args['force_gpu_idx'] == -1:
self.preview_history_path = self.model_path / (
'%s_history' % (self.get_model_name()))
else:
self.preview_history_path = self.model_path / (
'%d_%s_history' % (self.device_args['force_gpu_idx'],
self.get_model_name()))
if not self.preview_history_path.exists():
self.preview_history_path.mkdir(exist_ok=True)
else:
if self.epoch == 0:
for filename in Path_utils.get_image_paths(
self.preview_history_path):
Path(filename).unlink()
if self.generator_list is None:
raise ValueError('You didnt set_training_data_generators()')
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase):
raise ValueError(
'training data generator is not subclass of SampleGeneratorBase'
)
if (self.sample_for_preview is None) or (self.epoch == 0):
self.sample_for_preview = self.generate_next_sample()
model_summary_text = []
model_summary_text += ["===== Model summary ====="]
model_summary_text += ["== Model name: " + self.get_model_name()]
model_summary_text += ["=="]
model_summary_text += ["== Current epoch: " + str(self.epoch)]
model_summary_text += ["=="]
model_summary_text += ["== Model options:"]
for key in self.options.keys():
model_summary_text += ["== |== %s : %s" % (key, self.options[key])]
if self.device_config.multi_gpu:
model_summary_text += ["== |== multi_gpu : True "]
model_summary_text += ["== Running on:"]
if self.device_config.cpu_only:
model_summary_text += ["== |== [CPU]"]
else:
for idx in self.device_config.gpu_idxs:
model_summary_text += [
"== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx))
]
if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[
0] == 2:
model_summary_text += ["=="]
model_summary_text += [
"== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."
]
model_summary_text += [
"== If training does not start, close all programs and try again."
]
model_summary_text += [
"== Also you can disable Windows Aero Desktop to get extra free VRAM."
]
model_summary_text += ["=="]
model_summary_text += ["========================="]
model_summary_text = "\r\n".join(model_summary_text)
self.model_summary_text = model_summary_text
io.log_info(model_summary_text)
def onInitializeOptions(self, is_first_run, ask_override):
yn_str = {True: 'y', False: 'n'}
default_resolution = 128
default_archi = 'df'
default_face_type = 'f'
if is_first_run:
resolution = io.input_int(
"Resolution ( 64-256 ?:help skip:128) : ",
default_resolution,
help_message=
"More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16."
)
resolution = np.clip(resolution, 64, 256)
while np.modf(resolution / 16)[0] != 0.0:
resolution -= 1
self.options['resolution'] = resolution
self.options['face_type'] = io.input_str(
"Half or Full face? (h/f, ?:help skip:f) : ",
default_face_type, ['h', 'f'],
help_message=
"Half face has better resolution, but covers less area of cheeks."
).lower()
self.options['learn_mask'] = io.input_bool(
"Learn mask? (y/n, ?:help skip:y) : ",
True,
help_message=
"Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result."
)
else:
self.options['resolution'] = self.options.get(
'resolution', default_resolution)
self.options['face_type'] = self.options.get(
'face_type', default_face_type)
self.options['learn_mask'] = self.options.get('learn_mask', True)
if (is_first_run or
ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int(
"Optimizer mode? ( 1,2,3 ?:help skip:%d) : " %
(def_optimizer_mode),
def_optimizer_mode,
help_message=
"1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power."
)
else:
self.options['optimizer_mode'] = self.options.get(
'optimizer_mode', 1)
if is_first_run:
self.options['archi'] = io.input_str(
"AE architecture (df, liae ?:help skip:%s) : " %
(default_archi),
default_archi, ['df', 'liae'],
help_message=
"'df' keeps faces more natural. 'liae' can fix overly different face shapes."
).lower(
) #-s version is slower, but has decreased change to collapse.
else:
self.options['archi'] = self.options.get('archi', default_archi)
default_ae_dims = 256 if 'liae' in self.options['archi'] else 512
default_e_ch_dims = 42
default_d_ch_dims = default_e_ch_dims // 2
def_ca_weights = False
if is_first_run:
self.options['ae_dims'] = np.clip(
io.input_int(
"AutoEncoder dims (32-1024 ?:help skip:%d) : " %
(default_ae_dims),
default_ae_dims,
help_message=
"All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU."
), 32, 1024)
self.options['e_ch_dims'] = np.clip(
io.input_int(
"Encoder dims per channel (21-85 ?:help skip:%d) : " %
(default_e_ch_dims),
default_e_ch_dims,
help_message=
"More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU."
), 21, 85)
default_d_ch_dims = self.options['e_ch_dims'] // 2
self.options['d_ch_dims'] = np.clip(
io.input_int(
"Decoder dims per channel (10-85 ?:help skip:%d) : " %
(default_d_ch_dims),
default_d_ch_dims,
help_message=
"More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU."
), 10, 85)
self.options['multiscale_decoder'] = io.input_bool(
"Use multiscale decoder? (y/n, ?:help skip:n) : ",
False,
help_message="Multiscale decoder helps to get better details.")
self.options['ca_weights'] = io.input_bool(
"Use CA weights? (y/n, ?:help skip: %s ) : " %
(yn_str[def_ca_weights]),
def_ca_weights,
help_message=
"Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run."
)
else:
self.options['ae_dims'] = self.options.get('ae_dims',
default_ae_dims)
self.options['e_ch_dims'] = self.options.get(
'e_ch_dims', default_e_ch_dims)
self.options['d_ch_dims'] = self.options.get(
'd_ch_dims', default_d_ch_dims)
self.options['multiscale_decoder'] = self.options.get(
'multiscale_decoder', False)
self.options['ca_weights'] = self.options.get(
'ca_weights', def_ca_weights)
default_face_style_power = 0.0
default_bg_style_power = 0.0
if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = io.input_bool(
"Use pixel loss? (y/n, ?:help skip: %s ) : " %
(yn_str[def_pixel_loss]),
def_pixel_loss,
help_message=
"Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. Enabling this option too early increases the chance of model collapse."
)
default_face_style_power = default_face_style_power if is_first_run else self.options.get(
'face_style_power', default_face_style_power)
self.options['face_style_power'] = np.clip(
io.input_number(
"Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " %
(default_face_style_power),
default_face_style_power,
help_message=
"Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
default_bg_style_power = default_bg_style_power if is_first_run else self.options.get(
'bg_style_power', default_bg_style_power)
self.options['bg_style_power'] = np.clip(
io.input_number(
"Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : "
% (default_bg_style_power),
default_bg_style_power,
help_message=
"Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
self.options['face_style_power'] = self.options.get(
'face_style_power', default_face_style_power)
self.options['bg_style_power'] = self.options.get(
'bg_style_power', default_bg_style_power)
if is_first_run:
self.options['pretrain'] = io.input_bool(
"Pretrain the model? (y/n, ?:help skip:n) : ",
False,
help_message=
"Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again."
)
else:
self.options['pretrain'] = False
def onInitializeOptions(self, is_first_run, ask_override):
yn_str = {True: 'y', False: 'n'}
default_resolution = 128
default_archi = 'df'
default_face_type = 'f'
if is_first_run:
resolution = io.input_int(
"Resolution ( 64-256 ?:help skip:128) : ",
default_resolution,
help_message=
"More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16."
)
resolution = np.clip(resolution, 64, 256)
while np.modf(resolution / 16)[0] != 0.0:
resolution -= 1
self.options['resolution'] = resolution
self.options['face_type'] = io.input_str(
"Half or Full face? (h/f, ?:help skip:f) : ",
default_face_type, ['h', 'f'],
help_message=
"Half face has better resolution, but covers less area of cheeks."
).lower()
self.options['learn_mask'] = io.input_bool(
"Learn mask? (y/n, ?:help skip:y) : ",
True,
help_message=
"Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result."
)
else:
self.options['resolution'] = self.options.get(
'resolution', default_resolution)
self.options['face_type'] = self.options.get(
'face_type', default_face_type)
self.options['learn_mask'] = self.options.get('learn_mask', True)
if (is_first_run or
ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int(
"Optimizer mode? ( 1,2,3 ?:help skip:%d) : " %
(def_optimizer_mode),
def_optimizer_mode,
help_message=
"1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power."
)
else:
self.options['optimizer_mode'] = self.options.get(
'optimizer_mode', 1)
if is_first_run:
self.options['archi'] = io.input_str(
"AE architecture (df, liae ?:help skip:%s) : " %
(default_archi),
default_archi, ['df', 'liae'],
help_message=
"'df' keeps faces more natural. 'liae' can fix overly different face shapes."
).lower()
else:
self.options['archi'] = self.options.get('archi', default_archi)
default_ae_dims = 256 if self.options['archi'] == 'liae' else 512
default_e_ch_dims = 42
default_d_ch_dims = default_e_ch_dims // 2
if is_first_run:
self.options['ae_dims'] = np.clip(
io.input_int(
"AutoEncoder dims (32-1024 ?:help skip:%d) : " %
(default_ae_dims),
default_ae_dims,
help_message=
"All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU."
), 32, 1024)
self.options['e_ch_dims'] = np.clip(
io.input_int(
"Encoder dims per channel (21-85 ?:help skip:%d) : " %
(default_e_ch_dims),
default_e_ch_dims,
help_message=
"More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU."
), 21, 85)
default_d_ch_dims = self.options['e_ch_dims'] // 2
self.options['d_ch_dims'] = np.clip(
io.input_int(
"Decoder dims per channel (10-85 ?:help skip:%d) : " %
(default_d_ch_dims),
default_d_ch_dims,
help_message=
"More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU."
), 10, 85)
self.options['d_residual_blocks'] = io.input_bool(
"Add residual blocks to decoder? (y/n, ?:help skip:n) : ",
False,
help_message=
"These blocks help to get better details, but require more computing time."
)
self.options['remove_gray_border'] = io.input_bool(
"Remove gray border? (y/n, ?:help skip:n) : ",
False,
help_message=
"Removes gray border of predicted face, but requires more computing resources."
)
else:
self.options['ae_dims'] = self.options.get('ae_dims',
default_ae_dims)
self.options['e_ch_dims'] = self.options.get(
'e_ch_dims', default_e_ch_dims)
self.options['d_ch_dims'] = self.options.get(
'd_ch_dims', default_d_ch_dims)
self.options['d_residual_blocks'] = self.options.get(
'd_residual_blocks', False)
self.options['remove_gray_border'] = self.options.get(
'remove_gray_border', False)
if is_first_run:
self.options['multiscale_decoder'] = io.input_bool(
"Use multiscale decoder? (y/n, ?:help skip:n) : ",
False,
help_message="Multiscale decoder helps to get better details.")
else:
self.options['multiscale_decoder'] = self.options.get(
'multiscale_decoder', False)
default_face_style_power = 0.0
default_bg_style_power = 0.0
if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = io.input_bool(
"Use pixel loss? (y/n, ?:help skip: %s ) : " %
(yn_str[def_pixel_loss]),
def_pixel_loss,
help_message=
"Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 15-25k iters to enhance fine details and decrease face jitter."
)
default_face_style_power = default_face_style_power if is_first_run else self.options.get(
'face_style_power', default_face_style_power)
self.options['face_style_power'] = np.clip(
io.input_number(
"Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " %
(default_face_style_power),
default_face_style_power,
help_message=
"Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes."
), 0.0, 100.0)
default_bg_style_power = default_bg_style_power if is_first_run else self.options.get(
'bg_style_power', default_bg_style_power)
self.options['bg_style_power'] = np.clip(
io.input_number(
"Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : "
% (default_bg_style_power),
default_bg_style_power,
help_message=
"Learn to transfer image around face. This can make face more like dst."
), 0.0, 100.0)
else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
self.options['face_style_power'] = self.options.get(
'face_style_power', default_face_style_power)
self.options['bg_style_power'] = self.options.get(
'bg_style_power', default_bg_style_power)
def __init__(self,
model_path,
training_data_src_path=None,
training_data_dst_path=None,
pretraining_data_path=None,
debug=False,
device_args=None,
ask_enable_autobackup=True,
ask_write_preview_history=True,
ask_target_iter=True,
ask_batch_size=True,
ask_sort_by_yaw=True,
ask_random_flip=True,
ask_src_scale_mod=True):
device_args['force_gpu_idx'] = device_args.get('force_gpu_idx', -1)
device_args['cpu_only'] = device_args.get('cpu_only', False)
if device_args['force_gpu_idx'] == -1 and not device_args['cpu_only']:
idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList()
if len(idxs_names_list) > 1:
io.log_info("You have multi GPUs in a system: ")
for idx, name in idxs_names_list:
io.log_info("[%d] : %s" % (idx, name))
device_args['force_gpu_idx'] = io.input_int(
"Which GPU idx to choose? ( skip: best GPU ) : ", -1,
[x[0] for x in idxs_names_list])
self.device_args = device_args
self.device_config = nnlib.DeviceConfig(allow_growth=False,
**self.device_args)
io.log_info("Loading model...")
self.model_path = model_path
self.model_data_path = Path(
self.get_strpath_storage_for_file('data.dat'))
self.training_data_src_path = training_data_src_path
self.training_data_dst_path = training_data_dst_path
self.pretraining_data_path = pretraining_data_path
self.src_images_paths = None
self.dst_images_paths = None
self.src_yaw_images_paths = None
self.dst_yaw_images_paths = None
self.src_data_generator = None
self.dst_data_generator = None
self.debug = debug
self.is_training_mode = (training_data_src_path is not None
and training_data_dst_path is not None)
self.iter = 0
self.options = {}
self.loss_history = []
self.sample_for_preview = None
model_data = {}
if self.model_data_path.exists():
model_data = pickle.loads(self.model_data_path.read_bytes())
self.iter = max(model_data.get('iter', 0),
model_data.get('epoch', 0))
if 'epoch' in self.options:
self.options.pop('epoch')
if self.iter != 0:
self.options = model_data['options']
self.loss_history = model_data.get('loss_history', [])
self.sample_for_preview = model_data.get(
'sample_for_preview', None)
ask_override = self.is_training_mode and self.iter != 0 and io.input_in_time(
"Press enter in 2 seconds to override model settings.",
5 if io.is_colab() else 2)
yn_str = {True: 'y', False: 'n'}
if self.iter == 0:
io.log_info(
"\nModel first run. Enter model options as default for each run."
)
if ask_enable_autobackup and (self.iter == 0 or ask_override):
default_autobackup = False if self.iter == 0 else self.options.get(
'autobackup', False)
self.options['autobackup'] = io.input_bool(
"Enable autobackup? (y/n ?:help skip:%s) : " %
(yn_str[default_autobackup]),
default_autobackup,
help_message=
"Autobackup model files with preview every hour for last 15 hours. Latest backup located in model/<>_autobackups/01"
)
else:
self.options['autobackup'] = self.options.get('autobackup', False)
if ask_write_preview_history and (self.iter == 0 or ask_override):
default_write_preview_history = False if self.iter == 0 else self.options.get(
'write_preview_history', False)
self.options['write_preview_history'] = io.input_bool(
"Write preview history? (y/n ?:help skip:%s) : " %
(yn_str[default_write_preview_history]),
default_write_preview_history,
help_message=
"Preview history will be writed to <ModelName>_history folder."
)
else:
self.options['write_preview_history'] = self.options.get(
'write_preview_history', False)
if (self.iter == 0 or ask_override) and self.options[
'write_preview_history'] and io.is_support_windows():
choose_preview_history = io.input_bool(
"Choose image for the preview history? (y/n skip:%s) : " %
(yn_str[False]), False)
else:
choose_preview_history = False
if ask_target_iter:
if (self.iter == 0 or ask_override):
self.options['target_iter'] = max(
0,
io.input_int(
"Target iteration (skip:unlimited/default) : ", 0))
else:
self.options['target_iter'] = max(
model_data.get('target_iter', 0),
self.options.get('target_epoch', 0))
if 'target_epoch' in self.options:
self.options.pop('target_epoch')
if ask_batch_size and (self.iter == 0 or ask_override):
default_batch_size = 0 if self.iter == 0 else self.options.get(
'batch_size', 0)
self.options['batch_size'] = max(
0,
io.input_int(
"Batch_size (?:help skip:%d) : " % (default_batch_size),
default_batch_size,
help_message=
"Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."
))
else:
self.options['batch_size'] = self.options.get('batch_size', 0)
if ask_sort_by_yaw:
if (self.iter == 0 or ask_override):
default_sort_by_yaw = self.options.get('sort_by_yaw', False)
self.options['sort_by_yaw'] = io.input_bool(
"Feed faces to network sorted by yaw? (y/n ?:help skip:%s) : "
% (yn_str[default_sort_by_yaw]),
default_sort_by_yaw,
help_message=
"NN will not learn src face directions that don't match dst face directions. Do not use if the dst face has hair that covers the jaw."
)
else:
self.options['sort_by_yaw'] = self.options.get(
'sort_by_yaw', False)
if ask_random_flip:
if (self.iter == 0):
self.options['random_flip'] = io.input_bool(
"Flip faces randomly? (y/n ?:help skip:y) : ",
True,
help_message=
"Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset."
)
else:
self.options['random_flip'] = self.options.get(
'random_flip', True)
if ask_src_scale_mod:
if (self.iter == 0):
self.options['src_scale_mod'] = np.clip(
io.input_int(
"Src face scale modifier % ( -30...30, ?:help skip:0) : ",
0,
help_message=
"If src face shape is wider than dst, try to decrease this value to get a better result."
), -30, 30)
else:
self.options['src_scale_mod'] = self.options.get(
'src_scale_mod', 0)
self.autobackup = self.options.get('autobackup', False)
if not self.autobackup and 'autobackup' in self.options:
self.options.pop('autobackup')
self.write_preview_history = self.options.get('write_preview_history',
False)
if not self.write_preview_history and 'write_preview_history' in self.options:
self.options.pop('write_preview_history')
self.target_iter = self.options.get('target_iter', 0)
if self.target_iter == 0 and 'target_iter' in self.options:
self.options.pop('target_iter')
self.batch_size = self.options.get('batch_size', 0)
self.sort_by_yaw = self.options.get('sort_by_yaw', False)
self.random_flip = self.options.get('random_flip', True)
self.src_scale_mod = self.options.get('src_scale_mod', 0)
if self.src_scale_mod == 0 and 'src_scale_mod' in self.options:
self.options.pop('src_scale_mod')
self.onInitializeOptions(self.iter == 0, ask_override)
nnlib.import_all(self.device_config)
self.keras = nnlib.keras
self.K = nnlib.keras.backend
self.onInitialize()
self.options['batch_size'] = self.batch_size
if self.debug or self.batch_size == 0:
self.batch_size = 1
if self.is_training_mode:
if self.device_args['force_gpu_idx'] == -1:
self.preview_history_path = self.model_path / (
'%s_history' % (self.get_model_name()))
self.autobackups_path = self.model_path / (
'%s_autobackups' % (self.get_model_name()))
else:
self.preview_history_path = self.model_path / (
'%d_%s_history' %
(self.device_args['force_gpu_idx'], self.get_model_name()))
self.autobackups_path = self.model_path / (
'%d_%s_autobackups' %
(self.device_args['force_gpu_idx'], self.get_model_name()))
if self.autobackup:
self.autobackup_current_hour = time.localtime().tm_hour
if not self.autobackups_path.exists():
self.autobackups_path.mkdir(exist_ok=True)
if self.write_preview_history or io.is_colab():
if not self.preview_history_path.exists():
self.preview_history_path.mkdir(exist_ok=True)
else:
if self.iter == 0:
for filename in Path_utils.get_image_paths(
self.preview_history_path):
Path(filename).unlink()
if self.generator_list is None:
raise ValueError('You didnt set_training_data_generators()')
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase):
raise ValueError(
'training data generator is not subclass of SampleGeneratorBase'
)
if self.sample_for_preview is None or choose_preview_history:
if choose_preview_history and io.is_support_windows():
wnd_name = "[p] - next. [enter] - confirm."
io.named_window(wnd_name)
io.capture_keys(wnd_name)
choosed = False
while not choosed:
self.sample_for_preview = self.generate_next_sample()
preview = self.get_static_preview()
io.show_image(wnd_name,
(preview * 255).astype(np.uint8))
while True:
key_events = io.get_key_events(wnd_name)
key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[
-1] if len(key_events) > 0 else (0, 0, False,
False, False)
if key == ord('\n') or key == ord('\r'):
choosed = True
break
elif key == ord('p'):
break
try:
io.process_messages(0.1)
except KeyboardInterrupt:
choosed = True
io.destroy_window(wnd_name)
else:
self.sample_for_preview = self.generate_next_sample()
self.last_sample = self.sample_for_preview
model_summary_text = []
model_summary_text += ["===== Model summary ====="]
model_summary_text += ["== Model name: " + self.get_model_name()]
model_summary_text += ["=="]
model_summary_text += ["== Current iteration: " + str(self.iter)]
model_summary_text += ["=="]
model_summary_text += ["== Model options:"]
for key in self.options.keys():
model_summary_text += ["== |== %s : %s" % (key, self.options[key])]
if self.device_config.multi_gpu:
model_summary_text += ["== |== multi_gpu : True "]
model_summary_text += ["== Running on:"]
if self.device_config.cpu_only:
model_summary_text += ["== |== [CPU]"]
else:
for idx in self.device_config.gpu_idxs:
model_summary_text += [
"== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx))
]
if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[
0] == 2:
model_summary_text += ["=="]
model_summary_text += [
"== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."
]
model_summary_text += [
"== If training does not start, close all programs and try again."
]
model_summary_text += [
"== Also you can disable Windows Aero Desktop to get extra free VRAM."
]
model_summary_text += ["=="]
model_summary_text += ["========================="]
model_summary_text = "\r\n".join(model_summary_text)
self.model_summary_text = model_summary_text
io.log_info(model_summary_text)
def onInitializeOptions(self, is_first_run, ask_override):
yn_str = {True: 'y', False: 'n'}
default_resolution = 128
default_archi = 'df'
default_face_type = 'f'
if is_first_run:
resolution = io.input_int(
"Resolution ( 64-256 ?:help skip:128) : ",
default_resolution,
help_message=
"More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16."
)
resolution = np.clip(resolution, 64, 256)
while np.modf(resolution / 16)[0] != 0.0:
resolution -= 1
self.options['resolution'] = resolution
self.options['face_type'] = io.input_str(
"Half, mid full, or full face? (h/mf/f, ?:help skip:f) : ",
default_face_type, ['h', 'mf', 'f'],
help_message=
"Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face."
).lower()
else:
self.options['resolution'] = self.options.get(
'resolution', default_resolution)
self.options['face_type'] = self.options.get(
'face_type', default_face_type)
default_learn_mask = self.options.get('learn_mask', True)
if is_first_run or ask_override:
self.options['learn_mask'] = io.input_bool(
f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : ",
default_learn_mask,
help_message=
"Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted."
)
else:
self.options['learn_mask'] = self.options.get(
'learn_mask', default_learn_mask)
if (is_first_run or
ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int(
"Optimizer mode? ( 1,2,3 ?:help skip:%d) : " %
(def_optimizer_mode),
def_optimizer_mode,
help_message=
"1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power."
)
else:
self.options['optimizer_mode'] = self.options.get(
'optimizer_mode', 1)
if is_first_run:
self.options['archi'] = io.input_str(
"AE architecture (df, liae ?:help skip:%s) : " %
(default_archi),
default_archi, ['df', 'liae'],
help_message=
"'df' keeps faces more natural. 'liae' can fix overly different face shapes."
).lower(
) #-s version is slower, but has decreased change to collapse.
else:
self.options['archi'] = self.options.get('archi', default_archi)
default_ae_dims = 256
default_ed_ch_dims = 21
if is_first_run:
self.options['ae_dims'] = np.clip(
io.input_int(
"AutoEncoder dims (32-1024 ?:help skip:%d) : " %
(default_ae_dims),
default_ae_dims,
help_message=
"All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU."
), 32, 1024)
self.options['ed_ch_dims'] = np.clip(
io.input_int(
"Encoder/Decoder dims per channel (10-85 ?:help skip:%d) : "
% (default_ed_ch_dims),
default_ed_ch_dims,
help_message=
"More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU."
), 10, 85)
else:
self.options['ae_dims'] = self.options.get('ae_dims',
default_ae_dims)
self.options['ed_ch_dims'] = self.options.get(
'ed_ch_dims', default_ed_ch_dims)
default_true_face_training = self.options.get('true_face_training',
False)
default_face_style_power = self.options.get('face_style_power', 0.0)
default_bg_style_power = self.options.get('bg_style_power', 0.0)
if is_first_run or ask_override:
default_random_warp = self.options.get('random_warp', True)
self.options['random_warp'] = io.input_bool(
f"Enable random warp of samples? ( y/n, ?:help skip:{yn_str[default_random_warp]}) : ",
default_random_warp,
help_message=
"Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness for less amount of iterations."
)
self.options['true_face_training'] = io.input_bool(
f"Enable 'true face' training? (y/n, ?:help skip:{yn_str[default_true_face_training]}) : ",
default_true_face_training,
help_message=
"The result face will be more like src and will get extra sharpness. Enable it for last 10-20k iterations before conversion."
)
self.options['face_style_power'] = np.clip(
io.input_number(
"Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " %
(default_face_style_power),
default_face_style_power,
help_message=
"Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
self.options['bg_style_power'] = np.clip(
io.input_number(
"Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : "
% (default_bg_style_power),
default_bg_style_power,
help_message=
"Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."
), 0.0, 100.0)
default_ct_mode = self.options.get('ct_mode', 'none')
self.options['ct_mode'] = io.input_str(
f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt, ?:help skip:{default_ct_mode}) : ",
default_ct_mode, ['none', 'rct', 'lct', 'mkl', 'idt'],
help_message=
"Change color distribution of src samples close to dst samples. Try all modes to find the best."
)
if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301
default_clipgrad = False if is_first_run else self.options.get(
'clipgrad', False)
self.options['clipgrad'] = io.input_bool(
f"Enable gradient clipping? (y/n, ?:help skip:{yn_str[default_clipgrad]}) : ",
default_clipgrad,
help_message=
"Gradient clipping reduces chance of model collapse, sacrificing speed of training."
)
else:
self.options['clipgrad'] = False
else:
self.options['random_warp'] = self.options.get('random_warp', True)
self.options['true_face_training'] = self.options.get(
'true_face_training', default_true_face_training)
self.options['face_style_power'] = self.options.get(
'face_style_power', default_face_style_power)
self.options['bg_style_power'] = self.options.get(
'bg_style_power', default_bg_style_power)
self.options['ct_mode'] = self.options.get('ct_mode', 'none')
self.options['clipgrad'] = self.options.get('clipgrad', False)
if is_first_run:
self.options['pretrain'] = io.input_bool(
"Pretrain the model? (y/n, ?:help skip:n) : ",
False,
help_message=
"Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again."
)
else:
self.options['pretrain'] = False
def __init__(
self,
model_path,
training_data_src_path=None,
training_data_dst_path=None,
pretraining_data_path=None,
debug=False,
device_args=None,
ask_enable_autobackup=True,
ask_write_preview_history=True,
ask_target_iter=True,
ask_batch_size=True,
ask_sort_by_yaw=True,
ask_random_flip=True,
ask_src_scale_mod=True,
):
device_args["force_gpu_idx"] = device_args.get("force_gpu_idx", -1)
device_args["cpu_only"] = device_args.get("cpu_only", False)
if device_args["force_gpu_idx"] == -1 and not device_args["cpu_only"]:
idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList()
if len(idxs_names_list) > 1:
io.log_info("You have multi GPUs in a system: ")
for idx, name in idxs_names_list:
io.log_info("[%d] : %s" % (idx, name))
device_args["force_gpu_idx"] = io.input_int(
"Which GPU idx to choose? ( skip: best GPU ) : ",
-1,
[x[0] for x in idxs_names_list],
)
self.device_args = device_args
self.device_config = nnlib.DeviceConfig(allow_growth=True,
**self.device_args)
io.log_info("Loading model...")
self.model_path = model_path
self.model_data_path = Path(
self.get_strpath_storage_for_file("data.dat"))
self.training_data_src_path = training_data_src_path
self.training_data_dst_path = training_data_dst_path
self.pretraining_data_path = pretraining_data_path
self.src_images_paths = None
self.dst_images_paths = None
self.src_yaw_images_paths = None
self.dst_yaw_images_paths = None
self.src_data_generator = None
self.dst_data_generator = None
self.debug = debug
self.is_training_mode = (training_data_src_path is not None
and training_data_dst_path is not None)
self.iter = 0
self.options = {}
self.loss_history = []
self.sample_for_preview = None
model_data = {}
if self.model_data_path.exists():
model_data = pickle.loads(self.model_data_path.read_bytes())
self.iter = max(model_data.get("iter", 0),
model_data.get("epoch", 0))
if "epoch" in self.options:
self.options.pop("epoch")
if self.iter != 0:
self.options = model_data["options"]
self.loss_history = model_data.get("loss_history", [])
self.sample_for_preview = model_data.get(
"sample_for_preview", None)
ask_override = (
self.is_training_mode and self.iter != 0 and io.input_in_time(
"Press enter in 2 seconds to override model settings.",
5 if io.is_colab() else 2,
))
yn_str = {True: "y", False: "n"}
if self.iter == 0:
io.log_info(
"\nModel first run. Enter model options as default for each run."
)
if ask_enable_autobackup and (self.iter == 0 or ask_override):
default_autobackup = (False if self.iter == 0 else
self.options.get("autobackup", False))
self.options["autobackup"] = io.input_bool(
"Enable autobackup? (y/n ?:help skip:%s) : " %
(yn_str[default_autobackup]),
default_autobackup,
help_message=
"Autobackup model files with preview every hour for last 15 hours. Latest backup located in model/<>_autobackups/01",
)
else:
self.options["autobackup"] = self.options.get("autobackup", False)
if ask_write_preview_history and (self.iter == 0 or ask_override):
default_write_preview_history = (
False if self.iter == 0 else self.options.get(
"write_preview_history", False))
self.options["write_preview_history"] = io.input_bool(
"Write preview history? (y/n ?:help skip:%s) : " %
(yn_str[default_write_preview_history]),
default_write_preview_history,
help_message=
"Preview history will be writed to <ModelName>_history folder.",
)
else:
self.options["write_preview_history"] = self.options.get(
"write_preview_history", False)
if ((self.iter == 0 or ask_override)
and self.options["write_preview_history"]
and io.is_support_windows()):
choose_preview_history = io.input_bool(
"Choose image for the preview history? (y/n skip:%s) : " %
(yn_str[False]),
False,
)
elif ((self.iter == 0 or ask_override)
and self.options["write_preview_history"] and io.is_colab()):
choose_preview_history = io.input_bool(
"Randomly choose new image for preview history? (y/n ?:help skip:%s) : "
% (yn_str[False]),
False,
help_message=
"Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person",
)
else:
choose_preview_history = False
if ask_target_iter:
if self.iter == 0 or ask_override:
try:
iterations = int(os.getenv("SP_FaceLab_Iterations"))
except:
iterations = 0
self.options["target_iter"] = max(
0,
io.input_int(
"Target iteration (skip:unlimited/default) : ",
iterations),
)
else:
try:
iterations = int(os.getenv("SP_FaceLab_Iterations"))
except:
iterations = 0
self.options["target_iter"] = max(
model_data.get("target_iter", 0),
self.options.get("target_epoch", 0),
iterations,
)
if "target_epoch" in self.options:
self.options.pop("target_epoch")
if ask_batch_size and (self.iter == 0 or ask_override):
default_batch_size = (0 if self.iter == 0 else self.options.get(
"batch_size", 0))
try:
default_batch_size = int(os.getenv("SP_FaceLab_Batch_Size"))
except:
default_batch_size = (0 if self.iter == 0 else
self.options.get("batch_size", 0))
self.batch_size = max(
0,
io.input_int(
"Batch_size (?:help skip:%d) : " % (default_batch_size),
default_batch_size,
help_message=
"Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually.",
),
)
else:
try:
default_batch_size = int(os.getenv("SP_FaceLab_Batch_Size"))
except:
default_batch_size = self.options.get("batch_size", 0)
self.batch_size = default_batch_size
if ask_sort_by_yaw:
if self.iter == 0 or ask_override:
default_sort_by_yaw = self.options.get("sort_by_yaw", False)
self.options["sort_by_yaw"] = io.input_bool(
"Feed faces to network sorted by yaw? (y/n ?:help skip:%s) : "
% (yn_str[default_sort_by_yaw]),
default_sort_by_yaw,
help_message=
"NN will not learn src face directions that don't match dst face directions. Do not use if the dst face has hair that covers the jaw.",
)
else:
self.options["sort_by_yaw"] = self.options.get(
"sort_by_yaw", False)
if ask_random_flip:
if self.iter == 0:
self.options["random_flip"] = io.input_bool(
"Flip faces randomly? (y/n ?:help skip:y) : ",
True,
help_message=
"Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.",
)
else:
self.options["random_flip"] = self.options.get(
"random_flip", True)
if ask_src_scale_mod:
if self.iter == 0:
self.options["src_scale_mod"] = np.clip(
io.input_int(
"Src face scale modifier % ( -30...30, ?:help skip:0) : ",
0,
help_message=
"If src face shape is wider than dst, try to decrease this value to get a better result.",
),
-30,
30,
)
else:
self.options["src_scale_mod"] = self.options.get(
"src_scale_mod", 0)
self.autobackup = self.options.get("autobackup", False)
if not self.autobackup and "autobackup" in self.options:
self.options.pop("autobackup")
self.write_preview_history = self.options.get("write_preview_history",
False)
if not self.write_preview_history and "write_preview_history" in self.options:
self.options.pop("write_preview_history")
self.target_iter = self.options.get("target_iter", 0)
if self.target_iter == 0 and "target_iter" in self.options:
self.options.pop("target_iter")
# self.batch_size = self.options.get('batch_size',0)
self.sort_by_yaw = self.options.get("sort_by_yaw", False)
self.random_flip = self.options.get("random_flip", True)
self.src_scale_mod = self.options.get("src_scale_mod", 0)
if self.src_scale_mod == 0 and "src_scale_mod" in self.options:
self.options.pop("src_scale_mod")
self.onInitializeOptions(self.iter == 0, ask_override)
nnlib.import_all(self.device_config)
self.keras = nnlib.keras
self.K = nnlib.keras.backend
self.onInitialize()
self.options["batch_size"] = self.batch_size
if self.debug or self.batch_size == 0:
self.batch_size = 1
if self.is_training_mode:
if self.device_args["force_gpu_idx"] == -1:
self.preview_history_path = self.model_path / (
"%s_history" % (self.get_model_name()))
self.autobackups_path = self.model_path / (
"%s_autobackups" % (self.get_model_name()))
else:
self.preview_history_path = self.model_path / (
"%d_%s_history" %
(self.device_args["force_gpu_idx"], self.get_model_name()))
self.autobackups_path = self.model_path / (
"%d_%s_autobackups" %
(self.device_args["force_gpu_idx"], self.get_model_name()))
if self.autobackup:
self.autobackup_current_hour = time.localtime().tm_hour
if not self.autobackups_path.exists():
self.autobackups_path.mkdir(exist_ok=True)
if self.write_preview_history or io.is_colab():
if not self.preview_history_path.exists():
self.preview_history_path.mkdir(exist_ok=True)
else:
if self.iter == 0:
for filename in Path_utils.get_image_paths(
self.preview_history_path):
Path(filename).unlink()
if self.generator_list is None:
raise ValueError("You didnt set_training_data_generators()")
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase):
raise ValueError(
"training data generator is not subclass of SampleGeneratorBase"
)
if self.sample_for_preview is None or choose_preview_history:
if choose_preview_history and io.is_support_windows():
wnd_name = "[p] - next. [enter] - confirm."
io.named_window(wnd_name)
io.capture_keys(wnd_name)
choosed = False
while not choosed:
self.sample_for_preview = self.generate_next_sample()
preview = self.get_static_preview()
io.show_image(wnd_name,
(preview * 255).astype(np.uint8))
while True:
key_events = io.get_key_events(wnd_name)
key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = (
key_events[-1] if len(key_events) > 0 else
(0, 0, False, False, False))
if key == ord("\n") or key == ord("\r"):
choosed = True
break
elif key == ord("p"):
break
try:
io.process_messages(0.1)
except KeyboardInterrupt:
choosed = True
io.destroy_window(wnd_name)
else:
self.sample_for_preview = self.generate_next_sample()
self.last_sample = self.sample_for_preview
###Generate text summary of model hyperparameters
# Find the longest key name and value string. Used as column widths.
width_name = (
max([len(k) for k in self.options.keys()] + [17]) + 1
) # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration"
width_value = (max([len(str(x)) for x in self.options.values()] +
[len(str(self.iter)),
len(self.get_model_name())]) + 1
) # Single space buffer to right edge
if not self.device_config.cpu_only: # Check length of GPU names
width_value = max([
len(nnlib.device.getDeviceName(idx)) + 1
for idx in self.device_config.gpu_idxs
] + [width_value])
width_total = width_name + width_value + 2 # Plus 2 for ": "
model_summary_text = []
model_summary_text += [f'=={" Model Summary ":=^{width_total}}=='
] # Model/status summary
model_summary_text += [f'=={" "*width_total}==']
model_summary_text += [
f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='
] # Name
model_summary_text += [f'=={" "*width_total}==']
model_summary_text += [
f'=={"Current iteration": >{width_name}}: {str(self.iter): <{width_value}}=='
] # Iter
model_summary_text += [f'=={" "*width_total}==']
model_summary_text += [f'=={" Model Options ":-^{width_total}}=='
] # Model options
model_summary_text += [f'=={" "*width_total}==']
for key in self.options.keys():
model_summary_text += [
f"=={key: >{width_name}}: {str(self.options[key]): <{width_value}}=="
] # self.options key/value pairs
model_summary_text += [f'=={" "*width_total}==']
model_summary_text += [f'=={" Running On ":-^{width_total}}=='
] # Training hardware info
model_summary_text += [f'=={" "*width_total}==']
if self.device_config.multi_gpu:
model_summary_text += [
f'=={"Using multi_gpu": >{width_name}}: {"True": <{width_value}}=='
] # multi_gpu
model_summary_text += [f'=={" "*width_total}==']
if self.device_config.cpu_only:
model_summary_text += [
f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='
] # cpu_only
else:
for idx in self.device_config.gpu_idxs:
model_summary_text += [
f'=={"Device index": >{width_name}}: {idx: <{width_value}}=='
] # GPU hardware device index
model_summary_text += [
f'=={"Name": >{width_name}}: {nnlib.device.getDeviceName(idx): <{width_value}}=='
] # GPU name
vram_str = (f"{nnlib.device.getDeviceVRAMTotalGb(idx):.2f}GB"
) # GPU VRAM - Formated as #.## (or ##.##)
model_summary_text += [
f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}=='
]
model_summary_text += [f'=={" "*width_total}==']
model_summary_text += [f'=={"="*width_total}==']
if (not self.device_config.cpu_only and
self.device_config.gpu_vram_gb[0] <= 2): # Low VRAM warning
model_summary_text += ["/!\\"]
model_summary_text += ["/!\\ WARNING:"]
model_summary_text += [
"/!\\ You are using a GPU with 2GB or less VRAM. This may significantly reduce the quality of your result!"
]
model_summary_text += [
"/!\\ If training does not start, close all programs and try again."
]
model_summary_text += [
"/!\\ Also you can disable Windows Aero Desktop to increase available VRAM."
]
model_summary_text += ["/!\\"]
model_summary_text = "\n".join(model_summary_text)
self.model_summary_text = model_summary_text
io.log_info(model_summary_text)
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("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,
).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 relight(input_dir, lighten=None, random_one=None):
if lighten is None:
lighten = io.input_bool ("Lighten the faces? ( y/n default:n ?:help ) : ", False, help_message="Lighten the faces instead of shadow. May produce artifacts." )
if io.is_colab():
io.log_info("In colab version you cannot choose light directions manually.")
manual = False
else:
manual = io.input_bool ("Choose light directions manually? ( y/n default:y ) : ", True)
if not manual:
if random_one is None:
random_one = io.input_bool ("Relight the faces only with one random direction? ( y/n default:y ?:help) : ", True, help_message="Otherwise faceset will be relighted with predefined 7 light directions.")
image_paths = [Path(x) for x in Path_utils.get_image_paths(input_dir)]
filtered_image_paths = []
for filepath in io.progress_bar_generator(image_paths, "Collecting fileinfo"):
try:
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath) )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath) )
else:
dflimg = None
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
else:
if not dflimg.get_relighted():
filtered_image_paths += [filepath]
except:
io.log_err (f"Exception occured while processing file {filepath.name}. Error: {traceback.format_exc()}")
image_paths = filtered_image_paths
if len(image_paths) == 0:
io.log_info("No files to process.")
return
dpr = DeepPortraitRelighting()
if manual:
alt_azi_ar = RelightEditor(image_paths, dpr, lighten).run()
else:
if not random_one:
alt_azi_ar = [(60,0), (60,60), (0,60), (-60,60), (-60,0), (-60,-60), (0,-60), (60,-60)]
for filepath in io.progress_bar_generator(image_paths, "Relighting"):
try:
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath) )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath) )
else:
dflimg = None
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
else:
if dflimg.get_relighted():
continue
img = cv2_imread (str(filepath))
if random_one:
alt = np.random.randint(-90,91)
azi = np.random.randint(-90,91)
relighted_imgs = [dpr.relight(img,alt=alt,azi=azi,lighten=lighten)]
else:
relighted_imgs = [dpr.relight(img,alt=alt,azi=azi,lighten=lighten) for (alt,azi) in alt_azi_ar ]
i = 0
for i,relighted_img in enumerate(relighted_imgs):
im_flags = []
if filepath.suffix == '.jpg':
im_flags += [int(cv2.IMWRITE_JPEG_QUALITY), 100]
while True:
relighted_filepath = filepath.parent / (filepath.stem+f'_relighted_{i}'+filepath.suffix)
if not relighted_filepath.exists():
break
i += 1
cv2_imwrite (relighted_filepath, relighted_img )
dflimg.embed_and_set (relighted_filepath, source_filename="_", relighted=True )
except:
io.log_err (f"Exception occured while processing file {filepath.name}. Error: {traceback.format_exc()}")