generate.py

1import math2import mrlypy.gen3from automator.core.api import logger4from automator.core.errors import TaskAborted5from automator.core.models import Mockup, Task6from automator.core.s3 import cdn_key, put_png, s3_url, save_task7from automator.core.steps import Step8from io import BytesIO9from mrlypy.core.state import seed10from PIL import Image, ImageCms1112FORMAT = "PNG"13UNIT_SCALE = 114TILE_SCALE = 1015UNIT_IN = 0.2516DPI = 30017TILES = [1, 3, 5, 7, 9]18TILE_NAME = "tile"19OG_TILE = 320OG_NAME = "og"21OG_SIZE = 120022MAX_RENDERS = 32324PROFILE = ImageCms.createProfile("sRGB")25SRGB = ImageCms.ImageCmsProfile(PROFILE).tobytes()2627Image.MAX_IMAGE_PIXELS = None2829def crop(image: Image.Image, width: int, height: int) -> Image.Image:30    if image.width == width and image.height == height:31        return image32    w, h = image.size33    left = (w - width) // 234    top = (h - height) // 235    return image.crop(box=(left, top, left + width, top + height))3637def guard_renders(task: Task) -> Task:38    count = task.metadata.get("render_count", 0) + 139    task.metadata["render_count"] = count40    if count >= MAX_RENDERS:41        raise TaskAborted(f"render_count {count} for task {task.key}")42    save_task(task)43    return task4445def prepare_printfiles(task: Task, gen: mrlypy.gen.Gen) -> mrlypy.gen.Gen:46    tile_width, tile_height = gen.tile.unit_width, gen.tile.unit_height47    for pf in task.printfiles:48        grid_width = math.ceil(pf.width / (tile_width * UNIT_IN))49        grid_height = math.ceil(pf.height / (tile_height * UNIT_IN))50        gen.files.append(mrlypy.gen.File(width=grid_width, height=grid_height))51    return gen5253def prepare_tiles(task: Task, gen: mrlypy.gen.Gen) -> mrlypy.gen.Gen:54    for size in TILES:55        gen.files.append(mrlypy.gen.File(width=size, height=size))56    return gen5758def save_png(key: str, image: Image.Image, dpi: int = None) -> str:59    with BytesIO() as data:60        options = {"format": FORMAT, "optimize": True, "icc_profile": SRGB}61        if dpi:62            options["dpi"] = (dpi, dpi)63        image.save(data, **options)64        data.seek(0)65        put_png(key, data)66    return s3_url(key)6768def process_printfiles(task: Task, gen: mrlypy.gen.Gen, start: int) -> Task:69    for i, pf in enumerate(task.printfiles):70        pf.dpi = DPI71        image = gen.files[start + i].data72        full_width = round(image.width * UNIT_IN * pf.dpi)73        full_height = round(image.height * UNIT_IN * pf.dpi)74        image = image.resize(75            size=(full_width, full_height),76            resample=Image.Resampling.NEAREST,77        )78        image = crop(image, round(pf.width * pf.dpi), round(pf.height * pf.dpi))79        pf.url = save_png(cdn_key(task.key, pf.name), image, pf.dpi)80        logger.info(f"{task.desc} uploaded_printfile {pf.url}")81    return task8283def process_og(task: Task, image: Image.Image) -> None:84    og = image.resize(size=(OG_SIZE, OG_SIZE), resample=Image.Resampling.NEAREST)85    url = save_png(cdn_key(task.key, OG_NAME), og)86    logger.info(f"{task.desc} uploaded_og {url}")8788def process_tiles(task: Task, gen: mrlypy.gen.Gen, start: int) -> Task:89    for i, size in enumerate(TILES):90        raw = gen.files[start + i].data91        image = raw.resize(92            size=(raw.width * TILE_SCALE, raw.height * TILE_SCALE),93            resample=Image.Resampling.NEAREST,94        )95        name = f"{task.key}-{size}"96        url = save_png(cdn_key(task.key, f"{TILE_NAME}-{size}"), image)97        logger.info(f"{task.desc} uploaded_tile {url}")98        task.mockups.append(99            Mockup(100                id=size,101                key=task.key,102                name=name,103                category="Tile",104                title=f"{size}x{size}",105                url=url,106            )107        )108        if size == OG_TILE:109            process_og(task, raw)110    return task111112def mrly_generate(task: Task) -> Task:113    gen = mrlypy.gen.Gen.from_dict(task.variation)114    seed(gen.seed)115    gen = prepare_printfiles(task, gen)116    gen = prepare_tiles(task, gen)117    task = guard_renders(task)118    gen = mrlypy.gen.generate(gen)119    gen = mrlypy.gen.render(gen, scale=UNIT_SCALE)120    task = process_printfiles(task, gen, 0)121    task = process_tiles(task, gen, len(task.printfiles))122    task.variation = gen.to_dict()123    task.place(Step.MOCKUP)124    return task