# このスクリプトのライセンスは、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)