Spaces:
Running
Running
File size: 2,192 Bytes
e425192 ab4ea38 e425192 ab4ea38 e425192 ab4ea38 e425192 ab4ea38 e425192 ab4ea38 |
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 45 46 |
import os
from src.utils.face_alignment import image_align
from src.landmark_detector import LandmarksDetector
def align_images(img_files, aligned_dir, output_size=1024, x_scale=1, y_scale=1, em_scale=0.1, use_alpha=False):
"""
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
"""
ALIGNED_IMAGES_DIR = aligned_dir
# Create the directory if it doesn't exist
if not os.path.exists(ALIGNED_IMAGES_DIR):
os.makedirs(ALIGNED_IMAGES_DIR)
else: # Remove existing files in the directory
for file in os.listdir(ALIGNED_IMAGES_DIR):
os.remove(os.path.join(ALIGNED_IMAGES_DIR, file))
landmarks_detector = LandmarksDetector()
for idx, img_path in enumerate(img_files, start=1):
img_name = os.path.basename(img_path)
print('Aligning %s ...' % img_name)
try:
raw_img_path = img_path
fn = face_img_name = f'{idx:04d}_{os.path.splitext(img_name)[0]}_01.png'
if os.path.isfile(os.path.join(ALIGNED_IMAGES_DIR, 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 = f'{idx:04d}_{os.path.splitext(img_name)[0]}_{i:02d}.png'
aligned_face_path = os.path.join(ALIGNED_IMAGES_DIR, face_img_name)
image_align(raw_img_path, aligned_face_path, face_landmarks, output_size=output_size, x_scale=x_scale, y_scale=y_scale, em_scale=em_scale, alpha=use_alpha)
print('Wrote result %s\n' % aligned_face_path)
except Exception as e:
print(f"Exception in face alignment for {img_name}: {e}")
except Exception as e:
print(f"Exception in landmark detection for {img_name}: {e}")
# return absolute paths of aligned images
return [os.path.join(ALIGNED_IMAGES_DIR, img) for img in os.listdir(ALIGNED_IMAGES_DIR)]
|