skyeweave/atlasimagecomposer/backend/compose.py

import pathlib
from typing import List, Optional

import numpy as np
from PIL import Image
from tqdm.contrib import itertools as tqdm_itertools

from ..config import Paths
from .atlas import SpritesheetData, fetch_data, fetch_expression_sheets


def compose(servantid: int, filters: Optional[List[str]] = None):
    Paths.IMAGES.mkdir(exist_ok=True)
    Paths.OUTPUT.mkdir(exist_ok=True)

    sprites = fetch_data(servantid)

    if filters is not None:
        sprites = { key: value for key, value in sprites.items() if key in filters }

    for sprite in sprites:
        fetch_expression_sheets(servantid, sprite)

    path = Paths.IMAGES / str(servantid)

    savefolder = Paths.OUTPUT / str(servantid)
    if not savefolder.is_dir():
        savefolder.mkdir()

    for f in path.iterdir():
        if filters is not None and str(f.stem) not in filters:
            continue
        process_sprite(f, servantid, sprites[f.stem])

    print(f"Files have been saved at: {(Paths.OUTPUT / str(servantid)).absolute()}")


def calculate_counts(width: int, height: int, facesize: int):
    return height // facesize, width // facesize


def process_sprite(images_folder: pathlib.Path, servantid: int, configdata: SpritesheetData):
    main = Image.open(images_folder / "0.png")

    width, _ = main.size
    main_sprite = main.crop((0, 0, width, 756))

    save_sprite(main_sprite, servantid, f"{images_folder.stem}")

    for i in images_folder.iterdir():
        initial_row, index = 0, int(i.stem)
        expressions = Image.open(i)

        rowcount, colcount = calculate_counts(*expressions.size, configdata["facesize"])

        if i.name == "0.png":
            initial_row = 3

        for x, y in tqdm_itertools.product(range(initial_row, rowcount), range(0, colcount), ascii="-="):
            img = generate_sprite(main_sprite, expressions, x, y, configdata)
            if img is not None:
                save_sprite(img, servantid, f"{images_folder.stem}", (x, y, colcount, index))


def generate_sprite(main_sprite: Image.Image, expressions: Image.Image, row: int, col: int, configdata: SpritesheetData) -> Image.Image | None:
    position, facesize = configdata["position"], configdata["facesize"]
    roi = (
        col * facesize,
        row * facesize,
        (col + 1) * facesize - 1,
        (row + 1) * facesize - 1
    )

    expression = expressions.crop(roi)

    if is_empty(expression):
        return None

    composition = main_sprite.copy()
    composition.paste(expression, position, expression)
    return composition


def save_sprite(image: Image.Image, folder_id: int, name: str, info: tuple | None = None):
    savefolder = Paths.OUTPUT / str(folder_id) / name
    if not savefolder.is_dir():
        savefolder.mkdir()

    postfix = "0"

    if info is not None:
        (row, col, column_count, file_idx) = info

        if file_idx == 0 and column_count == 4:
            postfix = str(col + 1)
        elif file_idx == 0:
            raise ValueError("Should not have any faces")
        elif column_count == 4:
            postfix = str((column_count * row + col + 1) + pow(column_count, 2) * (file_idx - 1) + column_count)
        elif column_count < 4:
            postfix = str((column_count * row + col + 1) + pow(column_count, 2) * (file_idx - 1))
        else:
            raise ValueError("Unaccounted case")

    outfile = savefolder / f"{postfix}.png"

    with open(outfile, 'wb') as file:
        image.save(file)


def is_empty(img: Image.Image):
    data = np.asarray(img.crop((96, 96, 160, 160)).convert('LA'))
    np.reshape(data, (-1, 1))
    _, count_unique = np.unique(data, return_counts=True)
    if count_unique.size < 10:
        return True
    return False