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