Spaces:
Running
Running
File size: 3,622 Bytes
b7f710c |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
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
|