File size: 1,015 Bytes
2d7efb8 |
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 |
from configs import paths_config
import dlib
import glob
import os
from tqdm import tqdm
from utils.alignment import align_face
def pre_process_images(raw_images_path):
current_directory = os.getcwd()
IMAGE_SIZE = 256
predictor = dlib.shape_predictor(paths_config.dlib)
os.chdir(raw_images_path)
images_names = glob.glob(f'*')
aligned_images = []
for image_name in tqdm(images_names):
try:
aligned_image = align_face(filepath=f'{raw_images_path}/{image_name}',
predictor=predictor, output_size=IMAGE_SIZE)
aligned_images.append(aligned_image)
except Exception as e:
print(e)
os.makedirs(paths_config.input_data_path, exist_ok=True)
for image, name in zip(aligned_images, images_names):
real_name = name.split('.')[0]
image.save(f'{paths_config.input_data_path}/{real_name}.jpg')
os.chdir(current_directory)
if __name__ == "__main__":
pre_process_images('') |