feat(sort-images): sort images by size or aspect ratio
This commit is contained in:
Executable
+224
@@ -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()
|
||||
Reference in New Issue
Block a user