skyeweave/atlasimagecomposer/backend/compose.py

import math
import pathlib
import sys
from typing import Tuple

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

from .atlas import fetch_data, fetch_image
from ..config import ExpressionDefaults, Paths

def compose():
    Paths.IMAGES.mkdir(exist_ok=True)
    Paths.OUTPUT.mkdir(exist_ok=True)

    servantid = input("Enter servant ID: ")
    try:
        t = int(servantid)
        if t <= 0:
            raise ValueError
    except ValueError:
        print("Servant ID has to be a valid integer above 0")
        sys.exit(1)
    sprites = fetch_data(servantid)

    for sprite in sprites:
        fetch_image(servantid, sprite)

    path = Paths.IMAGES / servantid
    for f in path.iterdir():
        process_sprite(f, sprites[f.stem]["position"], sprites[f.stem]["faceSize"], servantid)
    print(f"Files have been saved at: {(Paths.OUTPUT / servantid).absolute()}")


def override_expressions(expressions: Image.Image, facesize: int) -> Tuple[Image.Image, int]:
    exp_w, exp_h = expressions.size
    composites, offset, floating_offset = [], facesize // 5, 0
    per_composite = (ExpressionDefaults.faceSize * 4) // facesize

    # Why ..........
    composite_count = math.ceil((exp_h - 1280) / (offset + facesize * per_composite)) + 1

    # Remove top offset
    expressions = expressions.crop((0, 256, exp_w, exp_h))

    for count in range(0, composite_count):
        # width_topleft, height_topleft, width_botright, height_botright
        composites.append((0, 3 * count * facesize + floating_offset, exp_w, (3 * count + 3) * facesize + floating_offset - 1 ))
        floating_offset += offset

    # Create new expression image, then add all compositions to it
    overrides = Image.new("RGBA", (exp_w, facesize * composite_count * 3), (255, 0, 0, 0))
    for idx, c in enumerate(composites):
        crop = expressions.crop(c)
        pos = (0, idx * 3 * facesize)
        overrides.paste(crop, pos, crop)

    # Override expressions sheet
    return overrides.copy(), per_composite


# pylint: disable=too-many-locals
def process_sprite(filepath: pathlib.Path, position: tuple, facesize: int, servantid: str):
    im = Image.open(filepath)
    width, height = im.size

    main_sprite = im.crop((0, 0, width, 768))
    save_sprite(main_sprite, servantid, filepath.stem)
    expressions = im.crop((0, 768, width, height))

    expressionwidth, expressionheight = expressions.size
    rows = expressionheight // facesize
    expressions_p_row = expressionwidth // facesize

    exp_per_row = 4

    if facesize != ExpressionDefaults.faceSize:
        expressions, exp_per_row = override_expressions(expressions, facesize)

    for x, y in tqdm_itertools.product(range(0, rows), range(0, expressions_p_row, 1), ascii="-="):
        img = generate_sprite(main_sprite, expressions, x, y, position, facesize)
        if img is not None:
            save_sprite(img, servantid, filepath.stem, (x, y, exp_per_row))

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

# pylint: disable=too-many-arguments
def generate_sprite(sprite: Image.Image, expressions: Image.Image, row: int, col: int, position: tuple, facesize: int) -> Image.Image | None:
    area = (
        col * facesize,
        row * facesize,
        (col + 1) * facesize - 1,
        (row + 1) * facesize - 1
    )

    expression = expressions.crop(area)
    if is_empty(expression):
        return None

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

def save_sprite(image: Image.Image, folder: str, name: str, rowcol: tuple | None = None):
    savefolder = Paths.OUTPUT / folder
    if not savefolder.is_dir():
        savefolder.mkdir()
    postfix = f"_{rowcol[0] * rowcol[2] + rowcol[1] + 1}" if rowcol is not None else "_0"
    outfile = savefolder / f"{name}{postfix}.png"
    with open(outfile, 'wb') as file:
        image.save(file)