Spaces:
Runtime error
Runtime error
| import argparse | |
| import json | |
| import os | |
| from os.path import join as pjoin | |
| import sys | |
| import bz2 | |
| import numpy as np | |
| import cv2 | |
| from tqdm import tqdm | |
| from tensorflow.keras.utils import get_file | |
| from utils.ffhq_dataset.face_alignment import image_align | |
| from utils.ffhq_dataset.landmarks_detector import LandmarksDetector | |
| LANDMARKS_MODEL_URL = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2' | |
| 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 | |
| class SizePathMap(dict): | |
| """{size: {aligned_face_path0, aligned_face_path1, ...}, ...}""" | |
| def add_item(self, size, path): | |
| if size not in self: | |
| self[size] = set() | |
| self[size].add(path) | |
| def get_sizes(self): | |
| sizes = [] | |
| for key, paths in self.items(): | |
| sizes.extend([key,]*len(paths)) | |
| return sizes | |
| def serialize(self): | |
| result = {} | |
| for key, paths in self.items(): | |
| result[key] = list(paths) | |
| return result | |
| def main(args): | |
| landmarks_model_path = unpack_bz2(get_file('shape_predictor_68_face_landmarks.dat.bz2', | |
| LANDMARKS_MODEL_URL, cache_subdir='temp')) | |
| landmarks_detector = LandmarksDetector(landmarks_model_path) | |
| face_sizes = SizePathMap() | |
| raw_img_dir = args.raw_image_dir | |
| img_names = [n for n in os.listdir(raw_img_dir) if os.path.isfile(pjoin(raw_img_dir, n))] | |
| aligned_image_dir = args.aligned_image_dir | |
| os.makedirs(aligned_image_dir, exist_ok=True) | |
| pbar = tqdm(img_names) | |
| for img_name in pbar: | |
| pbar.set_description(img_name) | |
| if os.path.splitext(img_name)[-1] == '.txt': | |
| continue | |
| raw_img_path = os.path.join(raw_img_dir, img_name) | |
| try: | |
| for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1): | |
| face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], i) | |
| aligned_face_path = os.path.join(aligned_image_dir, face_img_name) | |
| face_size = image_align( | |
| raw_img_path, aligned_face_path, face_landmarks, resize=args.resize | |
| ) | |
| face_sizes.add_item(face_size, aligned_face_path) | |
| pbar.set_description(f"{img_name}: {face_size}") | |
| if args.draw: | |
| visual = LandmarksDetector.draw(cv2.imread(raw_img_path), face_landmarks) | |
| cv2.imwrite( | |
| pjoin(args.aligned_image_dir, os.path.splitext(face_img_name)[0] + "_landmarks.png"), | |
| visual | |
| ) | |
| except Exception as e: | |
| print('[Error]', e, 'error happened when processing', raw_img_path) | |
| print(args.raw_image_dir, ':') | |
| sizes = face_sizes.get_sizes() | |
| results = { | |
| 'mean_size': np.mean(sizes), | |
| 'num_faces_detected': len(sizes), | |
| 'num_images': len(img_names), | |
| 'sizes': sizes, | |
| 'size_path_dict': face_sizes.serialize(), | |
| } | |
| print('\t', results) | |
| if args.out_stats is not None: | |
| os.makedirs(os.path.dirname(args.out_stats), exist_ok=True) | |
| with open(out_stats, 'w') as f: | |
| json.dump(results, f) | |
| def parse_args(args=None, namespace=None): | |
| parser = argparse.ArgumentParser(description=""" | |
| 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.add_argument('raw_image_dir') | |
| parser.add_argument('aligned_image_dir') | |
| parser.add_argument('--resize', | |
| help="True if want to resize to 1024", | |
| action='store_true') | |
| parser.add_argument('--draw', | |
| help="True if want to visualize landmarks", | |
| action='store_true') | |
| parser.add_argument('--out_stats', | |
| help="output_fn for statistics of faces", default=None) | |
| return parser.parse_args(args=args, namespace=namespace) | |
| if __name__ == "__main__": | |
| main(parse_args()) | |