| 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('') |