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 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)
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 match_by_pitch(data_src_path, data_dst_path):
r = 0.05
mn = 1
mx = 3
import cv
import shutil
# 准备各种路径
src_aligned_store = os.path.join(data_src_path, "aligned_store")
if not os.path.exists(src_aligned_store):
raise Exception("No Src Aligned Store")
src_aligned = os.path.join(data_dst_path, "src")
if os.path.exists(src_aligned):
shutil.rmtree(src_aligned)
os.mkdir(src_aligned)
dst_aligned = os.path.join(data_dst_path, "aligned")
dst_aligned_trash = os.path.join(data_dst_path, "aligned_trash")
if not os.path.exists(dst_aligned_trash):
os.mkdir(dst_aligned_trash)
# 读取角度信息
src_img_list = get_pitch_yaw_roll(src_aligned_store)
dst_img_list = get_pitch_yaw_roll(dst_aligned)
src_pitch = list([i[1] for i in src_img_list])
src_yaw = list([i[2] for i in src_img_list])
dst_pitch = list([i[1] for i in dst_img_list])
dst_yaw = list([i[2] for i in dst_img_list])
src_ps = np.array(list(zip(src_pitch, src_yaw)), "float")
dst_ps = np.array(list(zip(dst_pitch, dst_yaw)), "float")
# 计算最近的n个点
src_match = set()
dst_match = set()
for p, i in io.progress_bar_generator(zip(dst_ps, range(len(dst_ps))),
"Calculating"):
ds = np.linalg.norm(src_ps - p, axis=1, keepdims=True)
idxs = np.argsort(ds, axis=0)
min_idx = idxs[mn - 1][0]
# 极端情况所有距离都不满足半径范围
if ds[min_idx] > r:
continue
# 至少有一个满足半径条件了,dst_point可以留下
dst_match.add(i)
# 所有满足条件的加入到src_match
for idx in idxs[:mx]:
idx = idx[0]
if ds[idx] > r:
break
src_match.add(idx)
io.log_info("%s, %s, %s, %s" %
("Src Match", len(src_match), "Src All", len(src_img_list)))
io.log_info("%s, %s, %s, %s" %
("Dst Match", len(dst_match), "Dst All", len(dst_img_list)))
# 画图
width = 800
xycr = []
for idx in range(len(src_img_list)):
t = src_img_list[idx]
if idx in src_match:
xycr.append([t[1], t[2], (128, 128, 128),
int(r * width / 2)]) # 蓝色，匹配到的
shutil.copy(t[0], src_aligned)
else:
xycr.append([t[1], t[2], (128, 128, 128), 2]) # 灰色，没匹配到
for idx in range(len(dst_img_list)):
t = dst_img_list[idx]
if idx in dst_match:
xycr.append([t[1], t[2], (0, 255, 0), 2]) # 绿色，保留
else:
xycr.append([t[1], t[2], (0, 0, 255), 2]) # 红色，删除
shutil.move(t[0], dst_aligned_trash)
img = cv.cv_new((width, width))
xs = [i[0] for i in xycr]
ys = [i[1] for i in xycr]
cs = [i[2] for i in xycr]
rs = [i[3] for i in xycr]
cv.cv_scatter(img, xs, ys, [-1, 1], [-1, 1], cs, rs)
cv.cv_save(img, os.path.join(dst_aligned, "_match_by_pitch.bmp"))
# 加入base
base_dir = os.path.join(data_src_path, "aligned_base")
if os.path.exists(base_dir):
for img in os.listdir(base_dir):
if img.endswith(".jpg") or img.endswith(".png"):
img_path = os.path.join(base_dir, img)
shutil.copy(img_path, src_aligned)