246 lines
8.4 KiB
Python
Executable File
246 lines
8.4 KiB
Python
Executable File
#!/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 ---
|
|
|
|
|
|
class ImageFormat(enum.Enum):
|
|
JPEG = "jpeg"
|
|
PNG = "png"
|
|
WEBP = "webp"
|
|
|
|
@property
|
|
def suffix(self) -> str:
|
|
return ".jpg" if self is self.JPEG else f".{self.value}"
|
|
|
|
|
|
def sniff_image_format(image_path: Path) -> ImageFormat:
|
|
"""Identify a supported image format from its file signature."""
|
|
with image_path.open("rb") as file:
|
|
header = file.read(12)
|
|
|
|
if header.startswith(b"\xff\xd8\xff"):
|
|
return ImageFormat.JPEG
|
|
if header.startswith(b"\x89PNG\r\n\x1a\n"):
|
|
return ImageFormat.PNG
|
|
if len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP":
|
|
return ImageFormat.WEBP
|
|
raise ValueError(f"Unsupported or unrecognized image format: {image_path}")
|
|
|
|
|
|
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]:
|
|
image_format = sniff_image_format(image_path)
|
|
if image_format is ImageFormat.JPEG:
|
|
return read_dims_jpeg(image_path)
|
|
elif image_format is ImageFormat.PNG:
|
|
return read_dims_png(image_path)
|
|
elif image_format is ImageFormat.WEBP:
|
|
return read_dims_webp(image_path)
|
|
raise AssertionError(f"Unhandled image format: {image_format}")
|
|
|
|
|
|
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
|
|
|
|
try:
|
|
image_format = sniff_image_format(image_path)
|
|
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(image_format.suffix)
|
|
|
|
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()
|