playground/imgsheet.py

#!/usr/bin/env python3.9
import argparse
import dataclasses
import io
import itertools
import logging
import math
import statistics
import sys
import typing
from pathlib import Path

from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont

T = typing.TypeVar("T")

logger = logging.getLogger("imgsheet")


def chunk(it: typing.Iterable[T], size: int) -> typing.Iterable[typing.Iterable[T]]:
    it = iter(it)
    sentinel = ()
    return iter(lambda: tuple(itertools.islice(it, size)), sentinel)


@dataclasses.dataclass
class PositionedImage:
    path: Path
    image_size: typing.Tuple[int, int]
    size: typing.Tuple[int, int] = dataclasses.field(default_factory=lambda: (0, 0))
    position: typing.Tuple[int, int] = dataclasses.field(default_factory=lambda: (0, 0))

    @property
    def top(self) -> int:
        return self.position[1]

    @property
    def left(self) -> int:
        return self.position[0]

    @property
    def bottom(self) -> int:
        return self.height + self.top

    @property
    def right(self) -> int:
        return self.width + self.left

    @property
    def width(self) -> int:
        return self.size[0]

    @property
    def height(self) -> int:
        return self.size[1]

    def __post_init__(self):
        self.size = self.image_size

    def scale(self, scale: float) -> None:
        self.size = (math.ceil(self.width * scale), math.ceil(self.height * scale))


class ImageSheet:
    def __init__(
        self,
        width: int,
        columns: int = 6,
        font_path: typing.Optional[Path] = None,
        label_format: typing.Optional[str] = None,
        label_font_size: typing.Optional[int] = 10,
    ):
        self.width = width
        self.padding = 20
        self.columns = columns
        self.font_path = font_path
        self.label_format = label_format
        self.label_font_size = label_font_size

    def create(self, image_paths: typing.Iterable[Path]) -> io.BytesIO:
        images = []
        for it in image_paths:
            with Image.open(it) as img:
                images.append(PositionedImage(it, img.size))

        if self.font_path and self.font_path.is_file():
            font = ImageFont.truetype(str(self.font_path), self.label_font_size)
        else:
            font = None

        positioned = self._calculate_positions(images)
        image_sheet = Image.new("RGB", (self.width, 0))
        for row_images in positioned:
            height = row_images[0].top + row_images[0].size[1]

            extended_sheet = Image.new("RGB", (self.width, height + self.padding))
            extended_sheet.paste(image_sheet, (0, 0))
            image_sheet = extended_sheet

            draw = ImageDraw.Draw(image_sheet)
            for it in row_images:
                img = Image.open(it.path)
                img.thumbnail(it.size, Image.LANCZOS)
                image_sheet.paste(img, it.position)

                if font:
                    # write filename on image
                    text_x, text_y = it.left + 2, it.bottom + 2
                    filesize = it.path.stat().st_size / 1_048_576

                    text = self._render_text(
                        width=it.image_size[0],
                        height=it.image_size[1],
                        size_mb=filesize,
                        filename=it.path.name,
                    )
                    if text:
                        draw.text((text_x, text_y), text, (255, 255, 255), font=font, align="left", anchor="la")

                print(f"{it.path.name}\t{it.position}\t{it.size}")

        f = io.BytesIO()
        image_sheet.save(f, format="JPEG")
        return f

    def _render_text(self, width: int, height: int, size_mb: float, filename: str) -> typing.Optional[str]:
        if not self.label_format:
            return None

        return self.label_format.format(width=width, height=height, size_mb=size_mb, filename=filename)

    def _calculate_positions(self, images: list[PositionedImage]) -> list[list[PositionedImage]]:
        cols = self.columns
        width = self.width
        padding = self.padding

        positioned = []

        rows = list(chunk(images, cols))
        x, y = 0, 0
        for row, row_images in enumerate(rows):
            row_images = list(row_images)

            usable_width = width - (len(row_images) - 1) * padding
            max_height = max(image.height for image in row_images)
            # check if there are enough images to fill the row
            if len(row_images) < cols:
                mean_height = statistics.median(prow[0].height for prow in positioned)
                max_height = mean_height

            for img in row_images:
                y_scaling = max_height / img.height
                img.scale(y_scaling)

            if len(row_images) >= cols:
                total_width = sum(image.width for image in row_images)
                x_scaling = usable_width / total_width
                max_height = math.ceil(max_height * x_scaling)
                for img in row_images:
                    img.scale(x_scaling)

            is_last_row = row == len(rows) - 1
            x = 0
            for i, img in enumerate(row_images):
                img.position = (x, y)
                is_last_col = i == len(row_images) - 1
                x = x + img.width + (0 if is_last_col else padding)
            y = y + max_height + (0 if is_last_row else padding)

            positioned.append(row_images)

        return positioned


def parse_args(argv: typing.List[str]) -> argparse.Namespace:
    arger = argparse.ArgumentParser(
        description="Create a sheet of images", formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    arger.add_argument("--width", "-w", type=int, default=2000, help="Width of the sheet")
    arger.add_argument("--columns", "-c", type=int, default=4, help="Number of columns")
    arger.add_argument("--output", "-o", type=Path, default=Path("image_sheet.jpg"), help="Output file")
    arger.add_argument("--font", dest="font_path", type=Path, help="Label font to use")
    arger.add_argument("--font-size", dest="font_size", default=10, type=int, help="Label font size")
    arger.add_argument(
        "--label",
        dest="label_format",
        default="{width}x{height} {size_mb:.1f}MB {filename}",
        help="Label format. Available tokens: width, height, size_mb, filename",
    )
    arger.add_argument("images", nargs="+", type=Path, help="Images or image directory to include in the sheet")
    if len(argv) == 0:
        arger.print_help()
        sys.exit(1)
    return arger.parse_args(argv)


def main():
    args = parse_args(sys.argv[1:])

    images: list[Path] = args.images
    if len(images) == 1 and images[0].expanduser().is_dir():
        images = sorted(it for it in images[0].resolve().glob("*.jpg"))

    sheet = ImageSheet(
        width=args.width,
        columns=args.columns,
        font_path=args.font_path,
        label_font_size=args.font_size,
        label_format=args.label_format,
    )

    sheet_img = sheet.create(images)
    with args.output.open("wb") as f:
        f.write(sheet_img.getvalue())


if __name__ == "__main__":
    main()