def remove_ie_polys_folder(input_path):
input_path = Path(input_path)
io.log_info("Removing ie_polys...\r\n")
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Removing"):
filepath = Path(filepath)
remove_ie_polys_file(filepath)
def convert_png_to_jpg_folder(input_path):
input_path = Path(input_path)
io.log_info("Converting PNG to JPG...\r\n")
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Converting"):
filepath = Path(filepath)
convert_png_to_jpg_file(filepath)
def final_process(input_path, img_list, trash_img_list):
if len(trash_img_list) != 0:
parent_input_path = input_path.parent
trash_path = parent_input_path / (input_path.stem + '_trash')
trash_path.mkdir(exist_ok=True)
io.log_info("Trashing %d items to %s" %
(len(trash_img_list), str(trash_path)))
for filename in Path_utils.get_image_paths(trash_path):
Path(filename).unlink()
for i in io.progress_bar_generator(range(len(trash_img_list)),
"Moving trash",
leave=False):
src = Path(trash_img_list[i][0])
dst = trash_path / src.name
try:
src.rename(dst)
except:
io.log_info('fail to trashing %s' % (src.name))
io.log_info("")
if len(img_list) != 0:
for i in io.progress_bar_generator([*range(len(img_list))],
"Renaming",
leave=False):
src = Path(img_list[i][0])
dst = input_path / ('%.5d_%s' % (i, src.name))
try:
src.rename(dst)
except:
io.log_info('fail to rename %s' % (src.name))
for i in io.progress_bar_generator([*range(len(img_list))],
"Renaming"):
src = Path(img_list[i][0])
src = input_path / ('%.5d_%s' % (i, src.name))
dst = input_path / ('%.5d%s' % (i, src.suffix))
try:
src.rename(dst)
except:
io.log_info('fail to rename %s' % (src.name))
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({2:1,3:1,4:4,5:8,6:16})
resolution = self.options['resolution']
face_type = self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
self.model = FUNIT( face_type_str=FaceType.toString(face_type),
batch_size=self.batch_size,
encoder_nf=64,
encoder_downs=2,
encoder_res_blk=2,
class_downs=4,
class_nf=64,
class_latent=64,
mlp_nf=256,
mlp_blks=2,
dis_nf=64,
dis_res_blks=10,
num_classes=2,
subpixel_decoder=True,
initialize_weights=self.is_first_run(),
is_training=self.is_training_mode,
tf_cpu_mode=self.options['optimizer_mode']-1
)
if not self.is_first_run():
self.load_weights_safe(self.model.get_model_filename_list())
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
if self.is_training_mode:
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True},
]
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types )
])
else:
generator = SampleGeneratorFace(self.training_data_src_path, batch_size=1,
sample_process_options=SampleProcessor.Options(),
output_sample_types=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ] )
io.log_info("Calculating average src face style...")
codes = []
for i in io.progress_bar_generator(range(generator.get_total_sample_count())):
codes += self.model.get_average_class_code( generator.generate_next() )
self.average_class_code = np.mean ( np.array(codes), axis=0 )[None,...]
def sort_by_face(input_path):
io.log_info("根据相似度[similarity]排序...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
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 不是DeepFaceLab的图片格式，请使用DeepFaceLab提取脸图" %
(filepath.name))
trash_img_list.append([str(filepath)])
continue
img_list.append([str(filepath), dflimg.get_landmarks()])
img_list_len = len(img_list)
for i in io.progress_bar_generator(range(0, img_list_len - 1), "Sorting"):
min_score = float("inf")
j_min_score = i + 1
for j in range(i + 1, len(img_list)):
fl1 = img_list[i][1]
fl2 = img_list[j][1]
score = np.sum(np.absolute((fl2 - fl1).flatten()))
if score < min_score:
min_score = score
j_min_score = j
img_list[i +
1], img_list[j_min_score] = img_list[j_min_score], img_list[i
+
1]
return img_list, trash_img_list
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 sort_by_face_dissim(input_path):
io.log_info("根据差异度[dissimilarity]排序...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
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 不是DeepFaceLab的图片格式，请使用DeepFaceLab提取脸图" %
(filepath.name))
trash_img_list.append([str(filepath)])
continue
img_list.append([str(filepath), dflimg.get_landmarks(), 0])
img_list_len = len(img_list)
for i in io.progress_bar_generator(range(img_list_len - 1), "Sorting"):
score_total = 0
for j in range(i + 1, len(img_list)):
if i == j:
continue
fl1 = img_list[i][1]
fl2 = img_list[j][1]
score_total += np.sum(np.absolute((fl2 - fl1).flatten()))
img_list[i][2] = score_total
io.log_info("排序...")
img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True)
return img_list, trash_img_list
def get_pitch_yaw_roll(input_path, r=0.05):
import os
import numpy as np
import cv2
from shutil import copyfile
from pathlib import Path
from utils import Path_utils
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from facelib import LandmarksProcessor
from joblib import Subprocessor
import multiprocessing
from interact import interact as io
from imagelib import estimate_sharpness
io.log_info("Sorting by face yaw...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
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))
trash_img_list.append([str(filepath)])
continue
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll(
dflimg.get_landmarks())
img_list.append([str(filepath), pitch, yaw, roll])
img_list.sort(key=lambda item: item[1])
with open(os.path.join(input_path, "_pitch_yaw_roll.csv"), "w") as f:
for i in img_list:
f.write("%s,%f,%f,%f\n" %
(os.path.basename(i[0]), i[1], i[2], i[3]))
import cv
width = 800
img = cv.cv_new((width, width))
xs = [i[1] for i in img_list]
ys = [i[2] for i in img_list]
cs = [(128, 128, 128)] * len(xs)
rs = [int(r * width / 2)] * len(xs)
cv.cv_scatter(img, xs, ys, [-1, 1], [-1, 1], cs, rs)
cs = [(0xcc, 0x66, 0x33)] * len(xs)
rs = [2] * len(xs)
cv.cv_scatter(img, xs, ys, [-1, 1], [-1, 1], cs, rs)
cv.cv_save(img, os.path.join(input_path, "_pitch_yaw_roll.bmp"))
return img_list
def sort_by_black(input_path):
io.log_info ("Sorting by amount of black pixels...")
img_list = []
for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"):
img = cv2_imread(x)
img_list.append ([x, img[(img == 0)].size ])
io.log_info ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=False)
return img_list
def sync_trash(trash_path, pool_path):
import shutil
count = 0
for f in io.progress_bar_generator(os.listdir(trash_path), "Trash Files"):
if f.endswith(".jpg") or f.endswith(".png"):
img_name = f.split("_")[-1]
img_path = os.path.join(pool_path, img_name)
dst_path = os.path.join(trash_path, "_origin")
if os.path.exists(img_path):
shutil.move(img_path, dst_path)
count += 1
io.log_info("Trash %d" % count)
def delete_relighted(input_dir):
input_path = Path(input_dir)
image_paths = [Path(x) for x in Path_utils.get_image_paths(input_path)]
files_to_delete = []
for filepath in io.progress_bar_generator(image_paths, "Loading"):
if filepath.suffix == '.png':
dflimg = DFLPNG.load(str(filepath))
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load(str(filepath))
else:
dflimg = None
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
continue
else:
if dflimg.get_relighted():
files_to_delete += [filepath]
for file in io.progress_bar_generator(files_to_delete, "Deleting"):
file.unlink()
def unpack(samples_path):
samples_dat_path = samples_path / packed_faceset_filename
if not samples_dat_path.exists():
io.log_info(f"{samples_dat_path} : file not found.")
return
samples = PackedFaceset.load(samples_path)
for sample in io.progress_bar_generator(samples, "Unpacking"):
with open(samples_path / sample.filename, "wb") as f:
f.write( sample.read_raw_file() )
samples_dat_path.unlink()
def select(exists_path, pool_path, div=200):
# 先计算output_path的已有图像
import cv
import dfl
import random
width = 800
trans = cv.trans_fn(-1, 1, 0, width)
img = cv.cv_new((width, width))
for f in io.progress_bar_generator(os.listdir(exists_path),
"Existing Imgs"):
if f.endswith(".png") or f.endswith("jpg"):
img_path = os.path.join(exists_path, f)
dfl_img = dfl.dfl_load_img(img_path)
pitch, yaw, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
pitch = trans(pitch)
yaw = trans(yaw)
cv.cv_circle(img, (pitch, yaw), (128, 128, 128), width / div, -1)
time_str = get_time_str()
import shutil
pool_files = list(os.listdir(pool_path))
# random.shuffle(pool_files)
count = 0
for f in io.progress_bar_generator(pool_files,
os.path.basename(pool_path)):
if f.endswith(".png") or f.endswith(".jpg"):
img_path = os.path.join(pool_path, f)
dfl_img = dfl.dfl_load_img(img_path)
pitch, yaw, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
pitch = trans(pitch)
yaw = trans(yaw)
if sum(img[yaw][pitch]) == 255 * 3:
dst = os.path.join(exists_path, "%s_%s" % (time_str, f))
shutil.copy(img_path, dst)
count += 1
cv.cv_circle(img, (pitch, yaw), (0xcc, 0x66, 0x33),
width / div, -1)
cv.cv_save(img, os.path.join(exists_path, "_select.bmp"))
io.log_info("Copy %d, Total %d" % (count, len(pool_files)))
def sort_by_vggface(input_path):
io.log_info("Sorting by face similarity using VGGFace model...")
model = VGGFace()
final_img_list = []
trash_img_list = []
image_paths = Path_utils.get_image_paths(input_path)
img_list = [(x, ) for x in image_paths]
img_list_len = len(img_list)
img_list_range = [*range(img_list_len)]
feats = [None] * img_list_len
for i in io.progress_bar_generator(img_list_range, "Loading"):
img = cv2_imread(img_list[i][0]).astype(np.float32)
img = imagelib.normalize_channels(img, 3)
img = cv2.resize(img, (224, 224))
img = img[..., ::-1]
img[..., 0] -= 93.5940
img[..., 1] -= 104.7624
img[..., 2] -= 129.1863
feats[i] = model.predict(img[None, ...])[0]
tmp = np.zeros((img_list_len, ))
float_inf = float("inf")
for i in io.progress_bar_generator(range(img_list_len - 1), "Sorting"):
i_feat = feats[i]
for j in img_list_range:
tmp[j] = npla.norm(i_feat - feats[j]) if j >= i + 1 else float_inf
idx = np.argmin(tmp)
img_list[i + 1], img_list[idx] = img_list[idx], img_list[i + 1]
feats[i + 1], feats[idx] = feats[idx], feats[i + 1]
return img_list, trash_img_list
def sort_by_face(input_path):
io.log_info("Sorting by face similarity...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
filepath = Path(filepath)
dflimg = DFLIMG.load(filepath)
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
trash_img_list.append([str(filepath)])
continue
img_list.append([str(filepath), dflimg.get_landmarks()])
img_list_len = len(img_list)
for i in io.progress_bar_generator(range(0, img_list_len - 1), "Sorting"):
min_score = float("inf")
j_min_score = i + 1
for j in range(i + 1, len(img_list)):
fl1 = img_list[i][1]
fl2 = img_list[j][1]
score = np.sum(np.absolute((fl2 - fl1).flatten()))
if score < min_score:
min_score = score
j_min_score = j
img_list[i +
1], img_list[j_min_score] = img_list[j_min_score], img_list[i
+
1]
return img_list, trash_img_list
def load(sample_type, samples_path, target_samples_path=None):
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:
datas[sample_type] = SampleLoader.upgradeToFaceSamples([
Sample(filename=filename)
for filename in Path_utils.get_image_paths(samples_path)
])
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 reload_src():
nonlocal src_img_list
nonlocal src_cur_list
src_img_list = []
if src_path:
for f in io.progress_bar_generator(os.listdir(src_path),
"Loading"):
if f.endswith(".jpg") or f.endswith(".png"):
fpath = os.path.join(src_path, f)
dfl_img = dfl.dfl_load_img(fpath)
p, y, _ = dfl.dfl_estimate_pitch_yaw_roll(dfl_img)
src_img_list.append([fno, p, y])
src_img_list.append([fno, p, -y])
src_img_list = np.array(src_img_list, "float")
src_cur_list = src_img_list
def merge(input_path, target_path):
import os
import shutil
for f in os.listdir(input_path):
sub_path = os.path.join(input_path, f)
if os.path.abspath(sub_path) == os.path.abspath(target_path):
continue
if os.path.isdir(sub_path):
time_str = get_time_str()
for img in io.progress_bar_generator(os.listdir(sub_path), f):
if img.endswith(".png") or img.endswith(".jpg"):
img_path = os.path.join(sub_path, img)
dst_path = os.path.join(target_path,
"%s_%s" % (time_str, img))
shutil.move(img_path, dst_path)
def dev_test(input_dir):
input_path = Path(input_dir)
dir_names = Path_utils.get_all_dir_names(input_path)
for dir_name in io.progress_bar_generator(dir_names, desc="Processing"):
img_paths = Path_utils.get_image_paths (input_path / dir_name)
for filename in img_paths:
filepath = Path(filename)
dflimg = DFLIMG.load (filepath)
if dflimg is None:
raise ValueError
dflimg.embed_and_set(filename, person_name=dir_name)
def split(input_path, target_path, batch=3000):
import os
import shutil
count = 0
if not os.path.exists(target_path):
os.mkdir(target_path)
dst_dir = os.path.join(target_path, "split_%03d" % int(count / batch))
for f in io.progress_bar_generator(os.listdir(input_path), "Process"):
if not f.endswith(".jpg") and not f.endswith(".png"):
continue
if count % batch == 0:
dst_dir = os.path.join(target_path,
"split_%03d" % int(count / batch))
os.mkdir(dst_dir)
src = os.path.join(input_path, f)
shutil.move(src, dst_dir)
count += 1
def upgradeToFaceSamples(samples, silent=False):
sample_list = []
for s in (samples if silent else io.progress_bar_generator(
samples, "Loading")):
s_filename_path = Path(s.filename)
try:
if s_filename_path.suffix == '.png':
dflimg = DFLPNG.load(str(s_filename_path))
elif s_filename_path.suffix == '.jpg':
dflimg = DFLJPG.load(str(s_filename_path))
else:
dflimg = None
if dflimg is None:
print("%s is not a dfl image file required for training" %
(s_filename_path.name))
continue
landmarks = dflimg.get_landmarks()
pitch_yaw_roll = dflimg.get_pitch_yaw_roll()
eyebrows_expand_mod = dflimg.get_eyebrows_expand_mod()
if pitch_yaw_roll is None:
pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(
landmarks)
sample_list.append(
s.copy_and_set(
sample_type=SampleType.FACE,
face_type=FaceType.fromString(dflimg.get_face_type()),
shape=dflimg.get_shape(),
landmarks=landmarks,
ie_polys=dflimg.get_ie_polys(),
pitch_yaw_roll=pitch_yaw_roll,
eyebrows_expand_mod=eyebrows_expand_mod,
source_filename=dflimg.get_source_filename(),
fanseg_mask_exist=dflimg.get_fanseg_mask() is not None,
))
except:
print("Unable to load %s , error: %s" %
(str(s_filename_path), traceback.format_exc()))
return sample_list
def load(sample_type, samples_path):
samples_cache = SampleHost.samples_cache
if str(samples_path) not in samples_cache.keys():
samples_cache[str(samples_path)] = [None] * SampleType.QTY
samples = samples_cache[str(samples_path)]
if sample_type == SampleType.IMAGE:
if samples[sample_type] is None:
samples[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 samples[sample_type] is None:
result = None
try:
result = samplelib.PackedFaceset.load(samples_path)
except:
io.log_err(
f"Error occured while loading samplelib.PackedFaceset.load {str(samples_dat_path)}, {traceback.format_exc()}"
)
if result is not None:
io.log_info(
f"Loaded {len(result)} packed faces from {samples_path}"
)
if result is None:
result = SampleHost.load_face_samples(
Path_utils.get_image_paths(samples_path))
samples[sample_type] = result
elif sample_type == SampleType.FACE_TEMPORAL_SORTED:
if samples[sample_type] is None:
samples[
sample_type] = SampleHost.upgradeToFaceTemporalSortedSamples(
SampleHost.load(SampleType.FACE, samples_path))
return samples[sample_type]
def load_face_samples(image_paths):
sample_list = []
for filename in io.progress_bar_generator(image_paths, desc="Loading"):
dflimg = DFLIMG.load(Path(filename))
if dflimg is None:
io.log_err(f"{filename} is not a dfl image file.")
else:
sample_list.append(
Sample(
filename=filename,
sample_type=SampleType.FACE,
face_type=FaceType.fromString(dflimg.get_face_type()),
shape=dflimg.get_shape(),
landmarks=dflimg.get_landmarks(),
ie_polys=dflimg.get_ie_polys(),
eyebrows_expand_mod=dflimg.get_eyebrows_expand_mod(),
source_filename=dflimg.get_source_filename(),
))
return sample_list
def add_landmarks_debug_images(input_path):
io.log_info ("Adding landmarks debug images...")
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"):
filepath = Path(filepath)
img = cv2_imread(str(filepath))
dflimg = DFLIMG.load (filepath)
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
if img is not None:
face_landmarks = dflimg.get_landmarks()
LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True, ie_polys=IEPolys.load(dflimg.get_ie_polys()) )
output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem), '_debug.jpg')
cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
def sort_by_origname(input_path):
io.log_info("Sort by original filename...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
filepath = Path(filepath)
dflimg = DFLIMG.load(filepath)
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
trash_img_list.append([str(filepath)])
continue
img_list.append([str(filepath), dflimg.get_source_filename()])
io.log_info("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1))
return img_list, trash_img_list
def merge_dst_aligned(workspace):
import shutil
counter = 0
target_dst = os.path.join(workspace, "data_dst")
if os.path.exists(target_dst):
shutil.rmtree(target_dst)
target_dst_aligned = os.path.join(target_dst, "aligned")
os.makedirs(target_dst_aligned)
for f in os.listdir(workspace):
dst_path = os.path.join(workspace, f)
if os.path.isdir(dst_path) and f.startswith("data_dst_"):
counter += 1
dst_aligned = os.path.join(dst_path, "aligned")
for img in io.progress_bar_generator(os.listdir(dst_aligned),
"Process"):
if img.endswith(".png") or img.endswith(".jpg"):
img_path = os.path.join(dst_aligned, img)
base_name = os.path.basename(img_path)
dst_img_path = os.path.join(target_dst_aligned,
"%d_%s" % (counter, base_name))
shutil.copy(img_path, dst_img_path)
def unpack(samples_path):
samples_dat_path = samples_path / packed_faceset_filename
if not samples_dat_path.exists():
io.log_info(f"{samples_dat_path} : file not found.")
return
samples = PackedFaceset.load(samples_path)
for sample in io.progress_bar_generator(samples, "Unpacking"):
person_name = sample.person_name
if person_name is not None:
person_path = samples_path / person_name
person_path.mkdir(parents=True, exist_ok=True)
target_filepath = person_path / sample.filename
else:
target_filepath = samples_path / sample.filename
with open(target_filepath, "wb") as f:
f.write(sample.read_raw_file())
samples_dat_path.unlink()
def sort_by_hist_dissim(input_path):
io.log_info("根据[histogram dissimilarity]排序...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
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 不是DeepFaceLab的图片格式，请使用DeepFaceLab提取脸图" %
(filepath.name))
trash_img_list.append([str(filepath)])
continue
image = cv2_imread(str(filepath))
face_mask = LandmarksProcessor.get_image_hull_mask(
image.shape, dflimg.get_landmarks())
image = (image * face_mask).astype(np.uint8)
img_list.append([
str(filepath),
cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None,
[256], [0, 256]), 0
])
img_list = HistDissimSubprocessor(img_list).run()
io.log_info("排序...")
img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True)
return img_list, trash_img_list
def restore_faceset_metadata_folder(input_path):
input_path = Path(input_path)
metadata_filepath = input_path / 'meta.dat'
io.log_info(f"Restoring metadata from {str(metadata_filepath)}.\r\n")
if not metadata_filepath.exists():
io.log_err(f"Unable to find {str(metadata_filepath)}.")
try:
with open(metadata_filepath, "rb") as f:
d = pickle.loads(f.read())
except:
raise FileNotFoundError(filename)
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Processing"):
filepath = Path(filepath)
shape, dfl_dict = d.get(filepath.name, None)
img = cv2_imread(str(filepath))
if img.shape != shape:
img = cv2.resize(img, (shape[1], shape[0]), cv2.INTER_LANCZOS4)
if filepath.suffix == '.png':
cv2_imwrite(str(filepath), img)
elif filepath.suffix == '.jpg':
cv2_imwrite(str(filepath), img,
[int(cv2.IMWRITE_JPEG_QUALITY), 100])
if filepath.suffix == '.png':
DFLPNG.embed_dfldict(str(filepath), dfl_dict)
elif filepath.suffix == '.jpg':
DFLJPG.embed_dfldict(str(filepath), dfl_dict)
else:
continue
metadata_filepath.unlink()
def skip_no_face(data_dst_dir):
import os
import shutil
data_dst_aligned_dir = os.path.join(data_dst_dir, "aligned")
aligend = set([f.split('_')[0] for f in os.listdir(data_dst_aligned_dir)])
merged_dir = os.path.join(data_dst_dir, "merged")
merged_trash_dir = os.path.join(data_dst_dir, "merged_trash")
if os.path.exists(merged_trash_dir):
# raise Exception("Merge Dir Bak Exists")
shutil.rmtree(merged_trash_dir)
shutil.move(merged_dir, merged_trash_dir)
os.mkdir(merged_dir)
idx = 0
for f in io.progress_bar_generator(os.listdir(merged_trash_dir),
"Skip No Face"):
name = os.path.splitext(f)[0]
ext = os.path.splitext(f)[-1]
if name in aligend:
idx += 1
src = os.path.join(merged_trash_dir, f)
dst = os.path.join(merged_dir, "%05d" % idx + ext)
shutil.move(src, dst)
