pigway-frame-flow/pigway_frame_flow/extensions/auto_collage.py

"""Random selection and layout shared by preview, display and scheduled jobs."""

import random
from io import BytesIO

from PIL import Image, ImageOps

from pigway_frame_flow.core.photo_preferences import eligible
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.extensions.layout import LayoutInput, LayoutStudio
from pigway_frame_flow.models import TemplateContent


def prepare(service, content, device_id=None):
    rng = random.SystemRandom()
    ids = []
    hashes = set()
    if content.source_id:
        from pigway_frame_flow.extensions.quality import passes, photo_quality

        source = service.photo_sources.get(content.source_id)
        items = eligible(
            service.db,
            content.source_id,
            service.photo_sources.catalog(content.source_id),
        )
        candidates = rng.sample(items, min(len(items), content.count * 5))
        for item in candidates:
            try:
                data, name, _ = service.photo_sources.read(
                    content.source_id, item["id"]
                )
                quality = photo_quality(data)
                if not passes(quality, source) or quality["sha256"] in hashes:
                    continue
                if source.avoid_duplicates and quality["fingerprint"] in hashes:
                    continue
                hashes.update([quality["sha256"], quality["fingerprint"]])
                with Image.open(
                    BytesIO(data), formats=("JPEG", "PNG", "WEBP")
                ) as image:
                    photo = ImageOps.exif_transpose(image).convert("RGB")
                    photo.thumbnail((2400, 2400))
                    output = BytesIO()
                    photo.save(output, "JPEG", quality=92)
                asset = service.assets.put(output.getvalue(), name, "cache")
                if asset["asset_id"] not in ids:
                    ids.append(asset["asset_id"])
            except (OSError, ValueError):
                continue
            if len(ids) == content.count:
                break
    else:
        rows = service.db.all(
            "SELECT id AS asset_id FROM assets WHERE kind='images' AND NOT EXISTS (SELECT 1 FROM photo_preferences p WHERE p.source_id='@assets' AND p.item_id=assets.id AND p.excluded=1) ORDER BY created_at DESC LIMIT 10000"
        )
        ids = rng.sample(
            [row["asset_id"] for row in rows], min(content.count, len(rows))
        )
    if len(ids) != content.count:
        raise FlowError(
            E.INVALID_REQUEST, "Not enough distinct photos for requested count", 422
        )
    values = content.model_dump(exclude={"type", "source_id", "count", "profile"})
    if content.orientation == "auto":
        from pigway_frame_flow.renderers.frame_geometry import frame_canvas

        display = (
            service.registry.get(device_id or service.settings.device_id)
            .device()
            .display
        )
        width, height = frame_canvas(display)
        values["orientation"] = "landscape" if width > height else "portrait"
    plan = LayoutStudio(service.renderer).plan(
        LayoutInput(
            asset_ids=ids,
            **values,
        )
    )
    return TemplateContent(
        type="template",
        template="poster" if content.style == "poster" else "collage",
        data={"layout": plan.model_dump()},
        profile=content.profile,
    )