def sort_by_blur(input_path):
io.log_info("Sorting by blur...")
img_list = [(filename, [])
for filename in Path_utils.get_image_paths(input_path)]
img_list, trash_img_list = BlurEstimatorSubprocessor(img_list).run()
io.log_info("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
def refix(workspace):
dst = get_workspace_dst(workspace)
dst_aligned = os.path.join(dst, "aligned")
recover_filename_if_nessesary(dst_aligned)
extract_imgs = [
f if f.endswith(".jpg") or f.endswith(".png") else ""
for f in os.listdir(dst)
]
max_img_no = int(max(extract_imgs).split(".")[0])
ext = extract_imgs[0].split(".")[1]
aligned_imgs = list(
sorted(
filter(lambda x: x is not None, [
f if f.endswith(".jpg") or f.endswith(".png") else None
for f in os.listdir(dst_aligned)
])))
need_fix_no = []
i = 0 # 当前文件下标
j = 1 # 期望文件名
while i <= max_img_no and j <= max_img_no and i < len(aligned_imgs):
if aligned_imgs[i].startswith("%05d" % j):
i += 1
j += 1
else:
# print(aligned_imgs[i], j)
need_fix_no.append(j)
j += 1
for k in range(j, max_img_no + 1):
need_fix_no.append(k)
# print(k)
if len(need_fix_no) == 0:
return
fix_workspace = os.path.join(dst, "fix")
import shutil
if os.path.exists(fix_workspace):
shutil.rmtree(fix_workspace)
os.mkdir(fix_workspace)
for no in need_fix_no:
f = os.path.join(dst, "%05d.%s" % (no, ext))
io.log_info(f)
shutil.copy(f, fix_workspace)
fix_workspace_aligned = os.path.join(fix_workspace, "aligned")
from mainscripts import Extractor
Extractor.main(fix_workspace,
fix_workspace_aligned,
detector="manual",
manual_fix=False)
Extractor.extract_fanseg(fix_workspace_aligned)
for f in os.listdir(fix_workspace_aligned):
f = os.path.join(fix_workspace_aligned, f)
io.log_info(f)
shutil.move(f, dst_aligned)
shutil.rmtree(fix_workspace)
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 recover_original_aligned_filename(input_path):
io.log_info ("Recovering original aligned filename...")
files = []
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"):
filepath = Path(filepath)
dflimg = DFLIMG.load (filepath)
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
files += [ [filepath, None, dflimg.get_source_filename(), False] ]
files_len = len(files)
for i in io.progress_bar_generator( range(files_len), "Sorting" ):
fp, _, sf, converted = files[i]
if converted:
continue
sf_stem = Path(sf).stem
files[i][1] = fp.parent / ( sf_stem + '_0' + fp.suffix )
files[i][3] = True
c = 1
for j in range(i+1, files_len):
fp_j, _, sf_j, converted_j = files[j]
if converted_j:
continue
if sf_j == sf:
files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix )
files[j][3] = True
c += 1
for file in io.progress_bar_generator( files, "Renaming", leave=False ):
fs, _, _, _ = file
dst = fs.parent / ( fs.stem + '_tmp' + fs.suffix )
try:
fs.rename (dst)
except:
io.log_err ('fail to rename %s' % (fs.name) )
for file in io.progress_bar_generator( files, "Renaming" ):
fs, fd, _, _ = file
fs = fs.parent / ( fs.stem + '_tmp' + fs.suffix )
try:
fs.rename (fd)
except:
io.log_err ('fail to rename %s' % (fs.name) )
def auto(workspace):
import subprocess
for f in os.listdir(workspace):
if os.path.isdir(os.path.join(get_root_path(), "workspace",
f)) and f.startswith("data_dst_"):
train_bat = os.path.join(get_root_path(), "auto_train.bat")
convert_bat = os.path.join(get_root_path(), "auto_convert.bat")
step_bat = os.path.join(get_root_path(), "auto_step.bat")
subprocess.call([train_bat])
subprocess.call([convert_bat])
subprocess.call([step_bat])
io.log_info("Finish " + f)
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 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 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 save(self):
io.log_info ("Saving...")
Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text)
self.onSave()
model_data = {
'epoch': self.epoch,
'options': self.options,
'loss_history': self.loss_history,
'sample_for_preview' : self.sample_for_preview
}
self.model_data_path.write_bytes( pickle.dumps(model_data) )
def pack(samples_path):
samples_dat_path = samples_path / packed_faceset_filename
if samples_dat_path.exists():
io.log_info(f"{samples_dat_path} : file already exists !")
io.input_bool("Press enter to continue and overwrite.", False)
of = open(samples_dat_path, "wb")
image_paths = Path_utils.get_image_paths(samples_path)
samples = samplelib.SampleHost.load_face_samples(image_paths)
samples_len = len(samples)
samples_configs = []
for sample in samples:
sample.filename = str(Path(sample.filename).relative_to(samples_path))
samples_configs.append ( sample.get_config() )
samples_bytes = pickle.dumps(samples_configs, 4)
of.write ( struct.pack ("Q", PackedFaceset.VERSION ) )
of.write ( struct.pack ("Q", len(samples_bytes) ) )
of.write ( samples_bytes )
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:
with open( samples_path / sample.filename, "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.filename}")
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"):
Path(filename).unlink()
def ask_settings(self):
s = """Choose mode: \n"""
for key in self.mode_dict.keys():
s += f"""({key}) {self.mode_dict[key]}\n"""
s += f"""Default: {self.default_mode} : """
mode = io.input_int (s, self.default_mode)
self.mode = self.mode_dict.get (mode, self.mode_dict[self.default_mode] )
if 'raw' not in self.mode:
if self.mode == 'hist-match' or self.mode == 'hist-match-bw':
self.masked_hist_match = io.input_bool("Masked hist match? (y/n skip:y) : ", True)
if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match':
self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255)
if self.face_type == FaceType.FULL:
s = """Choose mask mode: \n"""
for key in self.full_face_mask_mode_dict.keys():
s += f"""({key}) {self.full_face_mask_mode_dict[key]}\n"""
s += f"""?:help Default: 1 : """
self.mask_mode = io.input_int (s, 1, valid_list=self.full_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks.")
else:
s = """Choose mask mode: \n"""
for key in self.half_face_mask_mode_dict.keys():
s += f"""({key}) {self.half_face_mask_mode_dict[key]}\n"""
s += f"""?:help , Default: 1 : """
self.mask_mode = io.input_int (s, 1, valid_list=self.half_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images.")
if 'raw' not in self.mode:
self.erode_mask_modifier = self.base_erode_mask_modifier + np.clip ( io.input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (self.default_erode_mask_modifier), self.default_erode_mask_modifier), -200, 200)
self.blur_mask_modifier = self.base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (self.default_blur_mask_modifier), self.default_blur_mask_modifier), -200, 200)
self.motion_blur_power = np.clip ( io.input_int ("Choose motion blur power [0..100] (skip:%d) : " % (0), 0), 0, 100)
self.output_face_scale = np.clip (io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0), -50, 50)
if 'raw' not in self.mode:
self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct'])
self.color_transfer_mode = self.ctm_str_dict[self.color_transfer_mode]
super().ask_settings()
if 'raw' not in self.mode:
self.color_degrade_power = np.clip ( io.input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100)
self.export_mask_alpha = io.input_bool("Export png with alpha channel of the mask? (y/n skip:n) : ", False)
io.log_info ("")
def process_folder(dirpath, multi_gpu=False, cpu_only=False):
output_dirpath = dirpath.parent / (dirpath.name + '_enhanced')
output_dirpath.mkdir(exist_ok=True, parents=True)
dirpath_parts = '/'.join(dirpath.parts[-2:])
output_dirpath_parts = '/'.join(output_dirpath.parts[-2:])
io.log_info(f"Enhancing faceset in {dirpath_parts}")
io.log_info(f"Processing to {output_dirpath_parts}")
output_images_paths = Path_utils.get_image_paths(output_dirpath)
if len(output_images_paths) > 0:
for filename in output_images_paths:
Path(filename).unlink()
image_paths = [Path(x) for x in Path_utils.get_image_paths(dirpath)]
result = FacesetEnhancerSubprocessor(image_paths,
output_dirpath,
multi_gpu=multi_gpu,
cpu_only=cpu_only).run()
is_merge = io.input_bool(
f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ? (y/n skip:y) : ",
True)
if is_merge:
io.log_info(f"Copying processed files to {dirpath_parts}")
for (filepath, output_filepath) in result:
try:
shutil.copy(output_filepath, filepath)
except:
pass
io.log_info(f"Removing {output_dirpath_parts}")
shutil.rmtree(output_dirpath)
def auto_extract_to_img():
workspace = os.path.join(get_root_path(), "workspace")
data_dst = None
for f in os.listdir(workspace):
if f.startswith("data_dst_"):
data_dst = f
break
io.log_info(data_dst)
video_name = None
if data_dst is not None:
name = "_".join(data_dst.split("_")[8:])
print(name)
for f in os.listdir(os.path.join(workspace, "data_trash")):
if f.startswith(name):
video_name = f
break
io.log_info(video_name)
if video_name is not None:
video_path = os.path.join(workspace, "data_trash", video_name)
data_dst_path = os.path.join(workspace, data_dst)
io.log_info(video_path)
io.log_info(data_dst_path)
for f in io.progress_bar_generator(os.listdir(data_dst_path),
"Remove"):
if f.endswith(".jpg") or f.endswith(".png"):
os.remove(os.path.join(data_dst_path, f))
dfl.dfl_extract_video(video_path, data_dst_path)
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 = "%s/%.5d_%s" % (input_path, i, src.name)
# 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 = "%s/%.5d_%s" % (input_path, i, src.name)
src = Path(src)
dst = "%s/%.5d%s" % (input_path, i, src.suffix)
# 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 sort_by_oneface_in_image(input_path):
io.log_info ("Sort by one face in images...")
image_paths = Path_utils.get_image_paths(input_path)
a = np.array ([ ( int(x[0]), int(x[1]) ) \
for x in [ Path(filepath).stem.split('_') for filepath in image_paths ] if len(x) == 2
])
if len(a) > 0:
idxs = np.ndarray.flatten ( np.argwhere ( a[:,1] != 0 ) )
idxs = np.unique ( a[idxs][:,0] )
idxs = np.ndarray.flatten ( np.argwhere ( np.array([ x[0] in idxs for x in a ]) == True ) )
if len(idxs) > 0:
io.log_info ("Found %d images." % (len(idxs)) )
img_list = [ (path,) for i,path in enumerate(image_paths) if i not in idxs ]
trash_img_list = [ (image_paths[x],) for x in idxs ]
return img_list, trash_img_list
return [], []
def train(workspace, target_loss=0.01):
import os
model_dir = os.path.join(workspace, "model")
for f in os.listdir(workspace):
if not os.path.isdir(os.path.join(workspace,
f)) or not f.startswith("data_dst_"):
continue
io.log_info(f)
data_dst = os.path.join(workspace, f)
data_src_aligned = os.path.join(data_dst, "src")
if not os.path.exists(data_src_aligned):
data_src_aligned = os.path.join(workspace, "data_src", "aligned")
data_dst_aligned = os.path.join(data_dst, "aligned")
# 训练
dfl.dfl_train(data_src_aligned, data_dst_aligned, model_dir)
return
def mp4(workspace, skip=False):
import os
for f in os.listdir(workspace):
if not os.path.isdir(os.path.join(workspace,
f)) or not f.startswith("data_dst_"):
continue
io.log_info(f)
data_dst = os.path.join(workspace, f)
data_dst_merged = os.path.join(data_dst, "merged")
data_dst_aligned = os.path.join(data_dst, "aligned")
data_dst_video = os.path.join(data_dst, "video")
refer_path = None
for v in os.listdir(data_dst_video):
if v.split(".")[-1] in ["mp4", "avi", "wmv", "mkv"]:
refer_path = os.path.join(data_dst_video, v)
break
if not refer_path:
io.log_err("No Refer File In " + data_dst_video)
return
io.log_info("Refer File " + refer_path)
# 恢复排序
need_recover = True
for img in os.listdir(data_dst_aligned):
if img.endswith("_0.jpg") or img.endswith("_0.png"):
need_recover = False
if need_recover:
recover_filename(data_dst_aligned)
# 如果data_dst里没有脸则extract
has_img = False
for img in os.listdir(data_dst):
if img.endswith(".jpg") or img.endswith(".png"):
has_img = True
break
if not has_img:
dfl.dfl_extract_video(refer_path, data_dst)
# 去掉没有脸的
# if skip:
# skip_no_face(data_dst)
# 转mp4
refer_name = ".".join(os.path.basename(refer_path).split(".")[:-1])
result_path = os.path.join(
workspace, "result_%s_%s.mp4" % (get_time_str(), refer_name))
dfl.dfl_video_from_sequence(data_dst_merged, result_path, refer_path)
# 移动到trash
trash_dir = os.path.join(workspace, "../trash_workspace")
import shutil
shutil.move(data_dst, trash_dir)
def on_initialize(self, client_dict):
io.log_info('Running on %s.' % (client_dict['device_name']))
self.device_idx = client_dict['device_idx']
self.device_name = client_dict['device_name']
self.converter = client_dict['converter']
self.output_path = Path(
client_dict['output_dir']) if 'output_dir' in client_dict.keys(
) else None
self.alignments = client_dict['alignments']
self.debug = client_dict['debug']
from nnlib import nnlib
#model process ate all GPU mem,
#so we cannot use GPU for any TF operations in converter processes
#therefore forcing active_DeviceConfig to CPU only
nnlib.active_DeviceConfig = nnlib.DeviceConfig(cpu_only=True)
return None
def extract_fanseg(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.')
paths_to_extract = []
for filename in Path_utils.get_image_paths(input_path) :
filepath = Path(filename)
dflimg = DFLIMG.load ( filepath )
if dflimg is not None:
paths_to_extract.append (filepath)
paths_to_extract_len = len(paths_to_extract)
if paths_to_extract_len > 0:
io.log_info ("Performing extract fanseg for %d files..." % (paths_to_extract_len) )
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in paths_to_extract ], 'fanseg', multi_gpu=multi_gpu, cpu_only=cpu_only).run()
def main(input_path, sort_by_method):
input_path = Path(input_path)
sort_by_method = sort_by_method.lower()
io.log_info("Running sort tool.\r\n")
img_list = []
trash_img_list = []
if sort_by_method == 'blur':
img_list, trash_img_list = sort_by_blur(input_path)
elif sort_by_method == 'face':
img_list, trash_img_list = sort_by_face(input_path)
elif sort_by_method == 'face-dissim':
img_list, trash_img_list = sort_by_face_dissim(input_path)
elif sort_by_method == 'face-yaw':
img_list, trash_img_list = sort_by_face_yaw(input_path)
elif sort_by_method == 'face-pitch':
img_list, trash_img_list = sort_by_face_pitch(input_path)
elif sort_by_method == 'hist':
img_list = sort_by_hist(input_path)
elif sort_by_method == 'hist-dissim':
img_list, trash_img_list = sort_by_hist_dissim(input_path)
elif sort_by_method == 'brightness':
img_list = sort_by_brightness(input_path)
elif sort_by_method == 'hue':
img_list = sort_by_hue(input_path)
elif sort_by_method == 'black':
img_list = sort_by_black(input_path)
elif sort_by_method == 'origname':
img_list, trash_img_list = sort_by_origname(input_path)
elif sort_by_method == 'oneface':
img_list, trash_img_list = sort_by_oneface_in_image(input_path)
elif sort_by_method == 'vggface':
img_list, trash_img_list = sort_by_vggface(input_path)
elif sort_by_method == 'absdiff':
img_list, trash_img_list = sort_by_absdiff(input_path)
elif sort_by_method == 'final':
img_list, trash_img_list = sort_final(input_path)
elif sort_by_method == 'final-no-blur':
img_list, trash_img_list = sort_final(input_path,
include_by_blur=False)
final_process(input_path, img_list, trash_img_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 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 on_clients_finalized(self):
io.progress_bar_close()
if self.is_interactive:
self.screen_manager.finalize()
for frame in self.frames:
frame.output_filename = None
frame.image = None
session_data = {
'frames': self.frames,
'frames_idxs': self.frames_idxs,
'frames_done_idxs': self.frames_done_idxs,
'model_iter' : self.model_iter,
}
self.converter_session_filepath.write_bytes( pickle.dumps(session_data) )
io.log_info ("Session is saved to " + '/'.join (self.converter_session_filepath.parts[-2:]) )
def step(workspace):
import shutil
for f in os.listdir(workspace):
if os.path.isdir(os.path.join(workspace,
f)) and f.startswith("data_dst_"):
model = os.path.join(workspace, "model")
model_dst = os.path.join(workspace, f, "model")
if not os.path.exists(model_dst):
io.log_info("Move Model Files To %s" % f)
os.mkdir(model_dst)
for m in os.listdir(model):
mf = os.path.join(model, m)
if os.path.isfile(mf):
shutil.copy(os.path.join(model, m), model_dst)
src = os.path.join(workspace, f)
dst = os.path.join(workspace, "data_trash")
io.log_info("Move %s To %s" % (src, dst))
shutil.move(src, dst)
return
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 change_workspace():
wss = []
for f in os.listdir(get_root_path()):
fpath = os.path.join(get_root_path(), f)
if os.path.isfile(fpath) and f.startswith("@workspace"):
os.remove(fpath)
elif os.path.isdir(fpath) and f.startswith("workspace"):
wss.append(f)
inputs = "1234567890"[0:len(wss)]
for i in range(0, len(wss)):
io.log_info("[ %s ] %s" % (inputs[i], wss[i]))
no = io.input_str("Select Workspace:", 0)[0]
idx = inputs.find(no)
if idx < 0:
raise Exception("Invalid Idx " + no)
ws = wss[idx]
io.log_info("Select " + ws)
f = open(os.path.join(get_root_path(), "@" + ws), 'w')
f.write(ws)
f.close()
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 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 skip_by_pitch(src_path, dst_path):
import os
import shutil
import cv
size = 800
r = 20
src_img_list = get_pitch_yaw_roll(src_path)
dst_img_list = get_pitch_yaw_roll(dst_path)
trash_path = dst_path + "_trash"
if not os.path.exists(trash_path):
os.makedirs(trash_path)
img = cv.cv_new((size + 1, size + 1))
trans: Callable[[Any], int] = lambda v: int((v + 1) * size / 2)
count = 0
for [_, pitch, yaw, _] in src_img_list:
x = trans(pitch)
y = trans(yaw)
cv.cv_point(img, (x, y), (128, 128, 128), r)
# cv.cv_show(img)
xys = []
for [path, pitch, yaw, _] in dst_img_list:
x = trans(pitch)
y = trans(yaw)
c = img[y, x]
c_ = img[-y, x]
if sum(c) == 255 * 3 and sum(c_) == 255 * 3:
xys.append((x, y, (0, 0, 0xff)))
if not os.path.exists(path) or not os.path.exists(trash_path):
continue
count += 1
shutil.move(path, trash_path)
else:
xys.append((x, y, (0xcc, 0x66, 0x33)))
for (x, y, color) in xys:
cv.cv_point(img, (x, y), color, 2)
# cv.cv_show(img)
io.log_info("Out Of Pitch, %d / %d" % (count, len(dst_img_list)))
save_path = os.path.join(dst_path, "_skip_by_pitch.bmp")
cv.cv_save(img, save_path)
