diff --git a/sort_images.py b/sort_images.py new file mode 100755 index 0000000..039e224 --- /dev/null +++ b/sort_images.py @@ -0,0 +1,224 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.13" +# dependencies = ["pillow"] +# /// +import argparse +import enum +from functools import cache +import logging +from pathlib import Path +import struct +from typing import Self +from PIL import Image +# --- Utility Functions --- + + +def read_dims_jpeg(image_path: Path) -> tuple[int, int]: + with image_path.open("rb") as file: + if file.read(2) != b"\xff\xd8": + raise ValueError(f"{image_path} is not a valid JPEG file") + + file.seek(0) + try: + img = Image.open(file) + width, height = img.size + return width, height + except Exception: + file.seek(0) + while True: + marker = file.read(1) + if not marker or marker != b"\xff": + raise ValueError(f"Invalid JPEG format in {image_path}") + marker_type = int.from_bytes(file.read(1), byteorder="big") + length = int.from_bytes(file.read(2), byteorder="big") - 2 + is_sof = 0xC0 <= marker_type <= 0xCF and marker_type not in (0xC4, 0xC8, 0xCC) + if is_sof: + file.seek(1, 1) + height = int.from_bytes(file.read(2), byteorder="big") + width = int.from_bytes(file.read(2), byteorder="big") + return width, height + file.seek(length, 1) + + +def read_dims_png(image_path: Path) -> tuple[int, int]: + """ + Reads PNG dimensions by first attempting PIL, then falling back to manual + IHDR parsing if the library fails to initialize. + """ + with image_path.open("rb") as file: + signature = file.read(8) + if signature != b"\x89PNG\r\n\x1a\n": + raise ValueError(f"{image_path} is not a valid PNG file") + + try: + file.seek(0) + with Image.open(file) as img: + return img.size + except Exception: + # PNG IHDR is always the first chunk. + # Offset 8: Chunk Length (4 bytes) + # Offset 12: Chunk Type "IHDR" (4 bytes) + # Offset 16: Width (4 bytes) + # Offset 20: Height (4 bytes) + file.seek(12) + if file.read(4) != b"IHDR": + raise ValueError(f"IHDR chunk not found in {image_path}") + + # Using struct for clean fixed-width big-endian unpacking + dims = file.read(8) + if len(dims) < 8: + raise ValueError(f"Truncated IHDR in {image_path}") + + width, height = struct.unpack(">II", dims) + return width, height + + +def read_dims_webp(image_path: Path) -> tuple[int, int]: + """ + Parses WebP dimensions directly from the RIFF container without + external dependencies. Handles VP8, VP8L, and VP8X bitstreams. + """ + with image_path.open("rb") as f: + header = f.read(12) + if len(header) < 12 or header[:4] != b"RIFF" or header[8:12] != b"WEBP": + raise ValueError(f"Not a valid WebP file: {image_path}") + + while True: + chunk_header = f.read(8) + if len(chunk_header) < 8: + break + + tag, length = struct.unpack("<4sI", chunk_header) + + if tag == b"VP8X": + # Extended Format: Canvas width/height are 24-bit integers + # starting at offset 4 of the chunk data. + data = f.read(10) + width = (int.from_bytes(data[4:7], "little") & 0xFFFFFF) + 1 + height = (int.from_bytes(data[7:10], "little") & 0xFFFFFF) + 1 + return width, height + + elif tag == b"VP8L": + # Lossless Format: 1 byte signature (0x2f), 14 bits width-1, + # 14 bits height-1, 1 bit alpha, 3 bits version. + data = f.read(5) + if data[0] != 0x2F: + raise ValueError("Invalid VP8L signature") + + # Unpack the 4 bytes following the signature + bits = struct.unpack("> 14) & 0x3FFF) + 1 + return width, height + + elif tag == b"VP8 ": + # Lossy Format: Sync code 0x9d012a starts at offset 3. + # Width/Height are 16-bit values (14 bits used) at offset 6. + data = f.read(10) + if data[3:6] != b"\x9d\x01\x2a": + raise ValueError("Invalid VP8 sync code") + + width, height = struct.unpack(" tuple[int, int]: + if image_path.suffix.lower() in [".jpg", ".jpeg"]: + return read_dims_jpeg(image_path) + elif image_path.suffix.lower() == ".png": + return read_dims_png(image_path) + elif image_path.suffix.lower() == ".webp": + return read_dims_webp(image_path) + else: + raise ValueError(f"Unsupported image format: {image_path.suffix}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Sort images based on size, aspect ratio, pose centering, or gaze direction.") + parser.add_argument("image_paths", nargs="+", type=Path, help="Path(s) to JPEG image(s)") + strategy = parser.add_mutually_exclusive_group(required=True) + strategy.add_argument("--by-size", action="store_true", help="Sort images by size") + strategy.add_argument("--by-aspect-ratio", action="store_true", help="Sort images by aspect ratio") + + return parser.parse_args() + + +class AspectRatio(enum.Enum): + PORTRAIT = "portrait" + WIDE = "wide" + + @classmethod + def from_dimensions(cls, width: int, height: int) -> Self: + return cls.PORTRAIT if height > width else cls.WIDE + + +class Size(enum.Enum): + SMALL = "small" + MEDIUM = "medium" + LARGE = "large" + + @classmethod + def from_dimensions(cls, width: int, height: int) -> Self: + smaller = min(width, height) + if smaller < 1000: + return cls.SMALL + elif smaller < 2000: + return cls.MEDIUM + return cls.LARGE + + +@cache +def make_dir(path: Path) -> None: + path.mkdir(parents=True, exist_ok=True) + + +def main(): + """Main function to process images.""" + args = parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") + + for image_path in args.image_paths: + if not image_path.is_file(): + logging.warning(f"{image_path} is not a file. Skipping.") + continue + + if image_path.suffix.lower() not in [".jpg", ".jpeg", ".png", ".webp"]: + logging.warning(f"{image_path} is not a JPEG, PNG, or WebP file. Skipping.") + continue + + try: + width, height = read_dims(image_path) + target_dir = None + + if args.by_size: + size_category = Size.from_dimensions(width, height) + target_dir = image_path.parent / f"_size_{size_category.value}" + elif args.by_aspect_ratio: + aspect_ratio_category = AspectRatio.from_dimensions(width, height) + target_dir = image_path.parent / f"_aspect_{aspect_ratio_category.value}" + else: + raise ValueError("No sorting strategy specified.") + + make_dir(target_dir) + + target_path = target_dir / image_path.name + target_path = target_path.with_suffix(".jpg") + + image_path.rename(target_path) + logging.info(f"Moved {image_path} to {target_path}") + + except KeyboardInterrupt: + raise + except ValueError: + logging.exception(f"Error processing {image_path}. Skipping.") + continue + + +if __name__ == "__main__": + main()