Spaces:
Build error
Build error
File size: 1,401 Bytes
6d314be |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import argparse
import os
import sys
from pathlib import Path
from PIL import Image
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from utils.seed import set_seed
from utils.image_utils import list_image_files
from utils.shape_predictor import align_face
def alignment_from_path(input_path, output_path, replace_cropped=False):
output_path.mkdir(parents=True, exist_ok=True)
image_files = list_image_files(input_path)
for img_path in image_files:
if (output_path / img_path).is_file() and not replace_cropped:
continue
image = Image.open(input_path / img_path)
# returns the largest face
crop_image = align_face(image, return_tensors=False)
if len(crop_image) == 1:
crop_image[0].save(output_path / img_path)
elif len(crop_image) > 1:
raise NotImplemented
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Align faces')
parser.add_argument('-unprocessed_dir', type=Path, default='unprocessed', help='directory with unprocessed images')
parser.add_argument('-output_dir', type=Path, default='input', help='output directory')
parser.add_argument('-replace_cropped', action='store_true')
args = parser.parse_args()
set_seed(3407)
alignment_from_path(args.unprocessed_dir, args.output_dir, args.replace_cropped)
|