Batch compression of images to Webp with specified quality

image-to-webp-compressor.py

import os
import signal
import argparse
from PIL import Image
from pathlib import Path
from tqdm import tqdm
from multiprocessing import Pool, cpu_count, current_process

SUPPORTED_FORMATS = (
    ".jpg",
    ".jpeg",
    ".png",
    ".bmp",
)

# Global flag for graceful shutdown
should_terminate = False


def init_worker():
    """Initialize worker process to ignore SIGINT"""
    signal.signal(signal.SIGINT, signal.SIG_IGN)


def signal_handler(signum, frame):
    """Handle interrupt signal"""
    global should_terminate
    if current_process().name == "MainProcess":
        print(
            "\nGracefully shutting down... Please wait for current operations to complete."
        )
        should_terminate = True


def convert_to_webp(args):
    """
    Converts an image to WebP format with specified quality.

    Args:
        args: tuple containing (input_path, output_path, quality)
    """
    if should_terminate:
        return None

    input_path, output_path, quality = args

    try:
        image = Image.open(input_path)

        # Convert to RGB if image has transparency
        if image.mode in ("RGBA", "LA"):
            background = Image.new("RGB", image.size, (255, 255, 255))
            background.paste(image, mask=image.split()[-1])
            image = background

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        image.save(output_path, "WEBP", quality=quality)

    except Exception as e:
        print(f"Error converting {input_path}: {str(e)}")
        return False

    return True


def process_directory(
    input_dir, output_dir, quality=65, num_processes=None, progress_update=1.0
):
    """
    Processes all images in the specified directory using multiple processes.

    Args:
        input_dir: path to directory with source images
        output_dir: path to directory for saving WebP files
        quality: compression quality (0-100)
        num_processes: number of processes to use (defaults to CPU count)
        progress_update: progress bar update interval in seconds (0 to disable, default: 1.0)
    """
    if num_processes is None:
        num_processes = cpu_count()

    Path(output_dir).mkdir(parents=True, exist_ok=True)

    image_files = [
        f for f in Path(input_dir).rglob("*") if f.suffix.lower() in SUPPORTED_FORMATS
    ]

    if not image_files:
        print(f"No supported images found in directory {input_dir}")
        return

    print(f"Found {len(image_files)} images for conversion")
    print(f"Using {num_processes} processes for parallel conversion")

    conversion_args = []
    for image_path in image_files:
        relative_path = image_path.relative_to(input_dir)
        output_path = Path(output_dir) / relative_path.with_suffix(".webp")
        conversion_args.append((str(image_path), str(output_path), quality))

    # Set up signal handling
    signal.signal(signal.SIGINT, signal_handler)

    with Pool(processes=num_processes, initializer=init_worker) as pool:
        try:
            if progress_update > 0:
                for _ in tqdm(
                    pool.imap_unordered(convert_to_webp, conversion_args),
                    total=len(conversion_args),
                    desc="Converting images",
                    mininterval=progress_update,
                ):
                    if should_terminate:
                        pool.terminate()
                        break
            else:
                for _ in pool.imap_unordered(convert_to_webp, conversion_args):
                    if should_terminate:
                        pool.terminate()
                        break

        except KeyboardInterrupt:
            pool.terminate()
        finally:
            pool.close()
            pool.join()

    if should_terminate:
        print(
            "\nConversion interrupted by user. Some images may not have been converted."
        )
    else:
        print("\nConversion completed successfully.")


def main():
    parser = argparse.ArgumentParser(
        description="Batch convert images to WebP format with parallel processing",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("input_dir", help="Source directory containing images")
    parser.add_argument("output_dir", help="Output directory for WebP files")
    parser.add_argument(
        "-q",
        "--quality",
        type=int,
        default=65,
        help="Compression quality (0-100, default: 65)",
    )
    parser.add_argument(
        "-p",
        "--processes",
        type=int,
        default=None,
        help="Number of processes to use (default: number of CPU cores)",
    )
    parser.add_argument(
        "-t",
        "--progress-update",
        type=float,
        default=1.0,
        help="Progress bar update interval in seconds (0 to disable, default: 1.0)",
    )

    args = parser.parse_args()

    if not os.path.exists(args.input_dir):
        print(f"Error: Directory {args.input_dir} does not exist")
        return

    if args.quality < 0 or args.quality > 100:
        print("Error: Quality must be between 0 and 100")
        return

    if args.processes is not None and args.processes < 1:
        print("Error: Number of processes must be at least 1")
        return

    process_directory(
        args.input_dir,
        args.output_dir,
        args.quality,
        args.processes,
        args.progress_update,
    )


if __name__ == "__main__":
    main()