Crop all images in a given directory to the same size

crop_images.py

"""
Crop all images in a given directory to the same size
"""

import argparse
import os
import concurrent.futures
from glob import glob
from enum import Enum

from PIL import Image

class Align(Enum):
    TOP_LEFT = "top-left"
    TOP_CENTER = "top-center"
    TOP_RIGHT = "top-right"
    CENTER_LEFT = "center-left"
    CENTER_CENTER = "center-center"
    CENTER_RIGHT = "center-right"
    BOTTOM_LEFT = "bottom-left"
    BOTTOM_CENTER = "bottom-center"
    BOTTOM_RIGHT = "bottom-right"


def crop_image(image_path: str, aspect_ratio: str, align: str, output_dir: str, verbose: bool = False):
    """
    Crop the image based on the given aspect ratio and alignment.
    """
    if verbose:
        print(f"Processing {image_path} with aspect ratio {aspect_ratio} and alignment {align}...")

    with Image.open(image_path) as img:
        # Calculate new dimensions
        width, height = img.size
        target_aspect_ratio = float(aspect_ratio.split(":")[0]) / float(aspect_ratio.split(":")[1])
        if width / height > target_aspect_ratio:
            new_width = int(height * target_aspect_ratio)
            new_height = height
        else:
            new_width = width
            new_height = int(width / target_aspect_ratio)

        # Find crop box based on alignment
        left = (width - new_width) // 2
        top = (height - new_height) // 2
        if "top" in align:
            top = 0
        elif "bottom" in align:
            top = height - new_height
        if "left" in align:
            left = 0
        elif "right" in align:
            left = width - new_width

        crop_box = (left, top, left + new_width, top + new_height)
        cropped_img = img.crop(crop_box)

        # Save cropped image
        output_path = os.path.join(output_dir, os.path.basename(image_path))
        os.makedirs(output_dir, exist_ok=True)
        cropped_img.save(output_path)

        if verbose:
            print(f"Saved cropped image to {output_path}")


def parse_args():
    """
    Parse command line arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("input_path", type=str, help="The input image. Supports glob patterns.")
    parser.add_argument("-o", "--output_dir", type=str, default="./outputs/",
                        help="The output directory to save images to.")
    parser.add_argument("-R", "--aspect-ratio", type=str, default="1:1",
                        help="The aspect ratio to crop images to.")
    parser.add_argument("-A", "--align", type=str, default="center-center",
                        help=f"The alignment of the crop. Options are: {', '.join([align.value for align in Align])}")
    parser.add_argument("-v", "--verbose", action="store_true", help="Print verbose output.")
    return parser.parse_args()


def main():
    """
    Process all images concurrently
    """
    args = parse_args()
    image_paths = glob(args.input_path)

    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = [
            executor.submit(crop_image, image_path, args.aspect_ratio, args.align, args.output_dir, args.verbose)
            for image_path in image_paths
        ]
        
        # Optionally wait for all futures to complete
        for future in concurrent.futures.as_completed(futures):
            try:
                future.result()  # Will raise an exception if the callable raised one
            except Exception as e:
                print(f"An error occurred: {e}")


if __name__ == "__main__":
    main()