1 Star 0 Fork 0

安家/BabyGAN

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
该仓库未声明开源许可证文件(LICENSE),使用请关注具体项目描述及其代码上游依赖。
克隆/下载
align_images.py 2.79 KB
一键复制 编辑 原始数据 按行查看 历史
Malimonov Denis 提交于 2020-08-25 02:39 . first commit
import os
import sys
import bz2
import argparse
from keras.utils import get_file
from ffhq_dataset.face_alignment import image_align
from ffhq_dataset.landmarks_detector import LandmarksDetector
import multiprocessing
def unpack_bz2(src_path):
data = bz2.BZ2File(src_path).read()
dst_path = src_path[:-4]
with open(dst_path, 'wb') as fp:
fp.write(data)
return dst_path
if __name__ == "__main__":
"""
Extracts and aligns all faces from images using DLib and a function from original FFHQ dataset preparation step
python align_images.py /raw_images /aligned_images
"""
parser = argparse.ArgumentParser(description='Align faces from input images', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('raw_dir', help='Directory with raw images for face alignment')
parser.add_argument('aligned_dir', help='Directory for storing aligned images')
parser.add_argument('--output_size', default=1024, help='The dimension of images for input to the model', type=int)
parser.add_argument('--x_scale', default=1, help='Scaling factor for x dimension', type=float)
parser.add_argument('--y_scale', default=1, help='Scaling factor for y dimension', type=float)
parser.add_argument('--em_scale', default=0.1, help='Scaling factor for eye-mouth distance', type=float)
parser.add_argument('--use_alpha', default=False, help='Add an alpha channel for masking', type=bool)
args, other_args = parser.parse_known_args()
landmarks_model_path = unpack_bz2("shape_predictor_68_face_landmarks.dat.bz2")
RAW_IMAGES_DIR = args.raw_dir
ALIGNED_IMAGES_DIR = args.aligned_dir
landmarks_detector = LandmarksDetector(landmarks_model_path)
for img_name in os.listdir(RAW_IMAGES_DIR):
print('Aligning %s ...' % img_name)
try:
raw_img_path = os.path.join(RAW_IMAGES_DIR, img_name)
fn = face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], 1)
if os.path.isfile(fn):
continue
print('Getting landmarks...')
for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1):
try:
print('Starting face alignment...')
face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], i)
aligned_face_path = os.path.join(ALIGNED_IMAGES_DIR, face_img_name)
image_align(raw_img_path, aligned_face_path, face_landmarks, output_size=args.output_size, x_scale=args.x_scale, y_scale=args.y_scale, em_scale=args.em_scale, alpha=args.use_alpha)
print('Wrote result %s' % aligned_face_path)
except:
print("Exception in face alignment!")
except:
print("Exception in landmark detection!")
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/anjia/BabyGAN.git
git@gitee.com:anjia/BabyGAN.git
anjia
BabyGAN
BabyGAN
master

搜索帮助

23e8dbc6 1850385 7e0993f3 1850385