def pack(samples_path):
samples_dat_path = samples_path / packed_faceset_filename
if samples_dat_path.exists():
io.log_info(f"{samples_dat_path} : file already exists !")
io.input("Press enter to continue and overwrite.")
as_person_faceset = False
dir_names = pathex.get_all_dir_names(samples_path)
if len(dir_names) != 0:
as_person_faceset = io.input_bool(
f"{len(dir_names)} subdirectories found, process as person faceset?",
True)
if as_person_faceset:
image_paths = []
for dir_name in dir_names:
image_paths += pathex.get_image_paths(samples_path / dir_name)
else:
image_paths = pathex.get_image_paths(samples_path)
samples = samplelib.SampleHost.load_face_samples(image_paths)
samples_len = len(samples)
samples_configs = []
for sample in io.progress_bar_generator(samples, "Processing"):
sample_filepath = Path(sample.filename)
sample.filename = sample_filepath.name
if as_person_faceset:
sample.person_name = sample_filepath.parent.name
samples_configs.append(sample.get_config())
samples_bytes = pickle.dumps(samples_configs, 4)
of = open(samples_dat_path, "wb")
of.write(struct.pack("Q", PackedFaceset.VERSION))
of.write(struct.pack("Q", len(samples_bytes)))
of.write(samples_bytes)
del samples_bytes #just free mem
del samples_configs
sample_data_table_offset = of.tell()
of.write(bytes(8 * (samples_len + 1))) #sample data offset table
data_start_offset = of.tell()
offsets = []
for sample in io.progress_bar_generator(samples, "Packing"):
try:
if sample.person_name is not None:
sample_path = samples_path / sample.person_name / sample.filename
else:
sample_path = samples_path / sample.filename
with open(sample_path, "rb") as f:
b = f.read()
offsets.append(of.tell() - data_start_offset)
of.write(b)
except:
raise Exception(f"error while processing sample {sample_path}")
offsets.append(of.tell())
of.seek(sample_data_table_offset, 0)
for offset in offsets:
of.write(struct.pack("Q", offset))
of.seek(0, 2)
of.close()
for filename in io.progress_bar_generator(image_paths,
"Deleting files"):
Path(filename).unlink()
if as_person_faceset:
for dir_name in io.progress_bar_generator(dir_names,
"Deleting dirs"):
dir_path = samples_path / dir_name
try:
shutil.rmtree(dir_path)
except:
io.log_info(f"unable to remove: {dir_path} ")
def main(model_class_name=None,
saved_models_path=None,
training_data_src_path=None,
force_model_name=None,
input_path=None,
output_path=None,
output_mask_path=None,
aligned_path=None,
force_gpu_idxs=None,
cpu_only=None):
io.log_info("Running merger.\r\n")
try:
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 output_mask_path.exists():
output_mask_path.mkdir(parents=True, exist_ok=True)
if not saved_models_path.exists():
io.log_err('Model directory not found. Please ensure it exists.')
return
# Initialize model
import models
model = models.import_model(model_class_name)(
is_training=False,
saved_models_path=saved_models_path,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only)
predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()
# Preparing MP functions
predictor_func = MPFunc(predictor_func)
run_on_cpu = len(nn.getCurrentDeviceConfig().devices) == 0
fanseg_full_face_256_extract_func = MPClassFuncOnDemand(
TernausNet,
'extract',
name=f'FANSeg_{FaceType.toString(FaceType.FULL)}',
resolution=256,
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
xseg_256_extract_func = MPClassFuncOnDemand(
XSegNet,
'extract',
name='XSeg',
resolution=256,
weights_file_root=saved_models_path,
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
face_enhancer_func = MPClassFuncOnDemand(FaceEnhancer,
'enhance',
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
is_interactive = io.input_bool("Use interactive merger?",
True) if not io.is_colab() else False
if not is_interactive:
cfg.ask_settings()
input_path_image_paths = pathex.get_image_paths(input_path)
if cfg.type == MergerConfig.TYPE_MASKED:
if not aligned_path.exists():
io.log_err(
'Aligned directory not found. Please ensure it exists.')
return
packed_samples = None
try:
packed_samples = samplelib.PackedFaceset.load(aligned_path)
except:
io.log_err(
f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}"
)
if packed_samples is not None:
io.log_info("Using packed faceset.")
def generator():
for sample in io.progress_bar_generator(
packed_samples, "Collecting alignments"):
filepath = Path(sample.filename)
yield filepath, DFLIMG.load(
filepath,
loader_func=lambda x: sample.read_raw_file())
else:
def generator():
for filepath in io.progress_bar_generator(
pathex.get_image_paths(aligned_path),
"Collecting alignments"):
filepath = Path(filepath)
yield filepath, DFLIMG.load(filepath)
alignments = {}
multiple_faces_detected = False
for filepath, dflimg in generator():
if dflimg is None or not dflimg.has_data():
io.log_err(f"{filepath.name} is not a dfl image file")
continue
source_filename = dflimg.get_source_filename()
if source_filename is None:
continue
source_filepath = Path(source_filename)
source_filename_stem = source_filepath.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(), filepath, source_filepath))
if len(alignments_ar) > 1:
multiple_faces_detected = True
if multiple_faces_detected:
io.log_info("")
io.log_info(
"Warning: multiple faces detected. Only one alignment file should refer one source file."
)
io.log_info("")
for a_key in list(alignments.keys()):
a_ar = alignments[a_key]
if len(a_ar) > 1:
for _, filepath, source_filepath in a_ar:
io.log_info(
f"alignment {filepath.name} refers to {source_filepath.name} "
)
io.log_info("")
alignments[a_key] = [a[0] for a in a_ar]
if multiple_faces_detected:
io.log_info(
"It is strongly recommended to process the faces separatelly."
)
io.log_info(
"Use 'recover original filename' to determine the exact duplicates."
)
io.log_info("")
frames = [
InteractiveMergerSubprocessor.Frame(frame_info=FrameInfo(
filepath=Path(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."
)
io.log_info("")
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
if len(frames) == 0:
io.log_info("No frames to merge in input_dir.")
else:
if False:
pass
else:
InteractiveMergerSubprocessor(
is_interactive=is_interactive,
merger_session_filepath=model.get_strpath_storage_for_file(
'merger_session.dat'),
predictor_func=predictor_func,
predictor_input_shape=predictor_input_shape,
face_enhancer_func=face_enhancer_func,
fanseg_full_face_256_extract_func=
fanseg_full_face_256_extract_func,
xseg_256_extract_func=xseg_256_extract_func,
merger_config=cfg,
frames=frames,
frames_root_path=input_path,
output_path=output_path,
output_mask_path=output_mask_path,
model_iter=model.get_iter()).run()
model.finalize()
except Exception as e:
print(traceback.format_exc())
def main(detector=None,
input_path=None,
output_path=None,
output_debug=None,
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
face_type='full_face',
max_faces_from_image=0,
cpu_only = False,
force_gpu_idxs = None,
):
face_type = FaceType.fromString(face_type)
image_size = 512 if face_type == FaceType.WHOLE_FACE else 256
if not input_path.exists():
io.log_err ('Input directory not found. Please ensure it exists.')
return
if detector is None:
io.log_info ("Choose detector type.")
io.log_info ("[0] S3FD")
io.log_info ("[1] manual")
detector = {0:'s3fd', 1:'manual'}[ io.input_int("", 0, [0,1]) ]
device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=detector=='manual', suggest_all_gpu=True) ) \
if not cpu_only else nn.DeviceConfig.CPU()
output_debug_path = output_path.parent / (output_path.name + '_debug')
if output_debug is None:
output_debug = io.input_bool (f"Write debug images to {output_debug_path.name}?", False)
if output_path.exists():
if not manual_output_debug_fix and input_path != output_path:
output_images_paths = pathex.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
input_path_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
if manual_output_debug_fix:
if not output_debug_path.exists():
io.log_err(f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.')
return
else:
detector = 'manual'
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, pathex.get_image_paths(output_debug_path) ).run()
input_path_image_paths = sorted (input_path_image_paths)
io.log_info('Found %d images.' % (len(input_path_image_paths)))
else:
if output_debug_path.exists():
for filename in pathex.get_image_paths(output_debug_path):
Path(filename).unlink()
else:
output_debug_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info ('Performing manual extract...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ], 'landmarks-manual', image_size, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
io.log_info ('Performing 3rd pass...')
data = ExtractSubprocessor (data, 'final', image_size, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run()
else:
io.log_info ('Extracting faces...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ],
'all',
image_size,
face_type,
output_debug_path if output_debug else None,
max_faces_from_image=max_faces_from_image,
final_output_path=output_path,
device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all ( np.array ( [ d.faces_detected > 0 for d in data] ) == True ):
io.log_info ('All faces are detected, manual fix not needed.')
else:
fix_data = [ ExtractSubprocessor.Data(d.filepath) for d in data if d.faces_detected == 0 ]
io.log_info ('Performing manual fix for %d images...' % (len(fix_data)) )
fix_data = ExtractSubprocessor (fix_data, 'landmarks-manual', image_size, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
fix_data = ExtractSubprocessor (fix_data, 'final', image_size, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info ('-------------------------')
io.log_info ('Images found: %d' % (images_found) )
io.log_info ('Faces detected: %d' % (faces_detected) )
io.log_info ('-------------------------')
def main(
detector=None,
input_path=None,
output_path=None,
output_debug=None,
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
face_type='full_face',
max_faces_from_image=None,
image_size=None,
jpeg_quality=None,
cpu_only=False,
force_gpu_idxs=None,
):
if not input_path.exists():
io.log_err('Input directory not found. Please ensure it exists.')
return
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
if face_type is not None:
face_type = FaceType.fromString(face_type)
if face_type is None:
if manual_output_debug_fix:
files = pathex.get_image_paths(output_path)
if len(files) != 0:
dflimg = DFLIMG.load(Path(files[0]))
if dflimg is not None and dflimg.has_data():
face_type = FaceType.fromString(dflimg.get_face_type())
input_image_paths = pathex.get_image_unique_filestem_paths(
input_path, verbose_print_func=io.log_info)
output_images_paths = pathex.get_image_paths(output_path)
output_debug_path = output_path.parent / (output_path.name + '_debug')
continue_extraction = False
if not manual_output_debug_fix and len(output_images_paths) > 0:
if len(output_images_paths) > 128:
continue_extraction = io.input_bool(
"Continue extraction?",
True,
help_message=
"Extraction can be continued, but you must specify the same options again."
)
if len(output_images_paths) > 128 and continue_extraction:
try:
input_image_paths = input_image_paths[
[Path(x).stem for x in input_image_paths].
index(Path(output_images_paths[-128]).stem.split('_')[0]):]
except:
io.log_err(
"Error in fetching the last index. Extraction cannot be continued."
)
return
elif input_path != output_path:
io.input(
f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n"
)
for filename in output_images_paths:
Path(filename).unlink()
device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=detector=='manual', suggest_all_gpu=True) ) \
if not cpu_only else nn.DeviceConfig.CPU()
if face_type is None:
face_type = io.input_str(
"Face type",
'wf', ['f', 'wf', 'head'],
help_message=
"Full face / whole face / head. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset."
).lower()
face_type = {
'f': FaceType.FULL,
'wf': FaceType.WHOLE_FACE,
'head': FaceType.HEAD
}[face_type]
if max_faces_from_image is None:
max_faces_from_image = io.input_int(
f"Max number of faces from image",
0,
help_message=
"If you extract a src faceset that has frames with a large number of faces, it is advisable to set max faces to 3 to speed up extraction. 0 - unlimited"
)
if image_size is None:
image_size = io.input_int(
f"Image size",
512 if face_type < FaceType.HEAD else 768,
valid_range=[256, 2048],
help_message=
"Output image size. The higher image size, the worse face-enhancer works. Use higher than 512 value only if the source image is sharp enough and the face does not need to be enhanced."
)
if jpeg_quality is None:
jpeg_quality = io.input_int(
f"Jpeg quality",
90,
valid_range=[1, 100],
help_message=
"Jpeg quality. The higher jpeg quality the larger the output file size."
)
if detector is None:
io.log_info("Choose detector type.")
io.log_info("[0] S3FD")
io.log_info("[1] manual")
detector = {0: 's3fd', 1: 'manual'}[io.input_int("", 0, [0, 1])]
if output_debug is None:
output_debug = io.input_bool(
f"Write debug images to {output_debug_path.name}?", False)
if output_debug:
output_debug_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
if not output_debug_path.exists():
io.log_err(
f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.'
)
return
else:
detector = 'manual'
io.log_info(
'Performing re-extract frames which were deleted from _debug directory.'
)
input_image_paths = DeletedFilesSearcherSubprocessor(
input_image_paths,
pathex.get_image_paths(output_debug_path)).run()
input_image_paths = sorted(input_image_paths)
io.log_info('Found %d images.' % (len(input_image_paths)))
else:
if not continue_extraction and output_debug_path.exists():
for filename in pathex.get_image_paths(output_debug_path):
Path(filename).unlink()
images_found = len(input_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info('Performing manual extract...')
data = ExtractSubprocessor(
[
ExtractSubprocessor.Data(Path(filename))
for filename in input_image_paths
],
'landmarks-manual',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
manual_window_size=manual_window_size,
device_config=device_config).run()
io.log_info('Performing 3rd pass...')
data = ExtractSubprocessor(
data,
'final',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
final_output_path=output_path,
device_config=device_config).run()
else:
io.log_info('Extracting faces...')
data = ExtractSubprocessor(
[
ExtractSubprocessor.Data(Path(filename))
for filename in input_image_paths
],
'all',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
max_faces_from_image=max_faces_from_image,
final_output_path=output_path,
device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all(np.array([d.faces_detected > 0 for d in data]) == True):
io.log_info('All faces are detected, manual fix not needed.')
else:
fix_data = [
ExtractSubprocessor.Data(d.filepath) for d in data
if d.faces_detected == 0
]
io.log_info('Performing manual fix for %d images...' %
(len(fix_data)))
fix_data = ExtractSubprocessor(
fix_data,
'landmarks-manual',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
manual_window_size=manual_window_size,
device_config=device_config).run()
fix_data = ExtractSubprocessor(
fix_data,
'final',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
final_output_path=output_path,
device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info('-------------------------')
io.log_info('Images found: %d' % (images_found))
io.log_info('Faces detected: %d' % (faces_detected))
io.log_info('-------------------------')
def dev_test_68(input_dir ):
# process 68 landmarks dataset with .pts files
input_path = Path(input_dir)
if not input_path.exists():
raise ValueError('input_dir not found. Please ensure it exists.')
output_path = input_path.parent / (input_path.name+'_aligned')
io.log_info(f'Output dir is % {output_path}')
if output_path.exists():
output_images_paths = pathex.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool("WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False )
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
images_paths = pathex.get_image_paths(input_path)
for filepath in io.progress_bar_generator(images_paths, "Processing"):
filepath = Path(filepath)
pts_filepath = filepath.parent / (filepath.stem+'.pts')
if pts_filepath.exists():
pts = pts_filepath.read_text()
pts_lines = pts.split('\n')
lmrk_lines = None
for pts_line in pts_lines:
if pts_line == '{':
lmrk_lines = []
elif pts_line == '}':
break
else:
if lmrk_lines is not None:
lmrk_lines.append (pts_line)
if lmrk_lines is not None and len(lmrk_lines) == 68:
try:
lmrks = [ np.array ( lmrk_line.strip().split(' ') ).astype(np.float32).tolist() for lmrk_line in lmrk_lines]
except Exception as e:
print(e)
print(filepath)
continue
rect = LandmarksProcessor.get_rect_from_landmarks(lmrks)
output_filepath = output_path / (filepath.stem+'.jpg')
img = cv2_imread(filepath)
img = imagelib.normalize_channels(img, 3)
cv2_imwrite(output_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 95] )
raise Exception("unimplemented")
#DFLJPG.x(output_filepath, face_type=FaceType.toString(FaceType.MARK_ONLY),
# landmarks=lmrks,
# source_filename=filepath.name,
# source_rect=rect,
# source_landmarks=lmrks
# )
io.log_info("Done.")
def 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( pathex.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 sort_by_absdiff(input_path):
io.log_info ("Sorting by absolute difference...")
is_sim = io.input_bool ("Sort by similar?", True, help_message="Otherwise sort by dissimilar.")
from core.leras import nn
device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
nn.initialize( device_config=device_config, data_format="NHWC" )
tf = nn.tf
image_paths = pathex.get_image_paths(input_path)
image_paths_len = len(image_paths)
batch_size = 512
batch_size_remain = image_paths_len % batch_size
i_t = tf.placeholder (tf.float32, (None,None,None,None) )
j_t = tf.placeholder (tf.float32, (None,None,None,None) )
outputs_full = []
outputs_remain = []
for i in range(batch_size):
diff_t = tf.reduce_sum( tf.abs(i_t-j_t[i]), axis=[1,2,3] )
outputs_full.append(diff_t)
if i < batch_size_remain:
outputs_remain.append(diff_t)
def func_bs_full(i,j):
return nn.tf_sess.run (outputs_full, feed_dict={i_t:i,j_t:j})
def func_bs_remain(i,j):
return nn.tf_sess.run (outputs_remain, feed_dict={i_t:i,j_t:j})
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 main(model_class_name=None,
saved_models_path=None,
training_data_src_path=None,
force_model_name=None,
input_path=None,
output_path=None,
output_mask_path=None,
aligned_path=None,
force_gpu_idxs=None,
cpu_only=None):
io.log_info("Running merger.\r\n")
try:
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 output_mask_path.exists():
output_mask_path.mkdir(parents=True, exist_ok=True)
if not saved_models_path.exists():
io.log_err('Model directory not found. Please ensure it exists.')
return
is_interactive = io.input_bool("Use interactive merger?",
True) if not io.is_colab() else False
import models
model = models.import_model(model_class_name)(
is_training=False,
saved_models_path=saved_models_path,
training_data_src_path=training_data_src_path,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only)
merger_session_filepath = model.get_strpath_storage_for_file(
'merger_session.dat')
predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()
if not is_interactive:
cfg.ask_settings()
input_path_image_paths = pathex.get_image_paths(input_path)
if cfg.type == MergerConfig.TYPE_MASKED:
if not aligned_path.exists():
io.log_err(
'Aligned directory not found. Please ensure it exists.')
return
packed_samples = None
try:
packed_samples = samplelib.PackedFaceset.load(aligned_path)
except:
io.log_err(
f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}"
)
if packed_samples is not None:
io.log_info("Using packed faceset.")
def generator():
for sample in io.progress_bar_generator(
packed_samples, "Collecting alignments"):
filepath = Path(sample.filename)
yield filepath, DFLIMG.load(
filepath,
loader_func=lambda x: sample.read_raw_file())
else:
def generator():
for filepath in io.progress_bar_generator(
pathex.get_image_paths(aligned_path),
"Collecting alignments"):
filepath = Path(filepath)
yield filepath, DFLIMG.load(filepath)
alignments = {}
multiple_faces_detected = False
for filepath, dflimg in generator():
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:
continue
source_filepath = Path(source_filename)
source_filename_stem = source_filepath.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(), filepath, source_filepath))
if len(alignments_ar) > 1:
multiple_faces_detected = True
if multiple_faces_detected:
io.log_info("")
io.log_info(
"Warning: multiple faces detected. Only one alignment file should refer one source file."
)
io.log_info("")
for a_key in list(alignments.keys()):
a_ar = alignments[a_key]
if len(a_ar) > 1:
for _, filepath, source_filepath in a_ar:
io.log_info(
f"alignment {filepath.name} refers to {source_filepath.name} "
)
io.log_info("")
alignments[a_key] = [a[0] for a in a_ar]
if multiple_faces_detected:
io.log_info(
"It is strongly recommended to process the faces separatelly."
)
io.log_info(
"Use 'recover original filename' to determine the exact duplicates."
)
io.log_info("")
filesdata = []
for filepath in io.progress_bar_generator(input_path_image_paths,
"Collecting info"):
filepath = Path(filepath)
filesdata += [
FrameInfo(filepath=filepath,
landmarks_list=alignments.get(
filepath.stem, None))
]
frames = []
filesdata_len = len(filesdata)
for i in range(len(filesdata)):
frame_info = filesdata[i]
if multiple_faces_detected:
prev_temporal_frame_infos = None
next_temporal_frame_infos = None
else:
prev_temporal_frame_infos = []
next_temporal_frame_infos = []
for t in range(1, 6):
prev_frame_info = filesdata[max(i - t, 0)]
next_frame_info = filesdata[min(
i + t, filesdata_len - 1)]
prev_temporal_frame_infos.insert(0, prev_frame_info)
next_temporal_frame_infos.append(next_frame_info)
frames.append(
MergeSubprocessor.Frame(
prev_temporal_frame_infos=prev_temporal_frame_infos,
frame_info=frame_info,
next_temporal_frame_infos=next_temporal_frame_infos))
if multiple_faces_detected:
io.log_info(
"Warning: multiple faces detected. Motion blur will not be used."
)
io.log_info("")
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 == MergerConfig.TYPE_FACE_AVATAR:
pass
"""
filesdata = []
for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"):
filepath = Path(filepath)
dflimg = DFLIMG.load(filepath)
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
filesdata += [ ( FrameInfo(filepath=filepath, landmarks_list=[dflimg.get_landmarks()] ), dflimg.get_source_filename() ) ]
filesdata = sorted(filesdata, key=operator.itemgetter(1)) #sort by source_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 ( MergeSubprocessor.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 merge in input_dir.")
else:
MergeSubprocessor(is_interactive=is_interactive,
merger_session_filepath=merger_session_filepath,
predictor_func=predictor_func,
predictor_input_shape=predictor_input_shape,
merger_config=cfg,
frames=frames,
frames_root_path=input_path,
output_path=output_path,
output_mask_path=output_mask_path,
model_iter=model.get_iter()).run()
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
def extract_vggface2_dataset(input_dir, device_args={}):
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
input_path = Path(input_dir)
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
bb_csv = input_path / 'loose_bb_train.csv'
if not bb_csv.exists():
raise ValueError('loose_bb_train.csv found. Please ensure it exists.')
bb_lines = bb_csv.read_text().split('\n')
bb_lines.pop(0)
bb_dict = {}
for line in bb_lines:
name, l, t, w, h = line.split(',')
name = name[1:-1]
l, t, w, h = [int(x) for x in (l, t, w, h)]
bb_dict[name] = (l, t, w, h)
output_path = input_path.parent / (input_path.name + '_out')
dir_names = pathex.get_all_dir_names(input_path)
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
data = []
for dir_name in io.progress_bar_generator(dir_names, "Collecting"):
cur_input_path = input_path / dir_name
cur_output_path = output_path / dir_name
if not cur_output_path.exists():
cur_output_path.mkdir(parents=True, exist_ok=True)
input_path_image_paths = pathex.get_image_paths(cur_input_path)
for filename in input_path_image_paths:
filename_path = Path(filename)
name = filename_path.parent.name + '/' + filename_path.stem
if name not in bb_dict:
continue
l, t, w, h = bb_dict[name]
if min(w, h) < 128:
continue
data += [
ExtractSubprocessor.Data(filename=filename,
rects=[(l, t, l + w, t + h)],
landmarks_accurate=False,
force_output_path=cur_output_path)
]
face_type = FaceType.fromString('full_face')
io.log_info('Performing 2nd pass...')
data = ExtractSubprocessor(data,
'landmarks',
256,
face_type,
debug_dir=None,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False).run()
io.log_info('Performing 3rd pass...')
ExtractSubprocessor(data,
'final',
256,
face_type,
debug_dir=None,
multi_gpu=multi_gpu,
cpu_only=cpu_only,
manual=False,
final_output_path=None).run()
def main(model_class_name=None,
saved_models_path=None,
training_data_src_path=None,
force_model_name=None,
input_path=None,
output_path=None,
output_mask_path=None,
aligned_path=None,
force_gpu_idxs=None,
cpu_only=None):
io.log_info("启动合成程序.\r\n")
try:
if not input_path.exists():
io.log_err('找不到输入文件')
return
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
if not output_mask_path.exists():
output_mask_path.mkdir(parents=True, exist_ok=True)
if not saved_models_path.exists():
io.log_err('找不到模型文件夹')
return
# Initialize model
import models
model = models.import_model(model_class_name)(
is_training=False,
saved_models_path=saved_models_path,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only)
predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()
# Preparing MP functions
predictor_func = MPFunc(predictor_func)
run_on_cpu = len(nn.getCurrentDeviceConfig().devices) == 0
xseg_256_extract_func = MPClassFuncOnDemand(
XSegNet,
'extract',
name='XSeg',
resolution=256,
weights_file_root=saved_models_path,
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
face_enhancer_func = MPClassFuncOnDemand(FaceEnhancer,
'enhance',
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
is_interactive = io.input_bool("是否启用交互式合成?",
True) if not io.is_colab() else False
if not is_interactive:
cfg.ask_settings()
subprocess_count = io.input_int(
"线程数量?",
max(8, multiprocessing.cpu_count()),
valid_range=[1, multiprocessing.cpu_count()],
help_message=
"Specify the number of threads to process. A low value may affect performance. A high value may result in memory error. The value may not be greater than CPU cores."
)
input_path_image_paths = pathex.get_image_paths(input_path)
if cfg.type == MergerConfig.TYPE_MASKED:
if not aligned_path.exists():
io.log_err('头像文件不存在')
return
packed_samples = None
try:
packed_samples = samplelib.PackedFaceset.load(aligned_path)
except:
io.log_err(
f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}"
)
if packed_samples is not None:
io.log_info("正在使用打包数据集.")
def generator():
for sample in io.progress_bar_generator(
packed_samples, "收集对齐头像"):
filepath = Path(sample.filename)
yield filepath, DFLIMG.load(
filepath,
loader_func=lambda x: sample.read_raw_file())
else:
def generator():
for filepath in io.progress_bar_generator(
pathex.get_image_paths(aligned_path), "收集对齐头像"):
filepath = Path(filepath)
yield filepath, DFLIMG.load(filepath)
alignments = {}
multiple_faces_detected = False
for filepath, dflimg in generator():
if dflimg is None or not dflimg.has_data():
io.log_err(f"{filepath.name} is not a dfl image file")
continue
source_filename = dflimg.get_source_filename()
if source_filename is None:
continue
source_filepath = Path(source_filename)
source_filename_stem = source_filepath.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(), filepath, source_filepath))
if len(alignments_ar) > 1:
multiple_faces_detected = True
if multiple_faces_detected:
io.log_info("")
io.log_info("警告: 检测到多个人脸. 一个对齐头像应该对应一个源文件.")
io.log_info("")
for a_key in list(alignments.keys()):
a_ar = alignments[a_key]
if len(a_ar) > 1:
for _, filepath, source_filepath in a_ar:
io.log_info(
f"alignment {filepath.name} refers to {source_filepath.name} "
)
io.log_info("")
alignments[a_key] = [a[0] for a in a_ar]
if multiple_faces_detected:
io.log_info("强烈建议分开处理人脸.")
io.log_info("使用 'recover original filename' 确定提取的唯一性.")
io.log_info("")
frames = [
InteractiveMergerSubprocessor.Frame(frame_info=FrameInfo(
filepath=Path(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."
)
io.log_info("")
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)),
"计算运动矢量"):
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
if len(frames) == 0:
io.log_info("没有可用图片")
else:
if False:
pass
else:
InteractiveMergerSubprocessor(
is_interactive=is_interactive,
merger_session_filepath=model.get_strpath_storage_for_file(
'merger_session.dat'),
predictor_func=predictor_func,
predictor_input_shape=predictor_input_shape,
face_enhancer_func=face_enhancer_func,
xseg_256_extract_func=xseg_256_extract_func,
merger_config=cfg,
frames=frames,
frames_root_path=input_path,
output_path=output_path,
output_mask_path=output_mask_path,
model_iter=model.get_iter(),
subprocess_count=subprocess_count,
).run()
model.finalize()
except Exception as e:
print(traceback.format_exc())
def __init__(self, is_interactive, merger_session_filepath, predictor_func,
predictor_input_shape, merger_config, frames,
frames_root_path, output_path, output_mask_path, model_iter):
if len(frames) == 0:
raise ValueError("len (frames) == 0")
super().__init__('Merger',
MergeSubprocessor.Cli,
86400 if MERGER_DEBUG else 60,
io_loop_sleep_time=0.001)
self.is_interactive = is_interactive
self.merger_session_filepath = Path(merger_session_filepath)
self.merger_config = merger_config
self.predictor_func_host, self.predictor_func = SubprocessFunctionCaller.make_pair(
predictor_func)
self.predictor_input_shape = predictor_input_shape
self.face_enhancer = None
def superres_func(face_bgr):
if self.face_enhancer is None:
self.face_enhancer = FaceEnhancer(place_model_on_cpu=True)
return self.face_enhancer.enhance(face_bgr,
is_tanh=True,
preserve_size=False)
self.superres_host, self.superres_func = SubprocessFunctionCaller.make_pair(
superres_func)
self.fanseg_by_face_type = {}
self.fanseg_input_size = 256
def fanseg_extract_func(face_type, *args, **kwargs):
fanseg = self.fanseg_by_face_type.get(face_type, None)
if self.fanseg_by_face_type.get(face_type, None) is None:
cpu_only = len(nn.getCurrentDeviceConfig().devices) == 0
with nn.tf.device('/CPU:0' if cpu_only else '/GPU:0'):
fanseg = TernausNet("FANSeg",
self.fanseg_input_size,
FaceType.toString(face_type),
place_model_on_cpu=True)
self.fanseg_by_face_type[face_type] = fanseg
return fanseg.extract(*args, **kwargs)
self.fanseg_host, self.fanseg_extract_func = SubprocessFunctionCaller.make_pair(
fanseg_extract_func)
self.frames_root_path = frames_root_path
self.output_path = output_path
self.output_mask_path = output_mask_path
self.model_iter = model_iter
self.prefetch_frame_count = self.process_count = multiprocessing.cpu_count(
)
session_data = None
if self.is_interactive and self.merger_session_filepath.exists():
io.input_skip_pending()
if io.input_bool("Use saved session?", True):
try:
with open(str(self.merger_session_filepath), "rb") as f:
session_data = pickle.loads(f.read())
except Exception as e:
pass
rewind_to_frame_idx = None
self.frames = frames
self.frames_idxs = [*range(len(self.frames))]
self.frames_done_idxs = []
if self.is_interactive and session_data is not None:
# Loaded session data, check it
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)) # frames count must match
if frames_equal:
for i in range(len(frames)):
frame = frames[i]
s_frame = s_frames[i]
# frames filenames must match
if frame.frame_info.filepath.name != s_frame.frame_info.filepath.name:
frames_equal = False
if not frames_equal:
break
if frames_equal:
io.log_info('Using saved session from ' +
'/'.join(self.merger_session_filepath.parts[-2:]))
for frame in s_frames:
if frame.cfg is not None:
# recreate MergerConfig class using constructor with get_config() as dict params
# so if any new param will be added, old merger 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 was more trained, recompute all frames
rewind_to_frame_idx = -1
for frame in self.frames:
frame.is_done = False
elif len(self.frames_idxs) == 0:
# all frames are done?
rewind_to_frame_idx = -1
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 pathex.get_image_paths(
self.output_path): #remove all images in output_path
Path(filename).unlink()
for filename in pathex.get_image_paths(
self.output_mask_path
): #remove all images in output_mask_path
Path(filename).unlink()
frames[0].cfg = self.merger_config.copy()
for i in range(len(self.frames)):
frame = self.frames[i]
frame.idx = i
frame.output_filepath = self.output_path / (
frame.frame_info.filepath.stem + '.png')
frame.output_mask_filepath = self.output_mask_path / (
frame.frame_info.filepath.stem + '.png')
if not frame.output_filepath.exists() or \
not frame.output_mask_filepath.exists():
# if some frame does not exist, recompute and rewind
frame.is_done = False
frame.is_shown = False
if rewind_to_frame_idx is None:
rewind_to_frame_idx = i - 1
else:
rewind_to_frame_idx = min(rewind_to_frame_idx, i - 1)
if rewind_to_frame_idx is not None:
while len(self.frames_done_idxs) > 0:
if self.frames_done_idxs[-1] > rewind_to_frame_idx:
prev_frame = self.frames[self.frames_done_idxs.pop()]
self.frames_idxs.insert(0, prev_frame.idx)
else:
break
def generator():
for filepath in io.progress_bar_generator(
pathex.get_image_paths(aligned_path), "收集对齐头像"):
filepath = Path(filepath)
yield filepath, DFLIMG.load(filepath)
def __init__(self,
root_path,
debug=False,
batch_size=1,
resolution=256,
face_type=None,
generators_count=4,
data_format="NHWC",
**kwargs):
super().__init__(debug, batch_size)
self.initialized = False
dataset_path = root_path / 'AvatarOperatorDataset'
if not dataset_path.exists():
raise ValueError(f'Unable to find {dataset_path}')
chains_dir_names = pathex.get_all_dir_names(dataset_path)
samples = SampleLoader.load(SampleType.FACE,
dataset_path,
subdirs=True)
sample_idx_by_path = {
sample.filename: i
for i, sample in enumerate(samples)
}
kf_idxs = []
for chain_dir_name in chains_dir_names:
chain_root_path = dataset_path / chain_dir_name
subchain_dir_names = pathex.get_all_dir_names(chain_root_path)
try:
subchain_dir_names.sort(key=int)
except:
raise Exception(
f'{chain_root_path} must contain only numerical name of directories'
)
chain_samples = []
for subchain_dir_name in subchain_dir_names:
subchain_root = chain_root_path / subchain_dir_name
subchain_samples = [ sample_idx_by_path[image_path] for image_path in pathex.get_image_paths(subchain_root) \
if image_path in sample_idx_by_path ]
if len(subchain_samples) < 3:
raise Exception(
f'subchain {subchain_dir_name} must contain at least 3 faces. If you delete this subchain, then th echain will be corrupted.'
)
chain_samples += [subchain_samples]
chain_samples_len = len(chain_samples)
for i in range(chain_samples_len - 1):
kf_idxs += [(chain_samples[i + 1][0], chain_samples[i][-1],
chain_samples[i][:-1])]
for i in range(1, chain_samples_len):
kf_idxs += [(chain_samples[i - 1][-1], chain_samples[i][0],
chain_samples[i][1:])]
if self.debug:
self.generators_count = 1
else:
self.generators_count = max(1, generators_count)
if self.debug:
self.generators = [
ThisThreadGenerator(
self.batch_func,
(samples, kf_idxs, resolution, face_type, data_format))
]
else:
self.generators = [SubprocessGenerator ( self.batch_func, (samples, kf_idxs, resolution, face_type, data_format), start_now=False ) \
for i in range(self.generators_count) ]
SubprocessGenerator.start_in_parallel(self.generators)
self.generator_counter = -1
self.initialized = True
def extract_video(self, input_file, output_dir, output_ext=None, fps=None):
input_file_path = Path(input_file)
output_path = Path(output_dir)
if not output_path.exists():
output_path.mkdir(exist_ok=True)
# self.g_progress_info.append("\n视频帧输出目录: " + str(Path(output_path).absolute()))
# self.NoteTipsEdit.append("\n视频帧输出目录: " + str(Path(output_path).absolute()))
InfoNotifier.InfoNotifier.g_progress_info.append(
"\n视频帧输出目录: " + str(Path(output_path).absolute()))
if input_file_path.suffix == '.*':
input_file_path = pathex.get_first_file_by_stem(
input_file_path.parent, input_file_path.stem)
else:
if not input_file_path.exists():
input_file_path = None
# self.g_progress_info.append("\n视频输入路径:" + str(input_file_path))
# self.NoteTipsEdit.append("\n视频输入路径:" + str(input_file_path))
InfoNotifier.InfoNotifier.g_progress_info.append("\n视频输入路径:" +
str(input_file_path))
if input_file_path is None:
io.log_err("input_file not found.")
# self.g_progress_info.append("\n视频输入路径不存在")
# self.NoteTipsEdit.append("\n视频输入路径不存在")
InfoNotifier.InfoNotifier.g_progress_info.append("\n视频输入路径不存在")
return
# if fps is None:
# fps = io.input_int("Enter FPS", 0,
# help_message="How many frames of every second of the video will be extracted. 0 - full fps")
# self.NoteTipsEdit.append("\n读取不到帧")
# self.g_progress_info.append("\n视频帧抽取频率: full fps")
# self.NoteTipsEdit.append("\n视频帧抽取频率: full fps")
if output_ext is None:
output_ext = io.input_str(
"Output image format",
"png", ["png", "jpg"],
help_message=
"png is lossless, but extraction is x10 slower for HDD, requires x10 more disk space than jpg."
)
# self.g_progress_info.append("\n视频帧输出格式频率: " + output_ext)
# self.NoteTipsEdit.append("\n视频帧输出格式频率: " + output_ext)
InfoNotifier.InfoNotifier.g_progress_info.append("\n视频帧输出格式频率: " +
output_ext)
filenames = pathex.get_image_paths(output_path, ['.' + output_ext])
if len(filenames) != 0:
# self.g_progress_info.append("\n视频帧输出目录不为空, 该目录将被清空!")
InfoNotifier.InfoNotifier.g_progress_info.append(
"\n视频帧输出目录不为空, 该目录将被清空!")
# self.NoteTipsEdit.append("\n视频帧输出目录不为空, 该目录将被清空!")
# Ui_MainWindow.setupUi(self.NoteTipsEdit.append())
for filename in filenames:
Path(filename).unlink()
QApplication.processEvents()
job = ffmpeg.input(str(input_file_path))
kwargs = {'pix_fmt': 'rgb24'}
# if fps !=0:
kwargs.update({'r': str(fps)})
if output_ext == 'jpg':
kwargs.update({'q:v': '2'}) # highest quality for jpg
job = job.output(str(output_path / ('%5d.' + output_ext)), **kwargs)
try:
job, err = job.run(cmd=self.ffmpeg_cmd_path)
except:
io.log_err("ffmpeg fail, job commandline:" + str(job.compile()))
def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default_mask=False):
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)
if not no_default_mask:
eyebrows_expand_mod = np.clip ( io.input_int ("Default eyebrows expand modifier?", 100, add_info="0..400"), 0, 400 ) / 100.0
else:
eyebrows_expand_mod = None
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 pathex.get_image_paths(input_path)]
done_paths = []
done_images_types = {}
image_paths_total = len(image_paths)
saved_ie_polys = IEPolys()
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:
dflimg = DFLIMG.load (filepath)
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
else:
lmrks = dflimg.get_landmarks()
ie_polys = IEPolys.load(dflimg.get_ie_polys())
fanseg_mask = dflimg.get_fanseg_mask()
if filepath.name in cached_images:
img = cached_images[filepath.name]
else:
img = cached_images[filepath.name] = cv2_imread(str(filepath)) / 255.0
if fanseg_mask is not None:
mask = fanseg_mask
else:
if no_default_mask:
mask = np.zeros ( (target_wh,target_wh,3) )
else:
mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks, eyebrows_expand_mod=eyebrows_expand_mod)
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',
'[r] - applies edits made to last saved image.',
'[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':
saved_ie_polys = ed.ie_polys
do_save_move_count = 1 if not shift_pressed else 10
elif chr_key == 'c':
saved_ie_polys = ed.ie_polys
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
elif chr_key == 'r' and saved_ie_polys != None:
ed.set_ie_polys(saved_ie_polys)
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(), eyebrows_expand_mod=eyebrows_expand_mod )
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(), eyebrows_expand_mod=eyebrows_expand_mod )
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 apply_celebamaskhq(input_dir):
input_path = Path(input_dir)
img_path = input_path / 'aligned'
mask_path = input_path / 'mask'
if not img_path.exists():
raise ValueError(
f'{str(img_path)} directory not found. Please ensure it exists.')
CelebAMASKHQSubprocessor(pathex.get_image_paths(img_path),
pathex.get_image_paths(mask_path,
subdirs=True)).run()
return
paths_to_extract = []
for filename in io.progress_bar_generator(pathex.get_image_paths(img_path),
desc="Processing"):
filepath = Path(filename)
dflimg = DFLIMG.load(filepath)
if dflimg is not None:
paths_to_extract.append(filepath)
image_to_face_mat = dflimg.get_image_to_face_mat()
src_filename = dflimg.get_source_filename()
#img = cv2_imread(filename)
h, w, c = dflimg.get_shape()
fanseg_mask = LandmarksProcessor.get_image_hull_mask(
(h, w, c), dflimg.get_landmarks())
idx_name = '%.5d' % int(src_filename.split('.')[0])
idx_files = [x for x in masks_files if idx_name in x]
skin_files = [x for x in idx_files if 'skin' in x]
eye_glass_files = [x for x in idx_files if 'eye_g' in x]
for files, is_invert in [(skin_files, False), (eye_glass_files, True)]:
if len(files) > 0:
mask = cv2_imread(files[0])
mask = mask[..., 0]
mask[mask == 255] = 1
mask = mask.astype(np.float32)
mask = cv2.resize(mask, (1024, 1024))
mask = cv2.warpAffine(mask, image_to_face_mat, (w, h),
cv2.INTER_LANCZOS4)
if not is_invert:
fanseg_mask *= mask[..., None]
else:
fanseg_mask *= (1 - mask[..., None])
#cv2.imshow("", (fanseg_mask*255).astype(np.uint8) )
#cv2.waitKey(0)
dflimg.embed_and_set(filename, fanseg_mask=fanseg_mask)
def sort_by_hist(input_path):
io.log_info ("Sorting by histogram similarity...")
img_list = HistSsimSubprocessor(pathex.get_image_paths(input_path)).run()
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 = pathex.get_image_paths(output_path)
if len(output_images_paths) > 0:
io.input_bool(
"WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue."
% (str(output_path)), False)
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
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']
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]])
]
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_best(input_path, faster=False):
target_count = io.input_int ("Target number of faces?", 2000)
io.log_info ("Performing sort by best faces.")
if faster:
io.log_info("Using faster algorithm. Faces will be sorted by source-rect-area instead of blur.")
img_list, trash_img_list = FinalLoaderSubprocessor( pathex.get_image_paths(input_path), faster ).run()
final_img_list = []
grads = 128
imgs_per_grad = round (target_count / grads)
#instead of math.pi / 2, using -1.2,+1.2 because actually maximum yaw for 2DFAN landmarks are -1.2+1.2
grads_space = np.linspace (-1.2, 1.2,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) > 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 (-math.pi / 2,math.pi / 2, 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 __init__(self,
is_interactive,
merger_session_filepath,
predictor_func,
predictor_input_shape,
face_enhancer_func,
xseg_256_extract_func,
merger_config,
frames,
frames_root_path,
output_path,
output_mask_path,
model_iter,
subprocess_count=4,
src_src=False):
if len(frames) == 0:
raise ValueError("len (frames) == 0")
super().__init__('Merger',
InteractiveMergerSubprocessor.Cli,
io_loop_sleep_time=0.001)
self.is_interactive = is_interactive
self.merger_session_filepath = Path(merger_session_filepath)
self.merger_config = merger_config
self.predictor_func = predictor_func
self.predictor_input_shape = predictor_input_shape
self.face_enhancer_func = face_enhancer_func
self.xseg_256_extract_func = xseg_256_extract_func
self.frames_root_path = frames_root_path
self.output_path = output_path
self.output_mask_path = output_mask_path
self.model_iter = model_iter
self.prefetch_frame_count = self.process_count = subprocess_count
self.src_src = src_src
session_data = None
if self.is_interactive and self.merger_session_filepath.exists():
io.input_skip_pending()
if io.input_bool("Use saved session?", True):
try:
with open(str(self.merger_session_filepath), "rb") as f:
session_data = pickle.loads(f.read())
except Exception as e:
pass
rewind_to_frame_idx = None
self.frames = frames
self.frames_idxs = [*range(len(self.frames))]
self.frames_done_idxs = []
if self.is_interactive and session_data is not None:
# Loaded session data, check it
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)) # frames count must match
if frames_equal:
for i in range(len(frames)):
frame = frames[i]
s_frame = s_frames[i]
# frames filenames must match
if frame.frame_info.filepath.name != s_frame.frame_info.filepath.name:
frames_equal = False
if not frames_equal:
break
if frames_equal:
io.log_info('Using saved session from ' +
'/'.join(self.merger_session_filepath.parts[-2:]))
for frame in s_frames:
if frame.cfg is not None:
# recreate MergerConfig class using constructor with get_config() as dict params
# so if any new param will be added, old merger 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 was more trained, recompute all frames
rewind_to_frame_idx = -1
for frame in self.frames:
frame.is_done = False
elif len(self.frames_idxs) == 0:
# all frames are done?
rewind_to_frame_idx = -1
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 pathex.get_image_paths(
self.output_path): #remove all images in output_path
Path(filename).unlink()
for filename in pathex.get_image_paths(
self.output_mask_path
): #remove all images in output_mask_path
Path(filename).unlink()
frames[0].cfg = self.merger_config.copy()
for i in range(len(self.frames)):
frame = self.frames[i]
frame.idx = i
frame.output_filepath = self.output_path / (
frame.frame_info.filepath.stem + '.png')
frame.output_mask_filepath = self.output_mask_path / (
frame.frame_info.filepath.stem + '.png')
if not frame.output_filepath.exists() or \
not frame.output_mask_filepath.exists():
# if some frame does not exist, recompute and rewind
frame.is_done = False
frame.is_shown = False
if rewind_to_frame_idx is None:
rewind_to_frame_idx = i - 1
else:
rewind_to_frame_idx = min(rewind_to_frame_idx, i - 1)
if rewind_to_frame_idx is not None:
while len(self.frames_done_idxs) > 0:
if self.frames_done_idxs[-1] > rewind_to_frame_idx:
prev_frame = self.frames[self.frames_done_idxs.pop()]
self.frames_idxs.insert(0, prev_frame.idx)
else:
break
def __init__(self,
is_training=False,
saved_models_path=None,
training_data_src_path=None,
training_data_dst_path=None,
pretraining_data_path=None,
pretrained_model_path=None,
no_preview=False,
force_model_name=None,
force_gpu_idxs=None,
cpu_only=False,
debug=False,
force_model_class_name=None,
silent_start=False,
**kwargs):
self.is_training = is_training
self.saved_models_path = saved_models_path
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.pretrained_model_path = pretrained_model_path
self.no_preview = no_preview
self.debug = debug
self.model_class_name = model_class_name = Path(
inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
if force_model_class_name is None:
if force_model_name is not None:
self.model_name = force_model_name
else:
while True:
# gather all model dat files
saved_models_names = []
for filepath in pathex.get_file_paths(saved_models_path):
filepath_name = filepath.name
if filepath_name.endswith(
f'{model_class_name}_data.dat'):
saved_models_names += [
(filepath_name.split('_')[0],
os.path.getmtime(filepath))
]
# sort by modified datetime
saved_models_names = sorted(saved_models_names,
key=operator.itemgetter(1),
reverse=True)
saved_models_names = [x[0] for x in saved_models_names]
if len(saved_models_names) != 0:
if silent_start:
self.model_name = saved_models_names[0]
io.log_info(
f'Silent start: choosed model "{self.model_name}"'
)
else:
io.log_info(
"Choose one of saved models, or enter a name to create a new model."
)
io.log_info("[r] : rename")
io.log_info("[d] : delete")
io.log_info("")
for i, model_name in enumerate(saved_models_names):
s = f"[{i}] : {model_name} "
if i == 0:
s += "- latest"
io.log_info(s)
inp = io.input_str(f"",
"0",
show_default_value=False)
model_idx = -1
try:
model_idx = np.clip(
int(inp), 0,
len(saved_models_names) - 1)
except:
pass
if model_idx == -1:
if len(inp) == 1:
is_rename = inp[0] == 'r'
is_delete = inp[0] == 'd'
if is_rename or is_delete:
if len(saved_models_names) != 0:
if is_rename:
name = io.input_str(
f"Enter the name of the model you want to rename"
)
elif is_delete:
name = io.input_str(
f"Enter the name of the model you want to delete"
)
if name in saved_models_names:
if is_rename:
new_model_name = io.input_str(
f"Enter new name of the model"
)
for filepath in pathex.get_paths(
saved_models_path):
filepath_name = filepath.name
model_filename, remain_filename = filepath_name.split(
'_', 1)
if model_filename == name:
if is_rename:
new_filepath = filepath.parent / (
new_model_name
+ '_' +
remain_filename
)
filepath.rename(
new_filepath)
elif is_delete:
filepath.unlink()
continue
self.model_name = inp
else:
self.model_name = saved_models_names[model_idx]
else:
self.model_name = io.input_str(
f"No saved models found. Enter a name of a new model",
"new")
self.model_name = self.model_name.replace('_', ' ')
break
self.model_name = self.model_name + '_' + self.model_class_name
else:
self.model_name = force_model_class_name
self.iter = 0
self.options = {}
self.loss_history = []
self.sample_for_preview = None
self.choosed_gpu_indexes = None
model_data = {}
self.model_data_path = Path(
self.get_strpath_storage_for_file('data.dat'))
if self.model_data_path.exists():
io.log_info(f"Loading {self.model_name} model...")
model_data = pickle.loads(self.model_data_path.read_bytes())
self.iter = model_data.get('iter', 0)
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)
self.choosed_gpu_indexes = model_data.get(
'choosed_gpu_indexes', None)
if self.is_first_run():
io.log_info("\nModel first run.")
if silent_start:
self.device_config = nn.DeviceConfig.BestGPU()
io.log_info(
f"Silent start: choosed device {'CPU' if self.device_config.cpu_only else self.device_config.devices[0].name}"
)
else:
self.device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_best_multi_gpu=True)) \
if not cpu_only else nn.DeviceConfig.CPU()
nn.initialize(self.device_config)
####
self.default_options_path = saved_models_path / f'{self.model_class_name}_default_options.dat'
self.default_options = {}
if self.default_options_path.exists():
try:
self.default_options = pickle.loads(
self.default_options_path.read_bytes())
except:
pass
self.choose_preview_history = False
self.batch_size = self.load_or_def_option(
'batch_size', int(params['suggest_batch_size']))
#####
io.input_skip_pending()
self.on_initialize_options()
if self.is_first_run():
# save as default options only for first run model initialize
self.default_options_path.write_bytes(pickle.dumps(self.options))
self.autobackup_hour = self.options.get('autobackup_hour', 0)
self.write_preview_history = self.options.get('write_preview_history',
False)
self.target_iter = self.options.get('target_iter', 0)
self.random_flip = self.options.get('random_flip', True)
self.on_initialize()
self.options['batch_size'] = self.batch_size
if self.is_training:
self.preview_history_path = self.saved_models_path / (
f'{self.get_model_name()}_history')
self.autobackups_path = self.saved_models_path / (
f'{self.get_model_name()}_autobackups')
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 pathex.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'
)
self.update_sample_for_preview(
choose_preview_history=self.choose_preview_history)
if self.autobackup_hour != 0:
self.autobackup_start_time = time.time()
if not self.autobackups_path.exists():
self.autobackups_path.mkdir(exist_ok=True)
io.log_info(self.get_summary_text())
def generator():
for filepath in io.progress_bar_generator(
pathex.get_image_paths(aligned_path),
"Collecting alignments"):
filepath = Path(filepath)
yield filepath, DFLIMG.load(filepath)
def denoise_image_sequence(input_dir, ext=None, factor=None):
input_path = Path(input_dir)
if not input_path.exists():
io.log_err("input_dir not found.")
return
image_paths = [
Path(filepath) for filepath in pathex.get_image_paths(input_path)
]
# Check extension of all images
image_paths_suffix = None
for filepath in image_paths:
if image_paths_suffix is None:
image_paths_suffix = filepath.suffix
else:
if filepath.suffix != image_paths_suffix:
io.log_err(
f"All images in {input_path.name} should be with the same extension."
)
return
if factor is None:
factor = np.clip(io.input_int("Denoise factor?", 7, add_info="1-20"),
1, 20)
# Rename to temporary filenames
for i, filepath in io.progress_bar_generator(enumerate(image_paths),
"Renaming",
leave=False):
src = filepath
dst = filepath.parent / (f'{i+1:06}_{filepath.name}')
try:
src.rename(dst)
except:
io.log_error('fail to rename %s' % (src.name))
return
# Rename to sequental filenames
for i, filepath in io.progress_bar_generator(enumerate(image_paths),
"Renaming",
leave=False):
src = filepath.parent / (f'{i+1:06}_{filepath.name}')
dst = filepath.parent / (f'{i+1:06}{filepath.suffix}')
try:
src.rename(dst)
except:
io.log_error('fail to rename %s' % (src.name))
return
# Process image sequence in ffmpeg
kwargs = {}
if image_paths_suffix == '.jpg':
kwargs.update({'q:v': '2'})
job = (ffmpeg.input(str(input_path / ('%6d' + image_paths_suffix))).filter(
"hqdn3d", factor, factor, 5,
5).output(str(input_path / ('%6d' + image_paths_suffix)), **kwargs))
try:
job = job.run()
except:
io.log_err("ffmpeg fail, job commandline:" + str(job.compile()))
# Rename to temporary filenames
for i, filepath in io.progress_bar_generator(enumerate(image_paths),
"Renaming",
leave=False):
src = filepath.parent / (f'{i+1:06}{filepath.suffix}')
dst = filepath.parent / (f'{i+1:06}_{filepath.name}')
try:
src.rename(dst)
except:
io.log_error('fail to rename %s' % (src.name))
return
# Rename to initial filenames
for i, filepath in io.progress_bar_generator(enumerate(image_paths),
"Renaming",
leave=False):
src = filepath.parent / (f'{i+1:06}_{filepath.name}')
dst = filepath
try:
src.rename(dst)
except:
io.log_error('fail to rename %s' % (src.name))
return
def __init__(self,
root_path,
debug=False,
batch_size=1,
resolution=256,
generators_count=4,
data_format="NHWC",
**kwargs):
super().__init__(debug, batch_size)
self.initialized = False
dataset_path = root_path / 'CelebAMask-HQ'
if not dataset_path.exists():
raise ValueError(f'Unable to find {dataset_path}')
images_path = dataset_path / 'CelebA-HQ-img'
if not images_path.exists():
raise ValueError(f'Unable to find {images_path}')
masks_path = dataset_path / 'CelebAMask-HQ-mask-anno'
if not masks_path.exists():
raise ValueError(f'Unable to find {masks_path}')
if self.debug:
self.generators_count = 1
else:
self.generators_count = max(1, generators_count)
source_images_paths = pathex.get_image_paths(images_path,
return_Path_class=True)
source_images_paths_len = len(source_images_paths)
mask_images_paths = pathex.get_image_paths(masks_path,
subdirs=True,
return_Path_class=True)
if source_images_paths_len == 0 or len(mask_images_paths) == 0:
raise ValueError('No training data provided.')
mask_file_id_hash = {}
for filepath in io.progress_bar_generator(mask_images_paths,
"Loading"):
stem = filepath.stem
file_id, mask_type = stem.split('_', 1)
file_id = int(file_id)
if file_id not in mask_file_id_hash:
mask_file_id_hash[file_id] = {}
mask_file_id_hash[file_id][MaskType_from_name[mask_type]] = str(
filepath.relative_to(masks_path))
source_file_id_set = set()
for filepath in source_images_paths:
stem = filepath.stem
file_id = int(stem)
source_file_id_set.update({file_id})
for k in mask_file_id_hash.keys():
if k not in source_file_id_set:
io.log_err(f"Corrupted dataset: {k} not in {images_path}")
if self.debug:
self.generators = [
ThisThreadGenerator(
self.batch_func,
(images_path, masks_path, mask_file_id_hash, data_format))
]
else:
self.generators = [SubprocessGenerator ( self.batch_func, (images_path, masks_path, mask_file_id_hash, data_format), start_now=False ) \
for i in range(self.generators_count) ]
SubprocessGenerator.start_in_parallel(self.generators)
self.generator_counter = -1
self.initialized = True