feat: Add thumbnail contact sheet generator
This commit is contained in:
abdus committed
1 parent
e12feae399
commit
1939ef0c0d
2 files changed
+889
No files matched your search
Executable
+413
@@ -0,0 +1,413 @@
|
||||
#!/usr/bin/env -S uv run --script
|
||||
# /// script
|
||||
# dependencies = ["Pillow"]
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from io import BytesIO
|
||||
from multiprocessing import get_context
|
||||
from pathlib import Path
|
||||
from time import perf_counter
|
||||
from typing import BinaryIO, Callable
|
||||
from zipfile import BadZipFile, ZipFile, ZipInfo, is_zipfile
|
||||
|
||||
from PIL import Image, ImageColor, ImageOps
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
worker_archive: ZipFile | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Source:
|
||||
name: str
|
||||
path: Path | None = None
|
||||
member: ZipInfo | None = None
|
||||
ratio: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class Row:
|
||||
sources: list[Source]
|
||||
height: int
|
||||
widths: list[int]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cell:
|
||||
source: Source
|
||||
position: tuple[int, int]
|
||||
size: tuple[int, int]
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Create a separate justified image contact sheet for each directory or ZIP.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument("inputs", nargs="+", type=Path, help="Image directories or ZIP archives; missing paths are skipped")
|
||||
parser.add_argument("--output", type=Path, help="JPEG output for a single existing input; defaults to <input-name>_thumbs.jpg beside each input")
|
||||
parser.add_argument("--width", type=int, default=2500, help="Sheet width in pixels")
|
||||
parser.add_argument("--target-height", type=int, default=300, help="Preferred row height in pixels; rows may be taller or shorter")
|
||||
parser.add_argument("--gap", type=int, default=8, help="Gap between images and at sheet edges in pixels")
|
||||
parser.add_argument("--bg-color", default="white", help="Background color as a color name or hex value")
|
||||
parser.add_argument("--workers", type=int, default=os.cpu_count() or 1, help="Number of resize worker processes")
|
||||
arguments = parser.parse_args()
|
||||
if arguments.target_height <= 0 or arguments.width <= 0 or arguments.gap < 0:
|
||||
parser.error("width and target height must be positive; gap must be nonnegative")
|
||||
if arguments.width <= 2 * arguments.gap:
|
||||
parser.error("width must exceed twice the gap")
|
||||
if arguments.workers <= 0:
|
||||
parser.error("workers must be positive")
|
||||
try:
|
||||
ImageColor.getcolor(color=arguments.bg_color, mode="RGB")
|
||||
except ValueError as error:
|
||||
parser.error(f"Invalid background color: {error}")
|
||||
return arguments
|
||||
|
||||
|
||||
def read_image(source: Source, archive: ZipFile | None) -> Image.Image:
|
||||
with ExitStack() as stack:
|
||||
if source.member is not None:
|
||||
if archive is None:
|
||||
raise ValueError("ZIP source requires an open archive")
|
||||
stream = stack.enter_context(archive.open(name=source.member))
|
||||
original = stack.enter_context(Image.open(fp=stream))
|
||||
else:
|
||||
original = stack.enter_context(Image.open(fp=source.path))
|
||||
oriented = ImageOps.exif_transpose(image=original)
|
||||
result = oriented.convert(mode="RGBA")
|
||||
if oriented is not original:
|
||||
oriented.close()
|
||||
return result
|
||||
|
||||
|
||||
def read_dims_jpeg(stream: BinaryIO) -> tuple[int, int]:
|
||||
EXIF_ORIENTATION_TAG = 274
|
||||
ORIENTATIONS_SWAPPING_WIDTH_HEIGHT = {5, 6, 7, 8}
|
||||
with Image.open(fp=stream) as image:
|
||||
width, height = image.size
|
||||
if image.getexif().get(EXIF_ORIENTATION_TAG) in ORIENTATIONS_SWAPPING_WIDTH_HEIGHT:
|
||||
return height, width
|
||||
return width, height
|
||||
|
||||
|
||||
def read_dims(image_path: Path | BinaryIO) -> tuple[int, int]:
|
||||
EXIF_ORIENTATION_TAG = 274
|
||||
ORIENTATIONS_SWAPPING_WIDTH_HEIGHT = {5, 6, 7, 8}
|
||||
with ExitStack() as stack:
|
||||
stream = stack.enter_context(image_path.open(mode="rb")) if isinstance(image_path, Path) else image_path
|
||||
signature = stream.read(12)
|
||||
stream.seek(0)
|
||||
if signature.startswith(b"\xff\xd8"):
|
||||
return read_dims_jpeg(stream=stream)
|
||||
if signature.startswith(b"\x89PNG\r\n\x1a\n"):
|
||||
raise ValueError("PNG images are not supported")
|
||||
if signature[:4] == b"RIFF" and signature[8:12] == b"WEBP":
|
||||
raise ValueError("WebP images are not supported")
|
||||
with Image.open(fp=stream) as image:
|
||||
if image.format == "TIFF":
|
||||
raise ValueError("TIFF images are not supported")
|
||||
width, height = image.size
|
||||
if image.getexif().get(EXIF_ORIENTATION_TAG) in ORIENTATIONS_SWAPPING_WIDTH_HEIGHT:
|
||||
return height, width
|
||||
return width, height
|
||||
|
||||
|
||||
def natural_sort_key(name: str) -> tuple[str | int, ...]:
|
||||
return tuple(int(part) if index % 2 else part.casefold() for index, part in enumerate(re.split(pattern=r"([0-9]+)", string=name)))
|
||||
|
||||
|
||||
def discover_sources(input_path: Path, output_path: Path, archive: ZipFile | None) -> list[Source]:
|
||||
Image.init()
|
||||
extensions = {extension for extension, image_format in Image.registered_extensions().items() if image_format not in {"PNG", "WEBP", "TIFF"}}
|
||||
if archive is not None:
|
||||
candidates = [
|
||||
Source(name=member.filename, member=member)
|
||||
for member in archive.infolist()
|
||||
if not member.is_dir() and Path(member.filename).suffix.lower() in extensions
|
||||
]
|
||||
else:
|
||||
candidates = [
|
||||
Source(name=path.relative_to(input_path).as_posix(), path=path)
|
||||
for path in input_path.rglob("*")
|
||||
if path.is_file() and path.suffix.lower() in extensions and path.resolve() != output_path.resolve()
|
||||
]
|
||||
sources = []
|
||||
logging.info(f"Reading metadata for {len(candidates)} images")
|
||||
for index, source in enumerate(sorted(candidates, key=lambda source: (natural_sort_key(name=source.name), source.name)), start=1):
|
||||
try:
|
||||
if source.member is not None and archive is not None:
|
||||
with archive.open(name=source.member) as stream:
|
||||
width, height = read_dims(image_path=stream)
|
||||
elif source.path is not None:
|
||||
width, height = read_dims(image_path=source.path)
|
||||
else:
|
||||
raise ValueError("Image source has no file or ZIP member")
|
||||
source.ratio = width / height
|
||||
sources.append(source)
|
||||
except (OSError, ValueError, SyntaxError, RuntimeError, BadZipFile, Image.DecompressionBombError) as error:
|
||||
logging.warning(f"Skipping {source.name}: {error}")
|
||||
if index % 100 == 0 or index == len(candidates):
|
||||
logging.info(f"Metadata {index}/{len(candidates)} ({index / len(candidates):.0%}); {len(sources)} usable images")
|
||||
if not sources:
|
||||
raise ValueError("No usable images found")
|
||||
return sources
|
||||
|
||||
|
||||
def fitted_height(sources: list[Source], width: int, gap: int) -> float:
|
||||
available = width - gap * (len(sources) + 1)
|
||||
return available / sum(source.ratio for source in sources)
|
||||
|
||||
|
||||
def make_row(sources: list[Source], width: int, gap: int, target_height: int, justify: bool) -> Row:
|
||||
available = width - gap * (len(sources) + 1)
|
||||
ideal_height = fitted_height(sources=sources, width=width, gap=gap)
|
||||
height = max(1, round(ideal_height if justify else min(target_height, ideal_height)))
|
||||
ratio_sum = sum(source.ratio for source in sources)
|
||||
total_width = available if justify else min(available, max(len(sources), round(ratio_sum * height)))
|
||||
widths = []
|
||||
cumulative = 0.0
|
||||
previous = 0
|
||||
for index, source in enumerate(sources):
|
||||
cumulative += source.ratio
|
||||
boundary = min(
|
||||
total_width - (len(sources) - index - 1),
|
||||
max(previous + 1, round(total_width * cumulative / ratio_sum)),
|
||||
)
|
||||
widths.append(boundary - previous)
|
||||
previous = boundary
|
||||
return Row(sources=sources, height=height, widths=widths)
|
||||
|
||||
|
||||
def layout_rows(sources: list[Source], width: int, gap: int, target_height: int) -> list[Row]:
|
||||
rows = []
|
||||
start = 0
|
||||
while start < len(sources):
|
||||
end = start + 1
|
||||
while end < len(sources) and fitted_height(sources=sources[start:end], width=width, gap=gap) > target_height:
|
||||
candidate = sources[start : end + 1]
|
||||
if width - gap * (len(candidate) + 1) < len(candidate):
|
||||
break
|
||||
candidate_height = fitted_height(sources=candidate, width=width, gap=gap)
|
||||
current_height = fitted_height(sources=sources[start:end], width=width, gap=gap)
|
||||
if candidate_height <= target_height and abs(current_height - target_height) <= abs(candidate_height - target_height):
|
||||
break
|
||||
end += 1
|
||||
group = sources[start:end]
|
||||
justify = end < len(sources) or fitted_height(sources=group, width=width, gap=gap) <= target_height
|
||||
rows.append(make_row(sources=group, width=width, gap=gap, target_height=target_height, justify=justify))
|
||||
start = end
|
||||
return rows
|
||||
|
||||
|
||||
def initialize_worker(archive_path: Path | None) -> None:
|
||||
global worker_archive
|
||||
if archive_path is not None:
|
||||
worker_archive = ZipFile(file=archive_path)
|
||||
atexit.register(worker_archive.close)
|
||||
|
||||
|
||||
def resize_with_pillow(cell: Cell) -> bytes:
|
||||
with read_image(source=cell.source, archive=worker_archive) as image:
|
||||
with image.resize(size=cell.size, resample=Image.Resampling.LANCZOS) as thumbnail:
|
||||
return thumbnail.tobytes()
|
||||
|
||||
|
||||
def resize_with_vips(cell: Cell, vips_binary: str) -> bytes:
|
||||
input_bytes = None
|
||||
if cell.source.member is not None:
|
||||
if worker_archive is None:
|
||||
raise ValueError("ZIP source requires an open archive")
|
||||
input_bytes = worker_archive.read(name=cell.source.member)
|
||||
operation, input_name = "thumbnail_source", "[descriptor=0]"
|
||||
elif cell.source.path is not None:
|
||||
operation, input_name = "thumbnail", str(cell.source.path)
|
||||
else:
|
||||
raise ValueError("Image source has no file or ZIP member")
|
||||
try:
|
||||
result = subprocess.run(
|
||||
args=[
|
||||
vips_binary,
|
||||
"--vips-concurrency=1",
|
||||
operation,
|
||||
input_name,
|
||||
".png",
|
||||
str(cell.size[0]),
|
||||
"--height",
|
||||
str(cell.size[1]),
|
||||
"--size",
|
||||
"both",
|
||||
],
|
||||
input=input_bytes,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
)
|
||||
except (OSError, subprocess.CalledProcessError) as error:
|
||||
detail = error.stderr.decode(errors="replace").strip() if isinstance(error, subprocess.CalledProcessError) and error.stderr else str(error)
|
||||
logging.warning(f"vips failed for {cell.source.name}; using Pillow: {detail}")
|
||||
return resize_with_pillow(cell=cell)
|
||||
with Image.open(fp=BytesIO(result.stdout)) as image:
|
||||
with image.convert(mode="RGBA") as thumbnail:
|
||||
if thumbnail.size == cell.size:
|
||||
return thumbnail.tobytes()
|
||||
with thumbnail.resize(size=cell.size, resample=Image.Resampling.LANCZOS) as exact_thumbnail:
|
||||
return exact_thumbnail.tobytes()
|
||||
|
||||
|
||||
def resize_thumbnail(cell: Cell, processor: Callable[..., bytes]) -> tuple[bytes, float]:
|
||||
started = perf_counter()
|
||||
logging.info(f"Resizing {cell.source.name} to {cell.size[0]}x{cell.size[1]}")
|
||||
pixels = processor(cell=cell)
|
||||
return pixels, perf_counter() - started
|
||||
|
||||
|
||||
def render_rows(
|
||||
sheet: Image.Image,
|
||||
rows: list[Row],
|
||||
input_path: Path,
|
||||
workers: int,
|
||||
gap: int,
|
||||
processor: Callable[..., bytes],
|
||||
) -> int:
|
||||
cells = []
|
||||
top = gap
|
||||
for row in rows:
|
||||
left = gap
|
||||
for source, cell_width in zip(row.sources, row.widths):
|
||||
cells.append(Cell(source=source, position=(left, top), size=(cell_width, row.height)))
|
||||
left += cell_width + gap
|
||||
top += row.height + gap
|
||||
remaining = iter(cells)
|
||||
pending: dict[Future[tuple[bytes, float]], Cell] = {}
|
||||
processed_count = rendered_count = 0
|
||||
archive_path = None if input_path.is_dir() else input_path
|
||||
with ProcessPoolExecutor(
|
||||
max_workers=workers,
|
||||
mp_context=get_context(method="spawn"),
|
||||
initializer=initialize_worker,
|
||||
initargs=(archive_path,),
|
||||
) as pool:
|
||||
while True:
|
||||
while len(pending) < workers * 2:
|
||||
cell = next(remaining, None)
|
||||
if cell is None:
|
||||
break
|
||||
future = pool.submit(resize_thumbnail, cell=cell, processor=processor)
|
||||
pending[future] = cell
|
||||
if not pending:
|
||||
break
|
||||
completed, _ = wait(fs=pending, return_when=FIRST_COMPLETED)
|
||||
for future in completed:
|
||||
cell = pending.pop(future)
|
||||
processed_count += 1
|
||||
try:
|
||||
pixels, duration = future.result()
|
||||
with Image.frombytes(mode="RGBA", size=cell.size, data=pixels) as thumbnail:
|
||||
sheet.paste(im=thumbnail, box=cell.position, mask=thumbnail)
|
||||
rendered_count += 1
|
||||
logging.info(f"Placed {cell.source.name} (resized in {duration:.1f}s)")
|
||||
except (OSError, ValueError, SyntaxError, RuntimeError, BadZipFile, Image.DecompressionBombError) as error:
|
||||
logging.warning(f"Skipping {cell.source.name} during rendering: {error}")
|
||||
logging.info(f"Progress {processed_count}/{len(cells)} ({processed_count / len(cells):.0%})")
|
||||
return rendered_count
|
||||
|
||||
|
||||
def create_sheet(
|
||||
input_path: Path,
|
||||
output_path: Path,
|
||||
width: int,
|
||||
target_height: int,
|
||||
gap: int,
|
||||
bg_color: str = "white",
|
||||
workers: int | None = None,
|
||||
processor: Callable[..., bytes] = resize_with_pillow,
|
||||
) -> None:
|
||||
if width <= 0 or target_height <= 0 or gap < 0 or width <= 2 * gap:
|
||||
raise ValueError("Invalid sheet dimensions or gap")
|
||||
if output_path.suffix.lower() not in {".jpg", ".jpeg"}:
|
||||
raise ValueError("Output must have a .jpg or .jpeg extension")
|
||||
if not input_path.is_dir() and not (input_path.is_file() and is_zipfile(filename=input_path)):
|
||||
raise ValueError(f"Expected a directory or ZIP archive: {input_path}")
|
||||
if input_path.resolve() == output_path.resolve():
|
||||
raise ValueError("Output must differ from input")
|
||||
if workers is not None and workers <= 0:
|
||||
raise ValueError("Workers must be positive")
|
||||
background = ImageColor.getcolor(color=bg_color, mode="RGB")
|
||||
with ExitStack() as stack:
|
||||
archive = None if input_path.is_dir() else stack.enter_context(ZipFile(file=input_path))
|
||||
sources = discover_sources(input_path=input_path, output_path=output_path, archive=archive)
|
||||
rows = layout_rows(sources=sources, width=width, gap=gap, target_height=target_height)
|
||||
height = sum(row.height for row in rows) + gap * (len(rows) + 1)
|
||||
worker_count = min(workers if workers is not None else os.cpu_count() or 1, len(sources))
|
||||
logging.info(f"Rendering {len(sources)} images in {len(rows)} rows on a {width}x{height} sheet using {worker_count} workers")
|
||||
sheet = stack.enter_context(Image.new(mode="RGB", size=(width, height), color=background))
|
||||
started = perf_counter()
|
||||
rendered_count = render_rows(
|
||||
sheet=sheet,
|
||||
rows=rows,
|
||||
input_path=input_path,
|
||||
workers=worker_count,
|
||||
gap=gap,
|
||||
processor=processor,
|
||||
)
|
||||
if not rendered_count:
|
||||
raise ValueError("No usable images could be decoded")
|
||||
logging.info(f"Rendering finished in {perf_counter() - started:.1f}s; saving {output_path}")
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
sheet.save(fp=output_path, quality=95)
|
||||
logging.info(f"Wrote {output_path} ({rendered_count} images, {width}x{height})")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
arguments = parse_args()
|
||||
input_paths = []
|
||||
for path in arguments.inputs:
|
||||
if path.exists():
|
||||
input_paths.append(path.resolve())
|
||||
else:
|
||||
logging.warning(f"Skipping non-existing input: {path}")
|
||||
if not input_paths:
|
||||
logging.error("No existing input paths")
|
||||
sys.exit(1)
|
||||
if arguments.output is not None and len(input_paths) > 1:
|
||||
logging.error("--output requires a single existing input")
|
||||
sys.exit(2)
|
||||
vips_binary = shutil.which(cmd="vips")
|
||||
processor = partial(resize_with_vips, vips_binary=vips_binary) if vips_binary else resize_with_pillow
|
||||
logging.info(f"Resize backend: {vips_binary or 'Pillow'}")
|
||||
failed_count = 0
|
||||
for index, input_path in enumerate(input_paths, start=1):
|
||||
logging.info(f"Processing input {index}/{len(input_paths)}: {input_path}")
|
||||
input_name = input_path.name if input_path.is_dir() else input_path.stem
|
||||
output_path = arguments.output or input_path.with_name(f"{input_name}_thumbs.jpg")
|
||||
try:
|
||||
create_sheet(
|
||||
input_path=input_path,
|
||||
output_path=output_path,
|
||||
width=arguments.width,
|
||||
target_height=arguments.target_height,
|
||||
gap=arguments.gap,
|
||||
bg_color=arguments.bg_color,
|
||||
workers=arguments.workers,
|
||||
processor=processor,
|
||||
)
|
||||
except (OSError, ValueError, RuntimeError, BadZipFile) as error:
|
||||
failed_count += 1
|
||||
logging.error(f"Could not create contact sheet for {input_path}: {error}")
|
||||
if failed_count:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user