grid_image_or_video.py

import math
from pathlib import Path

import numpy as np
from PIL import Image
import imageio.v3 as iio
from pillow_heif import register_heif_opener

register_heif_opener()

TARGET_SIZE = (256, 256)
START_OF_YEAR = 1672527600
OUTPUT_VIDEO = 'collage.mp4'
OUTPUT_COLLAGE = 'collage.jpg'
VIDEO_EXTENSIONS = {'.mp4', '.mov'}
IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.tif', '.tiff', '.heic'}
PHOTOS_DIR = Path('~/Pictures/Photos Library.photoslibrary/originals').expanduser().resolve()


def retrieve_images_and_videos_from_directory() -> list[Path]:
    """Retrieve the paths of all supported images & videos in a directory.

    This can be modified to enable media aggregation from multiple directories.
    """
    return [
        entry
        for entry in PHOTOS_DIR.glob('**/*')
        if entry.is_file()
        and int(entry.stat().st_birthtime) >= START_OF_YEAR
        and entry.suffix.lower() in IMAGE_EXTENSIONS | VIDEO_EXTENSIONS
    ]


def process_media_file(entry: Path) -> np.ndarray:
    if entry.suffix.lower() in VIDEO_EXTENSIONS:
        frame = iio.imread(entry, index=2, plugin='FFMPEG')
        return ensure_frame_shape(frame)

    if entry.suffix.lower() in IMAGE_EXTENSIONS:
        frame = iio.imread(entry, plugin='pillow')
        return ensure_frame_shape(frame)

    return None


def ensure_frame_shape(frame: np.ndarray) -> np.ndarray:
    """Ensures a consistent shape and channel count for frames."""

    # Convert to RGB if grayscale:
    if len(frame.shape) == 2:
        frame = np.stack((frame,) * 3, axis=-1)

    # Resize to target dimensions if necessary:
    if frame.shape[:2] != TARGET_SIZE:
        frame = Image.fromarray(frame).resize(TARGET_SIZE).convert('RGB')  # Ensure RGB format
        frame = np.array(frame)

    return frame


if __name__ == '__main__':
    media_files = retrieve_images_and_videos_from_directory()
    media_files_count = len(media_files)

    # image -- creates a grid image
    grid_size = math.ceil(math.sqrt(media_files_count))
    total_width = TARGET_SIZE[0] * grid_size
    total_height = TARGET_SIZE[1] * grid_size
    new_im = Image.new('RGB', (total_width, total_height))

    for idx, entry in enumerate(media_files):
        frame = process_media_file(entry)
        if frame is not None:
            img = Image.fromarray(frame)
            col = idx % grid_size
            row = idx // grid_size
            new_im.paste(img, (col * TARGET_SIZE[0], row * TARGET_SIZE[1]))

    new_im.save(OUTPUT_COLLAGE, 'JPEG')

    # video -- creates a sped up video of all images & video. Adjust `fps` to increase speed of video.
    writer = iio.imwrite(
        fps=20,
        uri=OUTPUT_VIDEO,
        codec='libx264',
        image=[frame for entry in media_files if (frame := process_media_file(entry)) is not None],
        plugin='pyav',
    )

pyproject.toml

[tool.poetry.dependencies]
python = "^3.11"
numpy = "^1.26.2"
imageio = { extras = ["ffmpeg"], version = "^2.33.1" }
av = "^11.0.0"
pillow = "^10.1.0"
pillow-heif = "^0.14.0"