247 lines
10 KiB
Python
247 lines
10 KiB
Python
# このスクリプトのライセンスは、train_dreambooth.pyと同じくApache License 2.0とします
|
||
# (c) 2022 Kohya S. @kohya_ss
|
||
|
||
# 横長の画像から顔検出して正立するように回転し、そこを中心に正方形に切り出す
|
||
|
||
# v2: extract max face if multiple faces are found
|
||
# v3: add crop_ratio option
|
||
# v4: add multiple faces extraction and min/max size
|
||
|
||
import argparse
|
||
import math
|
||
import cv2
|
||
import glob
|
||
import os
|
||
from anime_face_detector import create_detector
|
||
from tqdm import tqdm
|
||
import numpy as np
|
||
|
||
KP_REYE = 11
|
||
KP_LEYE = 19
|
||
|
||
SCORE_THRES = 0.90
|
||
|
||
|
||
def detect_faces(detector, image, min_size):
|
||
preds = detector(image) # bgr
|
||
# print(len(preds))
|
||
|
||
faces = []
|
||
for pred in preds:
|
||
bb = pred['bbox']
|
||
score = bb[-1]
|
||
if score < SCORE_THRES:
|
||
continue
|
||
|
||
left, top, right, bottom = bb[:4]
|
||
cx = int((left + right) / 2)
|
||
cy = int((top + bottom) / 2)
|
||
fw = int(right - left)
|
||
fh = int(bottom - top)
|
||
|
||
lex, ley = pred['keypoints'][KP_LEYE, 0:2]
|
||
rex, rey = pred['keypoints'][KP_REYE, 0:2]
|
||
angle = math.atan2(ley - rey, lex - rex)
|
||
angle = angle / math.pi * 180
|
||
|
||
faces.append((cx, cy, fw, fh, angle))
|
||
|
||
faces.sort(key=lambda x: max(x[2], x[3]), reverse=True) # 大きい順
|
||
return faces
|
||
|
||
|
||
def rotate_image(image, angle, cx, cy):
|
||
h, w = image.shape[0:2]
|
||
rot_mat = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
|
||
|
||
# # 回転する分、すこし画像サイズを大きくする→とりあえず無効化
|
||
# nh = max(h, int(w * math.sin(angle)))
|
||
# nw = max(w, int(h * math.sin(angle)))
|
||
# if nh > h or nw > w:
|
||
# pad_y = nh - h
|
||
# pad_t = pad_y // 2
|
||
# pad_x = nw - w
|
||
# pad_l = pad_x // 2
|
||
# m = np.array([[0, 0, pad_l],
|
||
# [0, 0, pad_t]])
|
||
# rot_mat = rot_mat + m
|
||
# h, w = nh, nw
|
||
# cx += pad_l
|
||
# cy += pad_t
|
||
|
||
result = cv2.warpAffine(image, rot_mat, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
|
||
return result, cx, cy
|
||
|
||
|
||
def process(args):
|
||
assert (not args.resize_fit) or args.resize_face_size is None, f"resize_fit and resize_face_size can't be specified both / resize_fitとresize_face_sizeはどちらか片方しか指定できません"
|
||
assert args.crop_ratio is None or args.resize_face_size is None, f"crop_ratio指定時はresize_face_sizeは指定できません"
|
||
|
||
# アニメ顔検出モデルを読み込む
|
||
print("loading face detector.")
|
||
detector = create_detector('yolov3')
|
||
|
||
# cropの引数を解析する
|
||
if args.crop_size is None:
|
||
crop_width = crop_height = None
|
||
else:
|
||
tokens = args.crop_size.split(',')
|
||
assert len(tokens) == 2, f"crop_size must be 'width,height' / crop_sizeは'幅,高さ'で指定してください"
|
||
crop_width, crop_height = [int(t) for t in tokens]
|
||
|
||
if args.crop_ratio is None:
|
||
crop_h_ratio = crop_v_ratio = None
|
||
else:
|
||
tokens = args.crop_ratio.split(',')
|
||
assert len(tokens) == 2, f"crop_ratio must be 'horizontal,vertical' / crop_ratioは'幅,高さ'の倍率で指定してください"
|
||
crop_h_ratio, crop_v_ratio = [float(t) for t in tokens]
|
||
|
||
# 画像を処理する
|
||
print("processing.")
|
||
output_extension = ".png"
|
||
|
||
os.makedirs(args.dst_dir, exist_ok=True)
|
||
paths = glob.glob(os.path.join(args.src_dir, "*.png")) + glob.glob(os.path.join(args.src_dir, "*.jpg")) + \
|
||
glob.glob(os.path.join(args.src_dir, "*.webp"))
|
||
for path in tqdm(paths):
|
||
basename = os.path.splitext(os.path.basename(path))[0]
|
||
|
||
# image = cv2.imread(path) # 日本語ファイル名でエラーになる
|
||
image = cv2.imdecode(np.fromfile(path, np.uint8), cv2.IMREAD_UNCHANGED)
|
||
if len(image.shape) == 2:
|
||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
||
if image.shape[2] == 4:
|
||
print(f"image has alpha. ignore / 画像の透明度が設定されているため無視します: {path}")
|
||
image = image[:, :, :3].copy() # copyをしないと内部的に透明度情報が付いたままになるらしい
|
||
|
||
h, w = image.shape[:2]
|
||
|
||
faces = detect_faces(detector, image, args.multiple_faces)
|
||
for i, face in enumerate(faces):
|
||
cx, cy, fw, fh, angle = face
|
||
face_size = max(fw, fh)
|
||
if args.min_size is not None and face_size < args.min_size:
|
||
continue
|
||
if args.max_size is not None and face_size >= args.max_size:
|
||
continue
|
||
face_suffix = f"_{i+1:02d}" if args.multiple_faces else ""
|
||
|
||
# オプション指定があれば回転する
|
||
face_img = image
|
||
if args.rotate:
|
||
face_img, cx, cy = rotate_image(face_img, angle, cx, cy)
|
||
|
||
# オプション指定があれば顔を中心に切り出す
|
||
if crop_width is not None or crop_h_ratio is not None:
|
||
cur_crop_width, cur_crop_height = crop_width, crop_height
|
||
if crop_h_ratio is not None:
|
||
cur_crop_width = int(face_size * crop_h_ratio + .5)
|
||
cur_crop_height = int(face_size * crop_v_ratio + .5)
|
||
|
||
# リサイズを必要なら行う
|
||
scale = 1.0
|
||
if args.resize_face_size is not None:
|
||
# 顔サイズを基準にリサイズする
|
||
scale = args.resize_face_size / face_size
|
||
if scale < cur_crop_width / w:
|
||
print(
|
||
f"image width too small in face size based resizing / 顔を基準にリサイズすると画像の幅がcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
|
||
scale = cur_crop_width / w
|
||
if scale < cur_crop_height / h:
|
||
print(
|
||
f"image height too small in face size based resizing / 顔を基準にリサイズすると画像の高さがcrop sizeより小さい（顔が相対的に大きすぎる）ので顔サイズが変わります: {path}")
|
||
scale = cur_crop_height / h
|
||
elif crop_h_ratio is not None:
|
||
# 倍率指定の時にはリサイズしない
|
||
pass
|
||
else:
|
||
# 切り出しサイズ指定あり
|
||
if w < cur_crop_width:
|
||
print(f"image width too small/ 画像の幅がcrop sizeより小さいので画質が劣化します: {path}")
|
||
scale = cur_crop_width / w
|
||
if h < cur_crop_height:
|
||
print(f"image height too small/ 画像の高さがcrop sizeより小さいので画質が劣化します: {path}")
|
||
scale = cur_crop_height / h
|
||
if args.resize_fit:
|
||
scale = max(cur_crop_width / w, cur_crop_height / h)
|
||
|
||
if scale != 1.0:
|
||
w = int(w * scale + .5)
|
||
h = int(h * scale + .5)
|
||
face_img = cv2.resize(face_img, (w, h), interpolation=cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LANCZOS4)
|
||
cx = int(cx * scale + .5)
|
||
cy = int(cy * scale + .5)
|
||
fw = int(fw * scale + .5)
|
||
fh = int(fh * scale + .5)
|
||
|
||
cur_crop_width = min(cur_crop_width, face_img.shape[1])
|
||
cur_crop_height = min(cur_crop_height, face_img.shape[0])
|
||
|
||
x = cx - cur_crop_width // 2
|
||
cx = cur_crop_width // 2
|
||
if x < 0:
|
||
cx = cx + x
|
||
x = 0
|
||
elif x + cur_crop_width > w:
|
||
cx = cx + (x + cur_crop_width - w)
|
||
x = w - cur_crop_width
|
||
face_img = face_img[:, x:x+cur_crop_width]
|
||
|
||
y = cy - cur_crop_height // 2
|
||
cy = cur_crop_height // 2
|
||
if y < 0:
|
||
cy = cy + y
|
||
y = 0
|
||
elif y + cur_crop_height > h:
|
||
cy = cy + (y + cur_crop_height - h)
|
||
y = h - cur_crop_height
|
||
face_img = face_img[y:y + cur_crop_height]
|
||
|
||
# # debug
|
||
# print(path, cx, cy, angle)
|
||
# crp = cv2.resize(image, (image.shape[1]//8, image.shape[0]//8))
|
||
# cv2.imshow("image", crp)
|
||
# if cv2.waitKey() == 27:
|
||
# break
|
||
# cv2.destroyAllWindows()
|
||
|
||
# debug
|
||
if args.debug:
|
||
cv2.rectangle(face_img, (cx-fw//2, cy-fh//2), (cx+fw//2, cy+fh//2), (255, 0, 255), fw//20)
|
||
|
||
_, buf = cv2.imencode(output_extension, face_img)
|
||
with open(os.path.join(args.dst_dir, f"{basename}{face_suffix}_{cx:04d}_{cy:04d}_{fw:04d}_{fh:04d}{output_extension}"), "wb") as f:
|
||
buf.tofile(f)
|
||
|
||
|
||
def setup_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--src_dir", type=str, help="directory to load images / 画像を読み込むディレクトリ")
|
||
parser.add_argument("--dst_dir", type=str, help="directory to save images / 画像を保存するディレクトリ")
|
||
parser.add_argument("--rotate", action="store_true", help="rotate images to align faces / 顔が正立するように画像を回転する")
|
||
parser.add_argument("--resize_fit", action="store_true",
|
||
help="resize to fit smaller side after cropping / 切り出し後の画像の短辺がcrop_sizeにあうようにリサイズする")
|
||
parser.add_argument("--resize_face_size", type=int, default=None,
|
||
help="resize image before cropping by face size / 切り出し前に顔がこのサイズになるようにリサイズする")
|
||
parser.add_argument("--crop_size", type=str, default=None,
|
||
help="crop images with 'width,height' pixels, face centered / 顔を中心として'幅,高さ'のサイズで切り出す")
|
||
parser.add_argument("--crop_ratio", type=str, default=None,
|
||
help="crop images with 'horizontal,vertical' ratio to face, face centered / 顔を中心として顔サイズの'幅倍率,高さ倍率'のサイズで切り出す")
|
||
parser.add_argument("--min_size", type=int, default=None,
|
||
help="minimum face size to output (included) / 処理対象とする顔の最小サイズ（この値以上）")
|
||
parser.add_argument("--max_size", type=int, default=None,
|
||
help="maximum face size to output (excluded) / 処理対象とする顔の最大サイズ（この値未満）")
|
||
parser.add_argument("--multiple_faces", action="store_true",
|
||
help="output each faces / 複数の顔が見つかった場合、それぞれを切り出す")
|
||
parser.add_argument("--debug", action="store_true", help="render rect for face / 処理後画像の顔位置に矩形を描画します")
|
||
|
||
return parser
|
||
|
||
|
||
if __name__ == '__main__':
|
||
parser = setup_parser()
|
||
|
||
args = parser.parse_args()
|
||
|
||
process(args)
|