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 convert(workspace, skip=True, manual=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)
model_dir = os.path.join(workspace, "model")
self_model_dir = os.path.join(workspace, f, "model")
if os.path.exists(self_model_dir):
io.log_info("Use Self Model")
model_dir = self_model_dir
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
# 恢复排序
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)
# 转换
dfl.dfl_convert(data_dst,
data_dst_merged,
data_dst_aligned,
model_dir,
enable_predef=not manual)
# ConverterMasked.enable_predef = enable_predef
# 去掉没有脸的
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, "data_trash")
import shutil
shutil.move(data_dst, trash_dir)
def prepare_vr(workspace):
import os
import shutil
for f in os.listdir(workspace):
ext = os.path.splitext(f)[-1]
if ext not in ['.mp4', '.avi']:
continue
if f.startswith("result"):
continue
# 获取所有的data_dst文件
tmp_dir = os.path.join(workspace, "_tmp")
tmp_aligned = os.path.join(tmp_dir, "aligned")
tmp_video_dir = os.path.join(tmp_dir, "video")
if os.path.exists(tmp_dir):
shutil.rmtree(tmp_dir)
if not os.path.exists(tmp_dir):
os.mkdir(tmp_dir)
os.mkdir(tmp_video_dir)
video = os.path.join(workspace, f)
# 提取帧
# VideoEd.extract_video(video, tmp_dir, "png", 0)
dfl.dfl_extract_video(video, tmp_dir, 0)
# 提取人脸
beep()
dfl.dfl_extract_faces(tmp_dir, tmp_aligned, "manual", False)
# aligned重命名
tmp_aligned2 = os.path.join(tmp_dir, "aligned2")
shutil.move(tmp_aligned, tmp_aligned2)
# 再次提取
beep()
dfl.dfl_extract_faces(tmp_dir, tmp_aligned, "manual", False)
# 两个aligned merge
for f2 in os.listdir(tmp_aligned2):
src = os.path.join(tmp_aligned2, f2)
dst = os.path.join(tmp_aligned, "0_" + f2)
shutil.move(src, dst)
# 保存video
shutil.copy(video, tmp_video_dir)
# 重命名
fname = f.replace(ext, "")
dst_dir = os.path.join(workspace,
"data_dst_%s_%s" % (get_time_str(), fname))
shutil.move(tmp_dir, dst_dir)
# 移动video
data_trash = os.path.join(workspace, "../trash_workspace")
if not os.path.exists(data_trash):
os.mkdir(data_trash)
shutil.move(video, data_trash)
beep()
def prepare(workspace, detector="s3fd", manual_fix=False):
import os
import shutil
for f in os.listdir(workspace):
if f.startswith("result"):
continue
ext = os.path.splitext(f)[-1]
if ext.lower() not in video_exts:
continue
# 获取所有的data_dst文件
tmp_dir = os.path.join(workspace, "_tmp")
tmp_aligned = os.path.join(tmp_dir, "aligned")
tmp_video_dir = os.path.join(tmp_dir, "video")
if os.path.exists(tmp_dir):
shutil.rmtree(tmp_dir)
if not os.path.exists(tmp_dir):
os.mkdir(tmp_dir)
os.mkdir(tmp_video_dir)
video_name = f
video = os.path.join(workspace, video_name)
# 提取帧
# VideoEd.extract_video(video, tmp_dir, "png", 0)
dfl.dfl_extract_video(video, tmp_dir, 0)
if detector == "manual":
beep()
# 提取人脸
dfl.dfl_extract_faces(tmp_dir,
tmp_aligned,
detector=detector,
manual_fix=manual_fix)
# 排序
if detector != "manual":
dfl.dfl_sort_by_hist(tmp_aligned)
# 保存video
shutil.copy(video, tmp_video_dir)
# 重命名
fname = f.replace(ext, "")
dst_dir = os.path.join(workspace,
"data_dst_%s_%s" % (get_time_str(), fname))
shutil.move(tmp_dir, dst_dir)
# 移动video
data_trash = os.path.join(workspace, "../trash_workspace")
if not os.path.exists(data_trash):
os.mkdir(data_trash)
trash_video = os.path.join(data_trash, video_name)
if os.path.exists(trash_video):
os.remove(trash_video)
shutil.move(video, trash_video)
beep()
def mp4(workspace):
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"]:
if testExt(v):
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 convert(workspace, model="SAEHD", force_recover=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)
model_dir = os.path.join(workspace, "model")
self_model_dir = os.path.join(workspace, f, "model")
if os.path.exists(self_model_dir):
io.log_info("Use Self Model")
model_dir = self_model_dir
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"]:
if testExt(v):
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
# 恢复排序
need_recover = True
for img in os.listdir(data_dst_aligned):
if img.endswith("_0.jpg") or img.endswith("_0.png"):
need_recover = False
break
if need_recover or force_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)
# 转换
dfl.dfl_merge(data_dst, data_dst_merged, data_dst_aligned, model_dir,
model)
def extract_src():
import os
import shutil
root_dir = get_root_path()
extract_workspace = os.path.join(root_dir, "extract_workspace")
target_dir = os.path.join(extract_workspace, "aligned_")
valid_exts = [".mp4", ".avi", ".wmv", ".mkv", ".ts"]
fps = io.input_int(
"Enter FPS ( ?:help skip:fullfps ) : ",
0,
help_message=
"How many frames of every second of the video will be extracted.")
def file_filter(file):
if os.path.isdir(os.path.join(extract_workspace, file)):
return False
ext = os.path.splitext(file)[-1]
if ext not in valid_exts:
return False
return True
files = list(filter(file_filter, os.listdir(extract_workspace)))
files.sort()
pos = 0
for file in files:
pos += 1
io.log_info("@@@@@ Start Process %s, %d / %d" %
(file, pos, len(files)))
# 提取图片
input_file = os.path.join(extract_workspace, file)
output_dir = os.path.join(extract_workspace, "extract_images")
if not os.path.exists(output_dir):
os.mkdir(output_dir)
for f in os.listdir(output_dir):
os.remove(os.path.join(output_dir, f))
dfl.dfl_extract_video(input_file, output_dir, fps)
io.log_info("@@@@@ Start Extract %s, %d / %d" %
(file, pos, len(files)))
# 提取人脸
input_dir = output_dir
output_dir = os.path.join(extract_workspace, "_current")
dfl.dfl_extract_faces(input_dir, output_dir, output_debug=True)
# 复制到结果集
io.log_info("@@@@@ Start Move %s, %d / %d" % (file, pos, len(files)))
if not os.path.exists(target_dir):
os.mkdir(target_dir)
ts = get_time_str()
for f in os.listdir(output_dir):
src = os.path.join(output_dir, f)
dst = os.path.join(target_dir, "%s_%s" % (ts, f))
shutil.move(src, dst)
# 全部做完，删除该文件
io.log_info("@@@@@ Finish %s, %d / %d" % (file, pos, len(files)))
os.remove(os.path.join(extract_workspace, file))
os.rmdir(output_dir)
# 做完后排序
io.log_info("@@@@@ Sort By Hist")
dfl.dfl_sort_by_hist(target_dir)
beep()