def get_person_id_max_count(samples_path):
return len ( Path_utils.get_all_dir_names(samples_path) )
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 mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
input_path = Path(input_dir)
confirmed_path = Path(confirmed_dir)
skipped_path = Path(skipped_dir)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
if not confirmed_path.exists():
confirmed_path.mkdir(parents=True)
if not skipped_path.exists():
skipped_path.mkdir(parents=True)
wnd_name = "MaskEditor tool"
io.named_window(wnd_name)
io.capture_mouse(wnd_name)
io.capture_keys(wnd_name)
cached_images = {}
image_paths = [Path(x) for x in Path_utils.get_image_paths(input_path)]
done_paths = []
done_images_types = {}
image_paths_total = len(image_paths)
zoom_factor = 1.0
preview_images_count = 9
target_wh = 256
do_prev_count = 0
do_save_move_count = 0
do_save_count = 0
do_skip_move_count = 0
do_skip_count = 0
def jobs_count():
return do_prev_count + do_save_move_count + do_save_count + do_skip_move_count + do_skip_count
is_exit = False
while not is_exit:
if len(image_paths) > 0:
filepath = image_paths.pop(0)
else:
filepath = None
next_image_paths = image_paths[0:preview_images_count]
next_image_paths_names = [path.name for path in next_image_paths]
prev_image_paths = done_paths[-preview_images_count:]
prev_image_paths_names = [path.name for path in prev_image_paths]
for key in list(cached_images.keys()):
if key not in prev_image_paths_names and \
key not in next_image_paths_names:
cached_images.pop(key)
for paths in [prev_image_paths, next_image_paths]:
for path in paths:
if path.name not in cached_images:
cached_images[path.name] = cv2_imread(str(path)) / 255.0
if filepath is not None:
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:
lmrks = dflimg.get_landmarks()
ie_polys = dflimg.get_ie_polys()
if filepath.name in cached_images:
img = cached_images[filepath.name]
else:
img = cached_images[filepath.name] = cv2_imread(
str(filepath)) / 255.0
mask = LandmarksProcessor.get_image_hull_mask(img.shape, lmrks)
else:
img = np.zeros((target_wh, target_wh, 3))
mask = np.ones((target_wh, target_wh, 3))
ie_polys = None
def get_status_lines_func():
return [
'Progress: %d / %d . Current file: %s' %
(len(done_paths), image_paths_total,
str(filepath.name) if filepath is not None else "end"),
'[Left mouse button] - mark include mask.',
'[Right mouse button] - mark exclude mask.',
'[Middle mouse button] - finish current poly.',
'[Mouse wheel] - undo/redo poly or point. [+ctrl] - undo to begin/redo to end',
'[q] - prev image. [w] - skip and move to %s. [e] - save and move to %s. '
% (skipped_path.name, confirmed_path.name),
'[z] - prev image. [x] - skip. [c] - save. ',
'hold [shift] - speed up the frame counter by 10.',
'[-/+] - window zoom [esc] - quit',
]
try:
ed = MaskEditor(img,
[(done_images_types[name], cached_images[name])
for name in prev_image_paths_names],
[(0, cached_images[name])
for name in next_image_paths_names], mask,
ie_polys, get_status_lines_func)
except Exception as e:
print(e)
continue
next = False
while not next:
io.process_messages(0.005)
if jobs_count() == 0:
for (x, y, ev, flags) in io.get_mouse_events(wnd_name):
x, y = int(x / zoom_factor), int(y / zoom_factor)
ed.set_mouse_pos(x, y)
if filepath is not None:
if ev == io.EVENT_LBUTTONDOWN:
ed.mask_point(1)
elif ev == io.EVENT_RBUTTONDOWN:
ed.mask_point(0)
elif ev == io.EVENT_MBUTTONDOWN:
ed.mask_finish()
elif ev == io.EVENT_MOUSEWHEEL:
if flags & 0x80000000 != 0:
if flags & 0x8 != 0:
ed.undo_to_begin_point()
else:
ed.undo_point()
else:
if flags & 0x8 != 0:
ed.redo_to_end_point()
else:
ed.redo_point()
for key, chr_key, ctrl_pressed, alt_pressed, shift_pressed in io.get_key_events(
wnd_name):
if chr_key == 'q' or chr_key == 'z':
do_prev_count = 1 if not shift_pressed else 10
elif chr_key == '-':
zoom_factor = np.clip(zoom_factor - 0.1, 0.1, 4.0)
ed.set_screen_changed()
elif chr_key == '+':
zoom_factor = np.clip(zoom_factor + 0.1, 0.1, 4.0)
ed.set_screen_changed()
elif key == 27: #esc
is_exit = True
next = True
break
elif filepath is not None:
if chr_key == 'e':
do_save_move_count = 1 if not shift_pressed else 10
elif chr_key == 'c':
do_save_count = 1 if not shift_pressed else 10
elif chr_key == 'w':
do_skip_move_count = 1 if not shift_pressed else 10
elif chr_key == 'x':
do_skip_count = 1 if not shift_pressed else 10
if do_prev_count > 0:
do_prev_count -= 1
if len(done_paths) > 0:
if filepath is not None:
image_paths.insert(0, filepath)
filepath = done_paths.pop(-1)
done_images_types[filepath.name] = 0
if filepath.parent != input_path:
new_filename_path = input_path / filepath.name
filepath.rename(new_filename_path)
image_paths.insert(0, new_filename_path)
else:
image_paths.insert(0, filepath)
next = True
elif filepath is not None:
if do_save_move_count > 0:
do_save_move_count -= 1
ed.mask_finish()
dflimg.embed_and_set(str(filepath),
ie_polys=ed.get_ie_polys())
done_paths += [confirmed_path / filepath.name]
done_images_types[filepath.name] = 2
filepath.rename(done_paths[-1])
next = True
elif do_save_count > 0:
do_save_count -= 1
ed.mask_finish()
dflimg.embed_and_set(str(filepath),
ie_polys=ed.get_ie_polys())
done_paths += [filepath]
done_images_types[filepath.name] = 2
next = True
elif do_skip_move_count > 0:
do_skip_move_count -= 1
done_paths += [skipped_path / filepath.name]
done_images_types[filepath.name] = 1
filepath.rename(done_paths[-1])
next = True
elif do_skip_count > 0:
do_skip_count -= 1
done_paths += [filepath]
done_images_types[filepath.name] = 1
next = True
else:
do_save_move_count = do_save_count = do_skip_move_count = do_skip_count = 0
if jobs_count() == 0:
if ed.switch_screen_changed():
screen = ed.make_screen()
if zoom_factor != 1.0:
h, w, c = screen.shape
screen = cv2.resize(
screen,
(int(w * zoom_factor), int(h * zoom_factor)))
io.show_image(wnd_name, screen)
io.process_messages(0.005)
io.destroy_all_windows()
def sort_final(input_path):
io.log_info ("Performing final sort.")
target_count = io.input_int ("Target number of images? (default:2000) : ", 2000)
img_list, trash_img_list = FinalLoaderSubprocessor( Path_utils.get_image_paths(input_path) ).run()
final_img_list = []
grads = 128
imgs_per_grad = round (target_count / grads)
grads_space = np.linspace (-1.0,1.0,grads)
yaws_sample_list = [None]*grads
for g in io.progress_bar_generator ( range(grads), "Sort by yaw"):
yaw = grads_space[g]
next_yaw = grads_space[g+1] if g < grads-1 else yaw
yaw_samples = []
for img in img_list:
s_yaw = -img[3]
if (g == 0 and s_yaw < next_yaw) or \
(g < grads-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
(g == grads-1 and s_yaw >= yaw):
yaw_samples += [ img ]
if len(yaw_samples) > 0:
yaws_sample_list[g] = yaw_samples
total_lack = 0
for g in io.progress_bar_generator ( range(grads), ""):
img_list = yaws_sample_list[g]
img_list_len = len(img_list) if img_list is not None else 0
lack = imgs_per_grad - img_list_len
total_lack += max(lack, 0)
imgs_per_grad += total_lack // grads
sharpned_imgs_per_grad = imgs_per_grad*10
for g in io.progress_bar_generator ( range (grads), "Sort by blur"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
if len(img_list) > imgs_per_grad*2:
trash_img_list += img_list[len(img_list) // 2:]
img_list = img_list[0: len(img_list) // 2]
if len(img_list) > sharpned_imgs_per_grad:
trash_img_list += img_list[sharpned_imgs_per_grad:]
img_list = img_list[0:sharpned_imgs_per_grad]
yaws_sample_list[g] = img_list
for g in io.progress_bar_generator ( range (grads), "Sort by hist"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
for i in range( len(img_list) ):
score_total = 0
for j in range( len(img_list) ):
if i == j:
continue
score_total += cv2.compareHist(img_list[i][2], img_list[j][2], cv2.HISTCMP_BHATTACHARYYA)
img_list[i][3] = score_total
yaws_sample_list[g] = sorted(img_list, key=operator.itemgetter(3), reverse=True)
for g in io.progress_bar_generator ( range (grads), "Fetching best"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
final_img_list += img_list[0:imgs_per_grad]
trash_img_list += img_list[imgs_per_grad:]
return final_img_list, trash_img_list
def main(input_dir,
output_dir,
debug,
detector='mt',
multi_gpu=True,
cpu_only=False,
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
image_size=256,
face_type='full_face'):
print("Running extractor.\r\n")
input_path = Path(input_dir)
output_path = Path(output_dir)
face_type = FaceType.fromString(face_type)
if not input_path.exists():
print('Input directory not found. Please ensure it exists.')
return
if output_path.exists():
if not manual_output_debug_fix:
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
debug = True
detector = 'manual'
print(
'Performing re-extract frames which were deleted from _debug directory.'
)
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(
input_path, verbose=True)
if debug:
debug_output_path = Path(str(output_path) + '_debug')
if manual_output_debug_fix:
if not debug_output_path.exists():
print("%s not found " % (str(debug_output_path)))
return
input_path_image_paths = DeletedFilesSearcherSubprocessor(
input_path_image_paths,
Path_utils.get_image_paths(debug_output_path)).process()
input_path_image_paths = sorted(input_path_image_paths)
else:
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
print('Performing manual extract...')
extracted_faces = ExtractSubprocessor(
[(filename, []) for filename in input_path_image_paths],
'landmarks',
image_size,
face_type,
debug,
cpu_only=cpu_only,
manual=True,
manual_window_size=manual_window_size).process()
else:
print('Performing 1st pass...')
extracted_rects = ExtractSubprocessor(
[(x, ) for x in input_path_image_paths],
'rects',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
detector=detector).process()
print('Performing 2nd pass...')
extracted_faces = ExtractSubprocessor(extracted_rects,
'landmarks',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False).process()
if manual_fix:
print('Performing manual fix...')
if all(
np.array(
[len(data[1]) > 0
for data in extracted_faces]) == True):
print('All faces are detected, manual fix not needed.')
else:
extracted_faces = ExtractSubprocessor(
extracted_faces,
'landmarks',
image_size,
face_type,
debug,
manual=True,
manual_window_size=manual_window_size).process()
if len(extracted_faces) > 0:
print('Performing 3rd pass...')
final_imgs_paths = ExtractSubprocessor(
extracted_faces,
'final',
image_size,
face_type,
debug,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
output_path=output_path).process()
faces_detected = len(final_imgs_paths)
print('-------------------------')
print('Images found: %d' % (images_found))
print('Faces detected: %d' % (faces_detected))
print('-------------------------')
def 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:
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
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)
else:
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info('Performing manual extract...')
extracted_faces = ExtractSubprocessor(
[(filename, []) for filename in input_path_image_paths],
'landmarks',
image_size,
face_type,
debug_dir,
cpu_only=cpu_only,
manual=True,
manual_window_size=manual_window_size).run()
else:
io.log_info('Performing 1st pass...')
extracted_rects = ExtractSubprocessor(
[(x, ) for x in input_path_image_paths],
'rects',
image_size,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
detector=detector).run()
io.log_info('Performing 2nd pass...')
extracted_faces = ExtractSubprocessor(extracted_rects,
'landmarks',
image_size,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False).run()
if manual_fix:
io.log_info('Performing manual fix...')
if all(
np.array(
[len(data[1]) > 0
for data in extracted_faces]) == True):
io.log_info(
'All faces are detected, manual fix not needed.')
else:
extracted_faces = ExtractSubprocessor(
extracted_faces,
'landmarks',
image_size,
face_type,
debug_dir,
manual=True,
manual_window_size=manual_window_size).run()
if len(extracted_faces) > 0:
io.log_info('Performing 3rd pass...')
final_imgs_paths = ExtractSubprocessor(
extracted_faces,
'final',
image_size,
face_type,
debug_dir,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
output_path=output_path).run()
faces_detected = len(final_imgs_paths)
io.log_info('-------------------------')
io.log_info('Images found: %d' % (images_found))
io.log_info('Faces detected: %d' % (faces_detected))
io.log_info('-------------------------')
def sort_final(input_path, include_by_blur=True):
io.log_info("运行终极[final]排序工具.")
target_count = io.input_int("Target number of images? (default:2000) : ",
2000)
img_list, trash_img_list = FinalLoaderSubprocessor(
Path_utils.get_image_paths(input_path), include_by_blur).run()
final_img_list = []
grads = 128
imgs_per_grad = round(target_count / grads)
grads_space = np.linspace(-1.0, 1.0, grads)
yaws_sample_list = [None] * grads
for g in io.progress_bar_generator(range(grads), "Sort by yaw"):
yaw = grads_space[g]
next_yaw = grads_space[g + 1] if g < grads - 1 else yaw
yaw_samples = []
for img in img_list:
s_yaw = -img[3]
if (g == 0 and s_yaw < next_yaw) or \
(g < grads-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
(g == grads-1 and s_yaw >= yaw):
yaw_samples += [img]
if len(yaw_samples) > 0:
yaws_sample_list[g] = yaw_samples
total_lack = 0
for g in io.progress_bar_generator(range(grads), ""):
img_list = yaws_sample_list[g]
img_list_len = len(img_list) if img_list is not None else 0
lack = imgs_per_grad - img_list_len
total_lack += max(lack, 0)
imgs_per_grad += total_lack // grads
if include_by_blur:
sharpned_imgs_per_grad = imgs_per_grad * 10
for g in io.progress_bar_generator(range(grads), "Sort by blur"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
img_list = sorted(img_list,
key=operator.itemgetter(1),
reverse=True)
if len(img_list) > sharpned_imgs_per_grad:
trash_img_list += img_list[sharpned_imgs_per_grad:]
img_list = img_list[0:sharpned_imgs_per_grad]
yaws_sample_list[g] = img_list
yaws_sample_list = FinalHistDissimSubprocessor(yaws_sample_list).run()
for g in io.progress_bar_generator(range(grads), "Fetching best"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
final_img_list += img_list[0:imgs_per_grad]
trash_img_list += img_list[imgs_per_grad:]
return final_img_list, trash_img_list
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 __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None,
batch_size=0,
write_preview_history = False,
debug = False, **in_options
):
print ("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.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
self.batch_size = batch_size
self.write_preview_history = write_preview_history
self.debug = debug
self.supress_std_once = ('TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1')
if self.model_data_path.exists():
model_data = pickle.loads ( self.model_data_path.read_bytes() )
self.epoch = model_data['epoch']
self.options = model_data['options']
self.loss_history = model_data['loss_history'] if 'loss_history' in model_data.keys() else []
self.generator_dict_states = model_data['generator_dict_states'] if 'generator_dict_states' in model_data.keys() else None
self.sample_for_preview = model_data['sample_for_preview'] if 'sample_for_preview' in model_data.keys() else None
else:
self.epoch = 0
self.options = {}
self.loss_history = []
self.generator_dict_states = None
self.sample_for_preview = None
if self.write_preview_history:
self.preview_history_path = self.model_path / ( '%s_history' % (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()
self.gpu_config = gpufmkmgr.GPUConfig(allow_growth=False, **in_options)
self.gpu_total_vram_gb = self.gpu_config.gpu_total_vram_gb
if self.epoch == 0:
#first run
self.options['created_vram_gb'] = self.gpu_total_vram_gb
self.created_vram_gb = self.gpu_total_vram_gb
else:
#not first run
if 'created_vram_gb' in self.options.keys():
self.created_vram_gb = self.options['created_vram_gb']
else:
self.options['created_vram_gb'] = self.gpu_total_vram_gb
self.created_vram_gb = self.gpu_total_vram_gb
self.tf = gpufmkmgr.import_tf( self.gpu_config )
self.tf_sess = gpufmkmgr.get_tf_session()
self.keras = gpufmkmgr.import_keras()
self.keras_contrib = gpufmkmgr.import_keras_contrib()
self.onInitialize(**in_options)
if self.debug or self.batch_size == 0:
self.batch_size = 1
if self.is_training_mode:
if self.generator_list is None:
raise Exception( 'You didnt set_training_data_generators()')
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, TrainingDataGeneratorBase):
raise Exception('training data generator is not subclass of TrainingDataGeneratorBase')
if self.generator_dict_states is not None and i < len(self.generator_dict_states):
generator.set_dict_state ( self.generator_dict_states[i] )
if self.sample_for_preview is None:
self.sample_for_preview = self.generate_next_sample()
print ("===== Model summary =====")
print ("== Model name: " + self.get_model_name())
print ("==")
print ("== Current epoch: " + str(self.epoch) )
print ("==")
print ("== Options:")
print ("== |== batch_size : %s " % (self.batch_size) )
print ("== |== multi_gpu : %s " % (self.gpu_config.multi_gpu) )
for key in self.options.keys():
print ("== |== %s : %s" % (key, self.options[key]) )
print ("== Running on:")
for idx in self.gpu_config.gpu_idxs:
print ("== |== [%d : %s]" % (idx, gpufmkmgr.getDeviceName(idx)) )
if self.gpu_total_vram_gb == 2:
print ("==")
print ("== WARNING: You are using 2GB GPU. Result quality may be significantly decreased.")
print ("== If training does not start, close all programs and try again.")
print ("== Also you can disable Windows Aero Desktop to get extra free VRAM.")
print ("==")
print ("=========================")
def sort_by_absdiff(input_path):
io.log_info("Sorting by absolute difference...")
is_sim = io.input_bool("Sort by similar? ( y/n ?:help skip:y ) : ",
True,
help_message="Otherwise sort by dissimilar.")
from nnlib import nnlib
exec(nnlib.import_all(device_config=nnlib.device.Config()), locals(),
globals())
image_paths = Path_utils.get_image_paths(input_path)
image_paths_len = len(image_paths)
batch_size = 1024
batch_size_remain = image_paths_len % batch_size
i_t = Input((256, 256, 3))
j_t = Input((256, 256, 3))
outputs = []
for i in range(batch_size):
outputs += [K.sum(K.abs(i_t - j_t[i]), axis=[1, 2, 3])]
func_bs_full = K.function([i_t, j_t], outputs)
outputs = []
for i in range(batch_size_remain):
outputs += [K.sum(K.abs(i_t - j_t[i]), axis=[1, 2, 3])]
func_bs_remain = K.function([i_t, j_t], outputs)
import h5py
db_file_path = Path(tempfile.gettempdir()) / 'sort_cache.hdf5'
db_file = h5py.File(str(db_file_path), "w")
db = db_file.create_dataset("results", (image_paths_len, image_paths_len),
compression="gzip")
pg_len = image_paths_len // batch_size
if batch_size_remain != 0:
pg_len += 1
pg_len = int((pg_len * pg_len - pg_len) / 2 + pg_len)
io.progress_bar("Computing", pg_len)
j = 0
while j < image_paths_len:
j_images = [cv2_imread(x) for x in image_paths[j:j + batch_size]]
j_images_len = len(j_images)
func = func_bs_remain if image_paths_len - j < batch_size else func_bs_full
i = 0
while i < image_paths_len:
if i >= j:
i_images = [
cv2_imread(x) for x in image_paths[i:i + batch_size]
]
i_images_len = len(i_images)
result = func([i_images, j_images])
db[j:j + j_images_len, i:i + i_images_len] = np.array(result)
io.progress_bar_inc(1)
i += batch_size
db_file.flush()
j += batch_size
io.progress_bar_close()
next_id = 0
sorted = [next_id]
for i in io.progress_bar_generator(range(image_paths_len - 1), "Sorting"):
id_ar = np.concatenate([db[:next_id, next_id], db[next_id, next_id:]])
id_ar = np.argsort(id_ar)
next_id = np.setdiff1d(id_ar, sorted, True)[0 if is_sim else -1]
sorted += [next_id]
db_file.close()
db_file_path.unlink()
img_list = [(image_paths[x], ) for x in sorted]
return img_list, []
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 sort_final(input_path, include_by_blur=True):
io.log_info("Performing final sort.")
target_count = io.input_int("Target number of images? (default:2000) : ",
2000)
img_list, trash_img_list = FinalLoaderSubprocessor(
Path_utils.get_image_paths(input_path), include_by_blur).run()
final_img_list = []
grads = 128
imgs_per_grad = round(target_count / grads)
grads_space = np.linspace(-1.0, 1.0, grads)
yaws_sample_list = [None] * grads
for g in io.progress_bar_generator(range(grads), "Sort by yaw"):
yaw = grads_space[g]
next_yaw = grads_space[g + 1] if g < grads - 1 else yaw
yaw_samples = []
for img in img_list:
s_yaw = -img[3]
if (g == 0 and s_yaw < next_yaw) or \
(g < grads-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
(g == grads-1 and s_yaw >= yaw):
yaw_samples += [img]
if len(yaw_samples) > 0:
yaws_sample_list[g] = yaw_samples
total_lack = 0
for g in io.progress_bar_generator(range(grads), ""):
img_list = yaws_sample_list[g]
img_list_len = len(img_list) if img_list is not None else 0
lack = imgs_per_grad - img_list_len
total_lack += max(lack, 0)
imgs_per_grad += total_lack // grads
if include_by_blur:
sharpned_imgs_per_grad = imgs_per_grad * 10
for g in io.progress_bar_generator(range(grads), "Sort by blur"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
img_list = sorted(img_list,
key=operator.itemgetter(1),
reverse=True)
if len(img_list) > sharpned_imgs_per_grad:
trash_img_list += img_list[sharpned_imgs_per_grad:]
img_list = img_list[0:sharpned_imgs_per_grad]
yaws_sample_list[g] = img_list
yaw_pitch_sample_list = [None] * grads
pitch_grads = imgs_per_grad
for g in io.progress_bar_generator(range(grads), "Sort by pitch"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
pitch_sample_list = [None] * pitch_grads
grads_space = np.linspace(-1.0, 1.0, pitch_grads)
for pg in range(pitch_grads):
pitch = grads_space[pg]
next_pitch = grads_space[pg + 1] if pg < pitch_grads - 1 else pitch
pitch_samples = []
for img in img_list:
s_pitch = img[4]
if (pg == 0 and s_pitch < next_pitch) or \
(pg < pitch_grads-1 and s_pitch >= pitch and s_pitch < next_pitch) or \
(pg == pitch_grads-1 and s_pitch >= pitch):
pitch_samples += [img]
if len(pitch_samples) > 0:
pitch_sample_list[pg] = pitch_samples
yaw_pitch_sample_list[g] = pitch_sample_list
yaw_pitch_sample_list = FinalHistDissimSubprocessor(
yaw_pitch_sample_list).run()
for g in io.progress_bar_generator(range(grads), "Fetching the best"):
pitch_sample_list = yaw_pitch_sample_list[g]
if pitch_sample_list is None:
continue
n = imgs_per_grad
while n > 0:
n_prev = n
for pg in range(pitch_grads):
img_list = pitch_sample_list[pg]
if img_list is None:
continue
final_img_list += [img_list.pop(0)]
if len(img_list) == 0:
pitch_sample_list[pg] = None
n -= 1
if n == 0:
break
if n_prev == n:
break
for pg in range(pitch_grads):
img_list = pitch_sample_list[pg]
if img_list is None:
continue
trash_img_list += img_list
return final_img_list, trash_img_list
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 and random intensity? ( y/n default:y ?:help) : ",
True,
help_message=
"Otherwise faceset will be relighted with predefined 7 light directions but with random intensity."
)
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:
dflimg = DFLIMG.load(Path(filepath))
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()
for filepath in io.progress_bar_generator(image_paths, "Relighting"):
try:
dflimg = DFLIMG.load(Path(filepath))
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)
inten = np.random.random() * 0.3 + 0.3
relighted_imgs = [
dpr.relight(img,
alt=alt,
azi=azi,
intensity=inten,
lighten=lighten)
]
else:
if not manual and not random_one:
inten = np.random.random() * 0.3 + 0.3
alt_azi_ar = [(60, 0, inten), (60, 60, inten),
(0, 60, inten), (-60, 60, inten),
(-60, 0, inten), (-60, -60, inten),
(0, -60, inten), (60, -60, inten)]
relighted_imgs = [
dpr.relight(img,
alt=alt,
azi=azi,
intensity=inten,
lighten=lighten)
for (alt, azi, inten) 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.remove_source_filename()
dflimg.embed_and_set(relighted_filepath, relighted=True)
except:
io.log_err(
f"Exception occured while processing file {filepath.name}. Error: {traceback.format_exc()}"
)
def load(sample_type, samples_path, target_samples_path=None, person_id_mode=False, use_caching=False):
cache = SampleLoader.cache
if str(samples_path) not in cache.keys():
cache[str(samples_path)] = [None]*SampleType.QTY
datas = cache[str(samples_path)]
if sample_type == SampleType.IMAGE:
if datas[sample_type] is None:
datas[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( Path_utils.get_image_paths(samples_path), "Loading") ]
elif sample_type == SampleType.FACE:
if datas[sample_type] is None:
if not use_caching:
datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] )
else:
samples_dat = samples_path / 'samples.dat'
if samples_dat.exists():
io.log_info (f"Using saved samples info from '{samples_dat}' ")
all_samples = pickle.loads(samples_dat.read_bytes())
if person_id_mode:
for samples in all_samples:
for sample in samples:
sample.filename = str( samples_path / Path(sample.filename) )
else:
for sample in all_samples:
sample.filename = str( samples_path / Path(sample.filename) )
datas[sample_type] = all_samples
else:
if person_id_mode:
dir_names = Path_utils.get_all_dir_names(samples_path)
all_samples = []
for i, dir_name in io.progress_bar_generator( [*enumerate(dir_names)] , "Loading"):
all_samples += [ SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename, person_id=i) for filename in Path_utils.get_image_paths( samples_path / dir_name ) ], silent=True ) ]
datas[sample_type] = all_samples
else:
datas[sample_type] = all_samples = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] )
if person_id_mode:
for samples in all_samples:
for sample in samples:
sample.filename = str(Path(sample.filename).relative_to(samples_path))
else:
for sample in all_samples:
sample.filename = str(Path(sample.filename).relative_to(samples_path))
samples_dat.write_bytes (pickle.dumps(all_samples))
if person_id_mode:
for samples in all_samples:
for sample in samples:
sample.filename = str( samples_path / Path(sample.filename) )
else:
for sample in all_samples:
sample.filename = str( samples_path / Path(sample.filename) )
elif sample_type == SampleType.FACE_TEMPORAL_SORTED:
if datas[sample_type] is None:
datas[sample_type] = SampleLoader.upgradeToFaceTemporalSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
elif sample_type == SampleType.FACE_YAW_SORTED:
if datas[sample_type] is None:
datas[sample_type] = SampleLoader.upgradeToFaceYawSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
elif sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
if datas[sample_type] is None:
if target_samples_path is None:
raise Exception('target_samples_path is None for FACE_YAW_SORTED_AS_TARGET')
datas[sample_type] = SampleLoader.upgradeToFaceYawSortedAsTargetSamples( SampleLoader.load(SampleType.FACE_YAW_SORTED, samples_path), SampleLoader.load(SampleType.FACE_YAW_SORTED, target_samples_path) )
return datas[sample_type]
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,
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.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[
'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.iter != 0 and io.input_in_time(
"Press enter in 2 seconds to override model settings.", 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_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 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 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 ask_sort_by_yaw:
if (self.iter == 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. 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.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()))
else:
self.preview_history_path = self.model_path / (
'%d_%s_history' %
(self.device_args['force_gpu_idx'], self.get_model_name()))
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 (self.iter == 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 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 main(args, device_args):
io.log_info("Running converter.\r\n")
training_data_src_dir = args.get('training_data_src_dir', None)
training_data_src_path = Path(
training_data_src_dir) if training_data_src_dir is not None else None
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,
training_data_src_path=training_data_src_path)
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 = dflimg.get_source_filename()
if source_filename is None or source_filename == "_":
continue
source_filename = Path(source_filename)
source_filename_stem = 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 main(input_dir, output_dir):
input_path = Path(input_dir)
output_path = Path(output_dir)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
if not output_path.exists():
output_path.mkdir(parents=True)
wnd_name = "Labeling tool"
io.named_window(wnd_name)
io.capture_mouse(wnd_name)
io.capture_keys(wnd_name)
#for filename in io.progress_bar_generator (Path_utils.get_image_paths(input_path), desc="Labeling"):
for filename in Path_utils.get_image_paths(input_path):
filepath = Path(filename)
if filepath.suffix == '.png':
dflimg = DFLPNG.load(str(filepath))
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load(str(filepath))
else:
dflimg = None
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
continue
lmrks = dflimg.get_landmarks()
lmrks_list = lmrks.tolist()
orig_img = cv2_imread(str(filepath))
h, w, c = orig_img.shape
mask_orig = LandmarksProcessor.get_image_hull_mask(
orig_img.shape, lmrks).astype(np.uint8)[:, :, 0]
ero_dil_rate = w // 8
mask_ero = cv2.erode(
mask_orig,
cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(ero_dil_rate, ero_dil_rate)),
iterations=1)
mask_dil = cv2.dilate(mask_orig,
cv2.getStructuringElement(
cv2.MORPH_ELLIPSE,
(ero_dil_rate, ero_dil_rate)),
iterations=1)
#mask_bg = np.zeros(orig_img.shape[:2],np.uint8)
mask_bg = 1 - mask_dil
mask_bgp = np.ones(orig_img.shape[:2],
np.uint8) #default - all background possible
mask_fg = np.zeros(orig_img.shape[:2], np.uint8)
mask_fgp = np.zeros(orig_img.shape[:2], np.uint8)
img = orig_img.copy()
l_thick = 2
def draw_4_lines(masks_out, pts, thickness=1):
fgp, fg, bg, bgp = masks_out
h, w = fg.shape
fgp_pts = []
fg_pts = np.array([pts[i:i + 2] for i in range(len(pts) - 1)])
bg_pts = []
bgp_pts = []
for i in range(len(fg_pts)):
a, b = line = fg_pts[i]
ba = b - a
v = ba / npl.norm(ba)
ccpv = np.array([v[1], -v[0]])
cpv = np.array([-v[1], v[0]])
step = 1 / max(np.abs(cpv))
fgp_pts.append(
np.clip(line + ccpv * step * thickness, 0,
w - 1).astype(np.int))
bg_pts.append(
np.clip(line + cpv * step * thickness, 0,
w - 1).astype(np.int))
bgp_pts.append(
np.clip(line + cpv * step * thickness * 2, 0,
w - 1).astype(np.int))
fgp_pts = np.array(fgp_pts)
bg_pts = np.array(bg_pts)
bgp_pts = np.array(bgp_pts)
cv2.polylines(fgp, fgp_pts, False, (1, ), thickness=thickness)
cv2.polylines(fg, fg_pts, False, (1, ), thickness=thickness)
cv2.polylines(bg, bg_pts, False, (1, ), thickness=thickness)
cv2.polylines(bgp, bgp_pts, False, (1, ), thickness=thickness)
def draw_lines(masks_steps, pts, thickness=1):
lines = np.array([pts[i:i + 2] for i in range(len(pts) - 1)])
for mask, step in masks_steps:
h, w = mask.shape
mask_lines = []
for i in range(len(lines)):
a, b = line = lines[i]
ba = b - a
ba_len = npl.norm(ba)
if ba_len != 0:
v = ba / ba_len
pv = np.array([-v[1], v[0]])
pv_inv_max = 1 / max(np.abs(pv))
mask_lines.append(
np.clip(line + pv * pv_inv_max * thickness * step,
0, w - 1).astype(np.int))
else:
mask_lines.append(np.array(line, dtype=np.int))
cv2.polylines(mask,
mask_lines,
False, (1, ),
thickness=thickness)
def draw_fill_convex(mask_out, pts, scale=1.0):
hull = cv2.convexHull(np.array(pts))
if scale != 1.0:
pts_count = hull.shape[0]
sum_x = np.sum(hull[:, 0, 0])
sum_y = np.sum(hull[:, 0, 1])
hull_center = np.array([sum_x / pts_count, sum_y / pts_count])
hull = hull_center + (hull - hull_center) * scale
hull = hull.astype(pts.dtype)
cv2.fillConvexPoly(mask_out, hull, (1, ))
def get_gc_mask_bgr(gc_mask):
h, w = gc_mask.shape
bgr = np.zeros((h, w, 3), dtype=np.uint8)
bgr[gc_mask == 0] = (0, 0, 0)
bgr[gc_mask == 1] = (255, 255, 255)
bgr[gc_mask == 2] = (0, 0, 255) #RED
bgr[gc_mask == 3] = (0, 255, 0) #GREEN
return bgr
def get_gc_mask_result(gc_mask):
return np.where((gc_mask == 1) + (gc_mask == 3), 1,
0).astype(np.int)
#convex inner of right chin to end of right eyebrow
#draw_fill_convex ( mask_fgp, lmrks_list[8:17]+lmrks_list[26:27] )
#convex inner of start right chin to right eyebrow
#draw_fill_convex ( mask_fgp, lmrks_list[8:9]+lmrks_list[22:27] )
#convex inner of nose
draw_fill_convex(mask_fgp, lmrks[27:36])
#convex inner of nose half
draw_fill_convex(mask_fg, lmrks[27:36], scale=0.5)
#left corner of mouth to left corner of nose
#draw_lines ( [ (mask_fg,0), ], lmrks_list[49:50]+lmrks_list[32:33], l_thick)
#convex inner: right corner of nose to centers of eyebrows
#draw_fill_convex ( mask_fgp, lmrks_list[35:36]+lmrks_list[19:20]+lmrks_list[24:25])
#right corner of mouth to right corner of nose
#draw_lines ( [ (mask_fg,0), ], lmrks_list[54:55]+lmrks_list[35:36], l_thick)
#left eye
#draw_fill_convex ( mask_fg, lmrks_list[36:40] )
#right eye
#draw_fill_convex ( mask_fg, lmrks_list[42:48] )
#right chin
draw_lines([
(mask_bg, 0),
(mask_fg, -1),
], lmrks[8:17], l_thick)
#left eyebrow center to right eyeprow center
draw_lines([
(mask_bg, -1),
(mask_fg, 0),
], lmrks_list[19:20] + lmrks_list[24:25], l_thick)
# #draw_lines ( [ (mask_bg,-1), (mask_fg,0), ], lmrks_list[24:25] + lmrks_list[19:17:-1], l_thick)
#half right eyebrow to end of right chin
draw_lines([
(mask_bg, -1),
(mask_fg, 0),
], lmrks_list[24:27] + lmrks_list[16:17], l_thick)
#import code
#code.interact(local=dict(globals(), **locals()))
#compose mask layers
gc_mask = np.zeros(orig_img.shape[:2], np.uint8)
gc_mask[mask_bgp == 1] = 2
gc_mask[mask_fgp == 1] = 3
gc_mask[mask_bg == 1] = 0
gc_mask[mask_fg == 1] = 1
gc_bgr_before = get_gc_mask_bgr(gc_mask)
#io.show_image (wnd_name, gc_mask )
##points, hierarcy = cv2.findContours(original_mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
##gc_mask = ( (1-erode_mask)*2 + erode_mask )# * dilate_mask
#gc_mask = (1-erode_mask)*2 + erode_mask
#cv2.addWeighted(
#gc_mask = mask_0_27 + (1-mask_0_27)*2
#
##import code
##code.interact(local=dict(globals(), **locals()))
#
#rect = (1,1,img.shape[1]-2,img.shape[0]-2)
#
#
cv2.grabCut(img, gc_mask, None, np.zeros((1, 65), np.float64),
np.zeros((1, 65), np.float64), 5, cv2.GC_INIT_WITH_MASK)
gc_bgr = get_gc_mask_bgr(gc_mask)
gc_mask_result = get_gc_mask_result(gc_mask)
gc_mask_result_1 = gc_mask_result[:, :, np.newaxis]
#import code
#code.interact(local=dict(globals(), **locals()))
orig_img_gc_layers_masked = (0.5 * orig_img + 0.5 * gc_bgr).astype(
np.uint8)
orig_img_gc_before_layers_masked = (0.5 * orig_img +
0.5 * gc_bgr_before).astype(
np.uint8)
pink_bg = np.full(orig_img.shape, (255, 0, 255), dtype=np.uint8)
orig_img_result = orig_img * gc_mask_result_1
orig_img_result_pinked = orig_img_result + pink_bg * (1 -
gc_mask_result_1)
#io.show_image (wnd_name, blended_img)
##gc_mask, bgdModel, fgdModel =
#
#mask2 = np.where((gc_mask==1) + (gc_mask==3),255,0).astype('uint8')[:,:,np.newaxis]
#mask2 = np.repeat(mask2, (3,), -1)
#
##mask2 = np.where(gc_mask!=0,255,0).astype('uint8')
#blended_img = orig_img #-\
# #0.3 * np.full(original_img.shape, (50,50,50)) * (1-mask_0_27)[:,:,np.newaxis]
# #0.3 * np.full(original_img.shape, (50,50,50)) * (1-dilate_mask)[:,:,np.newaxis] +\
# #0.3 * np.full(original_img.shape, (50,50,50)) * (erode_mask)[:,:,np.newaxis]
#blended_img = np.clip(blended_img, 0, 255).astype(np.uint8)
##import code
##code.interact(local=dict(globals(), **locals()))
orig_img_lmrked = orig_img.copy()
LandmarksProcessor.draw_landmarks(orig_img_lmrked,
lmrks,
transparent_mask=True)
screen = np.concatenate([
orig_img_gc_before_layers_masked,
orig_img_gc_layers_masked,
orig_img,
orig_img_lmrked,
orig_img_result_pinked,
orig_img_result,
],
axis=1)
io.show_image(wnd_name, screen.astype(np.uint8))
while True:
io.process_messages()
for (x, y, ev, flags) in io.get_mouse_events(wnd_name):
pass
#print (x,y,ev,flags)
key_events = [ev for ev, in io.get_key_events(wnd_name)]
for key in key_events:
if key == ord('1'):
pass
if key == ord('2'):
pass
if key == ord('3'):
pass
if ord(' ') in key_events:
break
import code
code.interact(local=dict(globals(), **locals()))
#original_mask = np.ones(original_img.shape[:2],np.uint8)*2
#cv2.drawContours(original_mask, points, -1, (1,), 1)
def sort_by_hist(input_path):
io.log_info("根据[histogram similarity]排序...")
img_list = HistSsimSubprocessor(
Path_utils.get_image_paths(input_path)).run()
return img_list
def main(args, device_args):
io.log_info("Running converter.\r\n")
aligned_dir = args.get('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 output_path.exists():
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if not model_path.exists():
io.log_err('Model directory not found. Please ensure it exists.')
return
import models
model = models.import_model(args['model_name'])(
model_path, device_args=device_args)
converter = model.get_converter()
converter.dummy_predict()
alignments = None
if converter.type == Converter.TYPE_FACE:
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 = {}
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[source_filename_stem].append(
dflimg.get_source_landmarks())
files_processed, faces_processed = ConvertSubprocessor(
converter=converter,
input_path_image_paths=Path_utils.get_image_paths(input_path),
output_path=output_path,
alignments=alignments,
debug=args.get('debug', False)).run()
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
def sort_by_hue(input_path):
io.log_info ("Sorting by hue...")
img_list = [ [x, np.mean ( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[...,0].flatten() )] for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading") ]
io.log_info ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, pretraining_data_path=None, is_training=False, debug = False, device_args = None,
ask_enable_autobackup=True,
ask_write_preview_history=True,
ask_target_iter=True,
ask_batch_size=True,
ask_random_flip=True, **kwargs):
device_args['force_gpu_idx'] = device_args.get('force_gpu_idx',-1)
device_args['cpu_only'] = True if debug else 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.debug = debug
self.is_training_mode = is_training
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.")
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):
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.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.batch_size = self.options.get('batch_size', 0)
if ask_random_flip:
default_random_flip = self.options.get('random_flip', True)
if (self.iter == 0 or ask_override):
self.options['random_flip'] = io.input_bool(f"Flip faces randomly? (y/n ?:help skip:{yn_str[default_random_flip]}) : ", default_random_flip, 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', default_random_flip)
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.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():
io.log_info ("Choose image for the preview history. [p] - next. [enter] - confirm.")
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()
try:
self.get_static_preview()
except:
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,
detector='mt',
multi_gpu=True,
manual_fix=False,
image_size=256,
face_type='full_face'):
print("Running extractor.\r\n")
input_path = Path(input_dir)
output_path = Path(output_dir)
if not input_path.exists():
print('Input directory not found. Please ensure it exists.')
return
if output_path.exists():
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if debug:
debug_output_path = Path(str(output_path) + '_debug')
if debug_output_path.exists():
for filename in Path_utils.get_image_paths(debug_output_path):
Path(filename).unlink()
else:
debug_output_path.mkdir(parents=True, exist_ok=True)
input_path_image_paths = Path_utils.get_image_paths(input_path)
images_found = len(input_path_image_paths)
if images_found != 0:
if detector == 'manual':
print('Performing manual extract...')
extracted_faces = [(filename, [])
for filename in input_path_image_paths]
extracted_faces = manual_pass(extracted_faces, image_size,
face_type)
else:
print('Performing 1st pass...')
extracted_rects = extract_pass([(x, )
for x in input_path_image_paths],
'rects',
image_size,
face_type,
debug,
multi_gpu,
detector=detector)
print('Performing 2nd pass...')
extracted_faces = extract_pass(extracted_rects, 'landmarks',
image_size, face_type, debug,
multi_gpu)
if manual_fix:
print('Performing manual fix...')
if all(
np.array(
[len(data[1]) > 0
for data in extracted_faces]) == True):
print('All faces are detected, manual fix not needed.')
else:
extracted_faces = manual_pass(extracted_faces, image_size,
face_type)
if len(extracted_faces) > 0:
print('Performing 3rd pass...')
final_imgs_paths = extract_pass(extracted_faces,
'final',
image_size,
face_type,
debug,
multi_gpu,
image_size,
output_path=output_path)
faces_detected = len(final_imgs_paths)
else:
faces_detected = 0
print('-------------------------')
print('Images found: %d' % (images_found))
print('Faces detected: %d' % (faces_detected))
print('-------------------------')
def main(input_dir, output_dir, aligned_dir, model_dir, model_name,
**in_options):
print("Running converter.\r\n")
debug = in_options['debug']
try:
input_path = Path(input_dir)
output_path = Path(output_dir)
aligned_path = Path(aligned_dir)
model_path = Path(model_dir)
if not input_path.exists():
print('Input directory not found. Please ensure it exists.')
return
if output_path.exists():
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if not aligned_path.exists():
print('Aligned directory not found. Please ensure it exists.')
return
if not model_path.exists():
print('Model directory not found. Please ensure it exists.')
return
model_sq = multiprocessing.Queue()
model_cq = multiprocessing.Queue()
model_lock = multiprocessing.Lock()
model_p = multiprocessing.Process(target=model_process,
args=(model_name, model_dir,
in_options, model_sq,
model_cq))
model_p.start()
while True:
if not model_cq.empty():
obj = model_cq.get()
obj_op = obj['op']
if obj_op == 'init':
converter = obj['converter']
break
alignments = {}
if converter.get_mode() == ConverterBase.MODE_FACE:
aligned_path_image_paths = Path_utils.get_image_paths(aligned_path)
for filename in tqdm(aligned_path_image_paths,
desc="Collecting alignments"):
a_png = AlignedPNG.load(str(filename))
if a_png is None:
print("%s - no embedded data found." % (filename))
continue
d = a_png.getFaceswapDictData()
if d is None or d['source_filename'] is None or d[
'source_rect'] is None or d['source_landmarks'] is None:
print("%s - no embedded data found." % (filename))
continue
source_filename_stem = Path(d['source_filename']).stem
if source_filename_stem not in alignments.keys():
alignments[source_filename_stem] = []
alignments[source_filename_stem].append(
np.array(d['source_landmarks']))
files_processed, faces_processed = ConvertSubprocessor(
converter=converter.copy_and_set_predictor(
model_process_predictor(model_sq, model_cq, model_lock)),
input_path_image_paths=Path_utils.get_image_paths(input_path),
output_path=output_path,
alignments=alignments,
**in_options).process()
model_sq.put({'op': 'close'})
model_p.join()
'''
if model_name == 'AVATAR':
output_path_image_paths = Path_utils.get_image_paths(output_path)
last_ok_frame = -1
for filename in output_path_image_paths:
filename_path = Path(filename)
stem = Path(filename).stem
try:
frame = int(stem)
except:
raise Exception ('Aligned avatars must be created from indexed sequence files.')
if frame-last_ok_frame > 1:
start = last_ok_frame + 1
end = frame - 1
print ("Filling gaps: [%d...%d]" % (start, end) )
for i in range (start, end+1):
shutil.copy ( str(filename), str( output_path / ('%.5d%s' % (i, filename_path.suffix )) ) )
last_ok_frame = frame
'''
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
def __init__(self,
model_path,
training_data_src_path=None,
training_data_dst_path=None,
debug=False,
force_gpu_idx=-1,
**in_options):
if force_gpu_idx == -1:
idxs_names_list = nnlib.device.getAllDevicesIdxsWithNamesList()
if len(idxs_names_list) > 1:
print("You have multi GPUs in a system: ")
for idx, name in idxs_names_list:
print("[%d] : %s" % (idx, name))
force_gpu_idx = input_int(
"Which GPU idx to choose? ( skip: best GPU ) : ", -1,
[x[0] for x in idxs_names_list])
self.force_gpu_idx = force_gpu_idx
print("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.supress_std_once = os.environ.get('TF_SUPPRESS_STD', '0') == '1'
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 input_in_time(
"Press enter in 2 seconds to override some model settings.", 2)
if self.epoch == 0:
print(
"\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'] = input_bool(
"Write preview history? (y/n ?:help skip:n/default) : ",
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, 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,
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'] = 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'] = 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(
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.drive = GoogleDriveSync(
model_dir=self.model_path,
key=input_str("Google Drive OAuth key (default: None) : ", None))
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']
if not self.sort_by_yaw:
self.options.pop('sort_by_yaw')
self.random_flip = self.options['random_flip']
if self.random_flip:
self.options.pop('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,
force_gpu_idx=self.force_gpu_idx,
**in_options))
self.device_config = nnlib.active_DeviceConfig
self.onInitialize(**in_options)
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.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.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 Exception('You didnt set_training_data_generators()')
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase):
raise Exception(
'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()
print("===== Model summary =====")
print("== Model name: " + self.get_model_name())
print("==")
print("== Current epoch: " + str(self.epoch))
print("==")
print("== Model options:")
for key in self.options.keys():
print("== |== %s : %s" % (key, self.options[key]))
if self.device_config.multi_gpu:
print("== |== multi_gpu : True ")
print("== Running on:")
if self.device_config.cpu_only:
print("== |== [CPU]")
else:
for idx in self.device_config.gpu_idxs:
print("== |== [%d : %s]" %
(idx, nnlib.device.getDeviceName(idx)))
if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[
0] == 2:
print("==")
print(
"== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."
)
print(
"== If training does not start, close all programs and try again."
)
print(
"== Also you can disable Windows Aero Desktop to get extra free VRAM."
)
print("==")
print("=========================")
'no_preview' : arguments.no_preview,
'debug' : arguments.debug,
'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ]
}
device_args = {'cpu_only' : arguments.cpu_only,
'force_gpu_idx' : arguments.force_gpu_idx,
}
from mainscripts import Trainer
Trainer.main(args, device_args)
p = subparsers.add_parser( "train", help="Trainer")
p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of extracted SRC faceset.")
p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of extracted DST faceset.")
p.add_argument('--pretraining-data-dir', action=fixPathAction, dest="pretraining_data_dir", default=None, help="Optional dir of extracted faceset that will be used in pretraining mode.")
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.")
p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.")
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.")
p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+')
p.set_defaults (func=process_train)
def process_convert(arguments):
os_utils.set_process_lowest_prio()
args = {'input_dir' : arguments.input_dir,
'output_dir' : arguments.output_dir,
'aligned_dir' : arguments.aligned_dir,
'avaperator_aligned_dir' : arguments.avaperator_aligned_dir,
'model_dir' : arguments.model_dir,
'model_name' : arguments.model_name,
'--pretraining-data-dir',
action=fixPathAction,
dest="pretraining_data_dir",
default=None,
help=
"Optional dir of extracted faceset that will be used in pretraining mode."
)
p.add_argument('--model-dir',
required=True,
action=fixPathAction,
dest="model_dir",
help="Model dir.")
p.add_argument('--model',
required=True,
dest="model_name",
choices=Path_utils.get_all_dir_names_startswith(
Path(__file__).parent / 'models', 'Model_'),
help="Type of model")
p.add_argument('--no-preview',
action="store_true",
dest="no_preview",
default=False,
help="Disable preview window.")
p.add_argument('--debug',
action="store_true",
dest="debug",
default=False,
help="Debug samples.")
p.add_argument('--cpu-only',
action="store_true",
dest="cpu_only",
default=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():
logger.error("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 = "png"
if lossless is None:
lossless = 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:
logger.error("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 = 25
if not lossless and bitrate is None:
bitrate = 16
input_image_paths = Path_utils.get_image_paths(input_path)
i_in = ffmpeg.input("pipe:", format="image2pipe", 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"})
job = ffmpeg.output(*output_args, **output_kwargs).overwrite_output()
try:
job_run = job.run_async(pipe_stdin=True)
for image_path in input_image_paths:
with open(image_path, "rb") as f:
image_bytes = f.read()
job_run.stdin.write(image_bytes)
job_run.stdin.close()
job_run.wait()
except:
logger.error("ffmpeg fail, job commandline:" + str(job.compile()))
def sort_by_hist(input_path):
print("Sorting by histogram similarity...")
img_list = HistSsimSubprocessor(
Path_utils.get_image_paths(input_path)).process()
return img_list
def save(self):
self.options['batch_size'] = self.batch_size
self.options['paddle'] = self.ping_pong_options.paddle
summary_path = self.get_strpath_storage_for_file('summary.txt')
Path(summary_path).write_text(self.model_summary_text)
self.onSave()
model_data = {
'iter': self.iter,
'options': self.options,
'loss_history': self.loss_history,
'sample_for_preview': self.sample_for_preview
}
self.model_data_path.write_bytes(pickle.dumps(model_data))
bckp_filename_list = [
self.get_strpath_storage_for_file(filename)
for _, filename in self.get_model_filename_list()
]
bckp_filename_list += [str(summary_path), str(self.model_data_path)]
if self.autobackup:
current_hour = time.localtime().tm_hour
if self.autobackup_current_hour != current_hour:
self.autobackup_current_hour = current_hour
for i in range(15, 0, -1):
idx_str = '%.2d' % i
next_idx_str = '%.2d' % (i + 1)
idx_backup_path = self.autobackups_path / idx_str
next_idx_packup_path = self.autobackups_path / next_idx_str
if idx_backup_path.exists():
if i == 15:
Path_utils.delete_all_files(idx_backup_path)
else:
next_idx_packup_path.mkdir(exist_ok=True)
Path_utils.move_all_files(idx_backup_path,
next_idx_packup_path)
if i == 1:
idx_backup_path.mkdir(exist_ok=True)
for filename in bckp_filename_list:
shutil.copy(
str(filename),
str(idx_backup_path / Path(filename).name))
previews = self.get_previews()
plist = []
for i in range(len(previews)):
name, bgr = previews[i]
plist += [(bgr, idx_backup_path /
(('preview_%s.jpg') % (name)))]
for preview, filepath in plist:
preview_lh = ModelBase.get_loss_history_preview(
self.loss_history, self.iter, self.batch_size,
preview.shape[1], preview.shape[2])
img = (
np.concatenate([preview_lh, preview], axis=0) *
255).astype(np.uint8)
cv2_imwrite(filepath, img)
