feat(sort-images): sort images by size or aspect ratio

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("<I", data[1:5])[0]
+                width = (bits & 0x3FFF) + 1
+                height = ((bits >> 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("<HH", data[6:10])
+                return width & 0x3FFF, height & 0x3FFF
+
+            # Skip chunk data + padding byte if length is odd
+            f.seek((length + 1) & ~1, 1)
+
+    raise ValueError(f"Could not find dimension chunks in {image_path}")
+
+
+def read_dims(image_path: Path) -> 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()