skyeweave/atlasimagecomposer/backend/compose.py

import logging
import pathlib
from typing import List, Optional
from collections import Counter

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

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

LOGGER = logging.getLogger(Logging.NAME)

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

    if input_id < 10000:
        chara_ids = fetch_data(input_id)
        savefolder = Paths.OUTPUT / str(input_id)
    else:
        LOGGER.info(f"Processing manually uploaded charaId {input_id}")
        savefolder = Paths.OUTPUT / "manual"
        chara_ids = [str(input_id)]

    if not savefolder.is_dir():
        savefolder.mkdir(parents=True, exist_ok=True)

    if filters is not None:
        chara_ids = [ v for v in chara_ids if v in filters ]
    LOGGER.debug(chara_ids)

    for char_id in chara_ids:
        expfolder = fetch_expression_sheets(savefolder.stem, char_id)
        config = fetch_config(char_id)
        process_sprite(expfolder, config, savefolder)

    LOGGER.info(f"Files have been saved at: {savefolder.absolute()}")


def calculate_counts(width: int, height: int, facesize: tuple[int, int]):
    rowcount, colcount = height // facesize[1], width // facesize[0]
    LOGGER.debug(f"{height} | {facesize[1]} --> {rowcount}")
    LOGGER.debug(f"{width} | {facesize[0]} --> {colcount}")
    return rowcount, colcount

def gen_main_sprite(folder: pathlib.Path):
    image = Image.open(folder / "0.png")
    width, height = image.size
    LOGGER.debug(f"Main sprite ({folder}): {width}:{height}")
    return image.crop((0, 0, width, height - 256))

def process_sprite(images_folder: pathlib.Path, configdata: SpritesheetData, outputfolder: pathlib.Path):
    main_sprite = gen_main_sprite(images_folder)
    image_idx = save_sprite(main_sprite, outputfolder, f"{images_folder.stem}")

    for i in images_folder.iterdir():
        LOGGER.debug(f"Idx: {image_idx}")
        initial_row = 0
        expressions = Image.open(i)

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

        if i.name == "0.png" and 256 < configdata["facesize"][1]:
            continue

        if i.name == "0.png":
            initial_row = rowcount - 1

        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:
                image_idx = save_sprite(img, outputfolder, f"{images_folder.stem}", image_idx)
            LOGGER.debug(f"{x}/{y} - {'Invalid' if img is None else 'Valid'} image")


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[0],
        row * facesize[1],
        (col + 1) * facesize[0],
        (row + 1) * facesize[1]
    )
    LOGGER.debug(roi)
    expression = expressions.crop(roi)

    if is_empty(expression):
        LOGGER.debug("Image empty")
        return None

    mask = Image.new("RGBA", (facesize[0], facesize[1]), (255,255,255,255))
    composition = main_sprite.copy()
    composition.paste(expression, position, mask)
    return composition


def save_sprite(image: Image.Image, outputfolder: pathlib.Path, name: str, idx: int = 0) -> int:
    savefolder = outputfolder / name
    if not savefolder.is_dir():
        savefolder.mkdir()
    outfile = savefolder / f"{idx}.png"

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

    return idx + 1

def is_empty(img: Image.Image):
    data = Counter(img.crop((96, 96, 160, 160)).convert('LA').getdata())
    LOGGER.debug(f"Counts: {len(data)}")
    if len(data) < 10:
        return True
    return False