Spaces:
Running
Running
| import os | |
| import sys | |
| from PIL import Image | |
| from tqdm import tqdm | |
| import warnings | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) | |
| from models.detection_models import align # Assuming this is available in your project | |
| def extract_and_save_faces(input_dir, output_dir, algorithm='yolo', resolution=224): | |
| """Preprocess images using face alignment and cache them with specified resolution.""" | |
| if align is None: | |
| raise ImportError("face_alignment package is required for preprocessing.") | |
| os.makedirs(output_dir, exist_ok=True) | |
| with warnings.catch_warnings(): | |
| warnings.filterwarnings("ignore", category=FutureWarning, message=".*rcond.*") | |
| for person in sorted(os.listdir(input_dir)): | |
| person_path = os.path.join(input_dir, person) | |
| if not os.path.isdir(person_path): | |
| continue | |
| output_person_path = os.path.join(output_dir, person) | |
| os.makedirs(output_person_path, exist_ok=True) | |
| skipped_count = 0 | |
| for img_name in tqdm(os.listdir(person_path), desc=f"Processing {person}"): | |
| if not img_name.endswith(('.jpg', '.jpeg', '.png')): | |
| continue | |
| img_path = os.path.join(person_path, img_name) | |
| output_img_path = os.path.join(output_person_path, img_name) | |
| if os.path.exists(output_img_path): | |
| skipped_count += 1 | |
| continue | |
| try: | |
| aligned_result = align.get_aligned_face([img_path], algorithm=algorithm) | |
| aligned_image = aligned_result[0][1] if aligned_result and len(aligned_result) > 0 else None | |
| if aligned_image is None: | |
| print(f"Face detection failed for {img_path}, using resized original image") | |
| aligned_image = Image.open(img_path).convert('RGB') | |
| aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS) | |
| aligned_image.save(output_img_path, quality=100) | |
| except Exception as e: | |
| print(f"Error processing {img_path}: {e}") | |
| aligned_image = Image.open(img_path).convert('RGB') | |
| aligned_image = aligned_image.resize((resolution, resolution), Image.Resampling.LANCZOS) | |
| aligned_image.save(output_img_path, quality=100) | |
| if skipped_count > 0: | |
| print(f"Skipped {skipped_count} images for {person} that were already processed.") | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser(description="Preprocess and cache images with face alignment.") | |
| parser.add_argument('--input_dir', type=str, required=True, help='Directory containing raw images.') | |
| parser.add_argument('--output_dir', type=str, required=True, help='Directory to save preprocessed images.') | |
| parser.add_argument('--algorithm', type=str, default='yolo', choices=['yolo', 'dlib'], help='Face detection algorithm to use.') | |
| parser.add_argument('--resolution', type=int, default=224, help='Resolution for the output images.') | |
| args = parser.parse_args() | |
| extract_and_save_faces(args.input_dir, args.output_dir, args.algorithm, args.resolution) | |
| # python src/slimface/data/preprocess.py \ | |
| # --input_dir "data/raw/Original Images/Original Images" \ | |
| # --output_dir "data/processed/Aligned Images" \ | |
| # --algorithm "yolo" \ | |
| # --resolution 224 | |