|
| 1 | +# Copyright 2020 NVIDIA Corporation |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +import argparse |
| 16 | +import json |
| 17 | +import numpy as np |
| 18 | +import os |
| 19 | +import PIL.Image |
| 20 | +import scipy.ndimage |
| 21 | +from tqdm import tqdm |
| 22 | + |
| 23 | +_examples = '''examples: |
| 24 | +
|
| 25 | + # Run x, y, z |
| 26 | + python %(prog)s --output=tmp |
| 27 | +''' |
| 28 | + |
| 29 | +def extract_face(face, source_images, output_dir, target_size=1024, supersampling=4, enable_padding=True): |
| 30 | + def rot90(v) -> np.ndarray: |
| 31 | + return np.array([-v[1], v[0]]) |
| 32 | + |
| 33 | + # Sanitize facial landmarks. |
| 34 | + face_spec = face['face_spec'] |
| 35 | + landmarks = (np.float32(face_spec['landmarks']) + 0.5) * face_spec['shrink'] |
| 36 | + assert landmarks.shape == (68, 2) |
| 37 | + lm_eye_left = landmarks[36 : 42] # left-clockwise |
| 38 | + lm_eye_right = landmarks[42 : 48] # left-clockwise |
| 39 | + lm_mouth_outer = landmarks[48 : 60] # left-clockwise |
| 40 | + |
| 41 | + # Calculate auxiliary vectors. |
| 42 | + eye_left = np.mean(lm_eye_left, axis=0) |
| 43 | + eye_right = np.mean(lm_eye_right, axis=0) |
| 44 | + eye_avg = (eye_left + eye_right) * 0.5 |
| 45 | + eye_to_eye = eye_right - eye_left |
| 46 | + mouth_left = lm_mouth_outer[0] |
| 47 | + mouth_right = lm_mouth_outer[6] |
| 48 | + mouth_avg = (mouth_left + mouth_right) * 0.5 |
| 49 | + eye_to_mouth = mouth_avg - eye_avg |
| 50 | + |
| 51 | + # Choose oriented crop rectangle. |
| 52 | + x = eye_to_eye - rot90(eye_to_mouth) |
| 53 | + x /= np.hypot(*x) |
| 54 | + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) |
| 55 | + y = rot90(x) |
| 56 | + c = eye_avg + eye_to_mouth * 0.1 |
| 57 | + |
| 58 | + # Calculate auxiliary data. |
| 59 | + qsize = np.hypot(*x) * 2 |
| 60 | + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) |
| 61 | + lo = np.min(quad, axis=0) |
| 62 | + hi = np.max(quad, axis=0) |
| 63 | + lm_rel = np.dot(landmarks - c, np.transpose([x, y])) / qsize**2 * 2 + 0.5 |
| 64 | + rp = np.dot(np.random.RandomState(123).uniform(-1, 1, size=(1024, 2)), [x, y]) + c |
| 65 | + |
| 66 | + # Load. |
| 67 | + img = PIL.Image.open(os.path.join(source_images, face['source_path'])).convert('RGB') |
| 68 | + |
| 69 | + # Shrink. |
| 70 | + shrink = int(np.floor(qsize / target_size * 0.5)) |
| 71 | + if shrink > 1: |
| 72 | + rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink))) |
| 73 | + img = img.resize(rsize, PIL.Image.ANTIALIAS) |
| 74 | + quad /= shrink |
| 75 | + qsize /= shrink |
| 76 | + |
| 77 | + # Crop. |
| 78 | + border = max(int(np.rint(qsize * 0.1)), 3) |
| 79 | + crop = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1])))) |
| 80 | + crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), min(crop[3] + border, img.size[1])) |
| 81 | + if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]: |
| 82 | + img = img.crop(crop) |
| 83 | + quad -= crop[0:2] |
| 84 | + |
| 85 | + # Pad. |
| 86 | + pad = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1])))) |
| 87 | + pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), max(pad[3] - img.size[1] + border, 0)) |
| 88 | + if enable_padding and max(pad) > border - 4: |
| 89 | + pad = np.maximum(pad, int(np.rint(qsize * 0.3))) |
| 90 | + img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect') |
| 91 | + h, w, _ = img.shape |
| 92 | + y, x, _ = np.ogrid[:h, :w, :1] |
| 93 | + mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w-1-x) / pad[2]), 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h-1-y) / pad[3])) |
| 94 | + blur = qsize * 0.02 |
| 95 | + img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0) |
| 96 | + img += (np.median(img, axis=(0,1)) - img) * np.clip(mask, 0.0, 1.0) |
| 97 | + img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB') |
| 98 | + quad += pad[:2] |
| 99 | + |
| 100 | + # Transform. |
| 101 | + super_size = target_size * supersampling |
| 102 | + img = img.transform((super_size, super_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR) |
| 103 | + if target_size < super_size: |
| 104 | + img = img.resize((target_size, target_size), PIL.Image.ANTIALIAS) |
| 105 | + |
| 106 | + # Save face image. |
| 107 | + img.save(os.path.join(output_dir, f"{face['obj_id']}-{face['face_idx']:02d}.png")) |
| 108 | + |
| 109 | + |
| 110 | +def main(): |
| 111 | + parser = argparse.ArgumentParser( |
| 112 | + description='MetFaces dataset processing tool', |
| 113 | + epilog=_examples, |
| 114 | + formatter_class=argparse.RawDescriptionHelpFormatter |
| 115 | + ) |
| 116 | + parser.add_argument('--json', help='MetFaces metadata json file path', required=True) |
| 117 | + parser.add_argument('--source-images', help='Location of MetFaces raw image data', required=True) |
| 118 | + parser.add_argument('--output-dir', help='Where to save output files') |
| 119 | + args = parser.parse_args() |
| 120 | + |
| 121 | + os.makedirs(args.output_dir, exist_ok=True) |
| 122 | + |
| 123 | + with open(args.json) as fin: |
| 124 | + faces = json.load(fin) |
| 125 | + for f in tqdm(faces): |
| 126 | + extract_face(f, source_images=args.source_images, output_dir=args.output_dir) |
| 127 | + |
| 128 | +if __name__ == "__main__": |
| 129 | + main() |
0 commit comments