【AIでRPG開発】AI生成したRPGツクール用の歩行CGを指定サイズに補正変換する（第1話）

"""ドット絵風画像を固定グリッドと共通パレットへ補正する。"""
from __future__ import annotations

import argparse
from collections import deque
from pathlib import Path

import numpy as np
from PIL import Image, ImageOps


def background_color(rgba: np.ndarray) -> np.ndarray:
    """画像の四隅の中央値から背景色を推定する。"""
    h, w = rgba.shape[:2]
    corners = rgba[[0, 0, h - 1, h - 1], [0, w - 1, 0, w - 1], :3]
    return np.median(corners, axis=0)


def remove_background(rgba: np.ndarray, tolerance: int) -> np.ndarray:
    """四隅の背景色に近く、外周につながる領域のみ透明にする。"""
    result = rgba.copy()
    h, w = result.shape[:2]
    background = background_color(result)
    eligible = (np.max(np.abs(result[..., :3].astype(float) - background), axis=2)
                <= tolerance) | (result[..., 3] == 0)
    seen = np.zeros((h, w), dtype=bool)
    queue = deque()
    for y, x in [(0, x) for x in range(w)] + [(h - 1, x) for x in range(w)] + [(y, 0) for y in range(h)] + [(y, w - 1) for y in range(h)]:
        if eligible[y, x] and not seen[y, x]:
            seen[y, x] = True
            queue.append((y, x))
    while queue:
        y, x = queue.popleft()
        for yy, xx in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
            if 0 <= yy < h and 0 <= xx < w and eligible[yy, xx] and not seen[yy, xx]:
                seen[yy, xx] = True
                queue.append((yy, xx))
    result[seen] = 0
    return result


def add_outline(rgba: np.ndarray) -> np.ndarray:
    """外周につながる透明領域に上下左右で接する不透明画素を黒くする。"""
    result = rgba.copy()
    transparent = rgba[..., 3] == 0
    h, w = transparent.shape
    exterior = np.zeros((h, w), dtype=bool)
    queue = deque()
    for y, x in [(0, x) for x in range(w)] + [(h - 1, x) for x in range(w)] + [(y, 0) for y in range(h)] + [(y, w - 1) for y in range(h)]:
        if transparent[y, x] and not exterior[y, x]:
            exterior[y, x] = True
            queue.append((y, x))
    while queue:
        y, x = queue.popleft()
        for yy, xx in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
            if 0 <= yy < h and 0 <= xx < w and transparent[yy, xx] and not exterior[yy, xx]:
                exterior[yy, xx] = True
                queue.append((yy, xx))
    adjacent = np.zeros((h, w), dtype=bool)
    adjacent[1:] |= exterior[:-1]
    adjacent[:-1] |= exterior[1:]
    adjacent[:, 1:] |= exterior[:, :-1]
    adjacent[:, :-1] |= exterior[:, 1:]
    result[adjacent & ~transparent, :3] = 0
    return result


def boundaries(profile: np.ndarray, length: int, count: int, snap: float) -> np.ndarray:
    """等間隔グリッドの近傍で強い色境界へ区画境界を移動する。"""
    step = length / count
    edges = [0]
    for i in range(1, count):
        ideal = i * step
        lo = max(edges[-1] + 1, int(np.ceil(ideal - step * snap)))
        hi = min(length - (count - i), int(np.floor(ideal + step * snap)))
        candidates = np.arange(lo, hi + 1)
        if candidates.size:
            scores = profile[candidates - 1] - 0.15 * np.abs(candidates - ideal) / step
            edges.append(int(candidates[np.argmax(scores)]))
        else:
            edges.append(int(round(ideal)))
    return np.array(edges + [length])


def convert(image: Image.Image, size=(144, 192), sheet=(3, 4), colors=32,
            background="auto", tolerance=24, snap=0.25, alpha_threshold=128,
            outline=False, frame_size=None) -> Image.Image:
    """フレームを個別に補正し、二値アルファのRGBA画像を返す。"""
    columns, rows = sheet
    if frame_size is not None:
        if min(frame_size) <= 0:
            raise ValueError("フレームの幅・高さは正の整数にしてください")
        size = (frame_size[0] * columns, frame_size[1] * rows)
    width, height = size
    if min(width, height, columns, rows) <= 0 or width % columns or height % rows:
        raise ValueError("出力サイズはシートの列数・行数で割り切れる正の整数にしてください")
    if not 2 <= colors <= 256 or not 0 <= snap <= 0.45:
        raise ValueError("色数は2～256、補正強度は0～0.45にしてください")
    if not 0 <= tolerance <= 255 or not 1 <= alpha_threshold <= 255:
        raise ValueError("背景許容差は0～255、アルファ閾値は1～255にしてください")
    if background not in ("auto", "keep"):
        raise ValueError("背景モードはautoまたはkeepにしてください")
    rgba = np.array(ImageOps.exif_transpose(image).convert("RGBA"))
    ih, iw = rgba.shape[:2]
    if iw < width or ih < height:
        raise ValueError("入力画像は出力サイズ以上の解像度が必要です")
    if background == "auto":
        rgba = remove_background(rgba, tolerance)
    opaque = rgba[..., 3] >= alpha_threshold
    output = np.zeros((height, width, 4), dtype=np.uint8)
    if not opaque.any():
        return Image.fromarray(output)
    # 透明領域を除いた色から全フレーム共通のパレットを作る。
    samples = rgba[..., :3][opaque]
    training = Image.fromarray(samples.reshape(1, -1, 3))
    palette_image = training.quantize(colors=colors, method=Image.Quantize.MEDIANCUT,
                                     dither=Image.Dither.NONE)
    mapped = Image.fromarray(rgba[..., :3]).quantize(palette=palette_image, dither=Image.Dither.NONE)
    labels = np.array(mapped).astype(np.int16)
    labels[~opaque] = 256
    palette = np.array(palette_image.getpalette(), dtype=np.uint8).reshape(-1, 3)
    fw, fh = width // columns, height // rows
    for row in range(rows):
        for col in range(columns):
            x0, x1 = round(col * iw / columns), round((col + 1) * iw / columns)
            y0, y1 = round(row * ih / rows), round((row + 1) * ih / rows)
            frame = labels[y0:y1, x0:x1]
            xp = np.mean(frame[:, 1:] != frame[:, :-1], axis=0)
            yp = np.mean(frame[1:, :] != frame[:-1, :], axis=1)
            xs = boundaries(xp, frame.shape[1], fw, snap)
            ys = boundaries(yp, frame.shape[0], fh, snap)
            for y in range(fh):
                for x in range(fw):
                    block = frame[ys[y]:ys[y + 1], xs[x]:xs[x + 1]]
                    # 中央を重く扱い、境界の混色を代表色へ持ち込まない。
                    by, bx = np.indices(block.shape)
                    weights = 1 + 2 * (1 - np.abs((bx + .5) / block.shape[1] * 2 - 1)) * (1 - np.abs((by + .5) / block.shape[0] * 2 - 1))
                    label = int(np.argmax(np.bincount(block.ravel(), weights=weights.ravel(), minlength=257)))
                    if label != 256:
                        output[row * fh + y, col * fw + x] = (*palette[label], 255)
            if outline:
                # 各フレーム内で判定し、隣のフレームの透明画素を参照しない。
                region = output[row * fh:(row + 1) * fh, col * fw:(col + 1) * fw]
                region[:] = add_outline(region)
    return Image.fromarray(output, "RGBA")


def indexed_png(image: Image.Image) -> Image.Image:
    """透明色を確保し、最大256色のパレット画像へ変換する。"""
    rgba = np.array(image.convert("RGBA"))
    opaque = rgba[..., 3] != 0
    offset = int(not opaque.all())
    budget = 256 - offset
    indices = np.zeros(opaque.shape, dtype=np.uint8)
    palette = [0] * 768
    if opaque.any():
        samples = rgba[..., :3][opaque]
        unique, inverse = np.unique(samples, axis=0, return_inverse=True)
        if len(unique) <= budget:
            indices[opaque] = inverse + offset
            palette[offset * 3:(offset + len(unique)) * 3] = unique.ravel().tolist()
        else:
            # 透明色を含めて256色を超える場合だけ追加で減色する。
            quantized = Image.fromarray(samples.reshape(1, -1, 3)).quantize(
                colors=budget, method=Image.Quantize.MEDIANCUT, dither=Image.Dither.NONE)
            indices[opaque] = np.array(quantized).ravel().astype(np.int16) + offset
            palette[offset * 3:] = quantized.getpalette()[:budget * 3]
    result = Image.fromarray(indices).convert("P")
    result.putpalette(palette)
    if offset:
        result.info["transparency"] = 0
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--width", type=int, help="シート全体の出力幅（既定144）")
    parser.add_argument("--height", type=int, help="シート全体の出力高さ（既定192）")
    parser.add_argument("--columns", type=int, default=3, help="シートの列数（既定3）")
    parser.add_argument("--rows", type=int, default=4, help="シートの行数（既定4）")
    parser.add_argument("--frame-width", type=int, help="出力フレーム1枚の幅")
    parser.add_argument("--frame-height", type=int, help="出力フレーム1枚の高さ")
    parser.add_argument("--alpha-map", action=argparse.BooleanOptionalAction, default=True,
                        help="アルファ付き32bit PNGで保存する（既定有効）。--no-alpha-mapで24bit RGB")
    parser.add_argument("--png8", action="store_true", help="アルファ設定を反映した8bitパレットPNGで保存する")
    parser.add_argument("--colors", type=int, default=32)
    parser.add_argument("--background", choices=("auto", "keep"), default="auto")
    parser.add_argument("--tolerance", type=int, default=24)
    parser.add_argument("--snap", type=float, default=.25)
    parser.add_argument("--alpha-threshold", type=int, default=128)
    parser.add_argument("--outline", action="store_true", help="外周の透明領域に接する内側1画素を黒い輪郭線にする")
    parser.add_argument("--preview", type=Path, help="最近傍で4倍拡大した確認画像の保存先")
    args = parser.parse_args()
    paths = [path.resolve() for path in (args.input, args.output, args.preview) if path is not None]
    if len(paths) != len(set(paths)):
        parser.error("入力・出力・プレビューには別々のパスを指定してください")
    frame_size = None
    if args.frame_width is not None or args.frame_height is not None:
        if args.frame_width is None or args.frame_height is None:
            parser.error("--frame-widthと--frame-heightは両方指定してください")
        if args.width is not None or args.height is not None:
            parser.error("フレームサイズと--width/--heightは同時に指定できません")
        frame_size = (args.frame_width, args.frame_height)
    try:
        with Image.open(args.input) as source:
            restore_background = not args.alpha_map and args.background == "keep" and args.outline
            fill_color = None
            if restore_background:
                # 輪郭判定用に背景を一時的に透明化し、保存前に元の背景色へ戻す。
                fill_color = background_color(np.array(ImageOps.exif_transpose(source).convert("RGBA")))
            result = convert(source, (args.width if args.width is not None else 144,
                                      args.height if args.height is not None else 192), (args.columns, args.rows),
                             args.colors, "auto" if restore_background else args.background,
                             args.tolerance, args.snap, args.alpha_threshold,
                             outline=args.outline, frame_size=frame_size)
            if fill_color is not None:
                pixels = np.array(result)
                pixels[pixels[..., 3] == 0, :3] = np.rint(fill_color).astype(np.uint8)
                result = Image.fromarray(pixels, "RGBA")
        args.output.parent.mkdir(parents=True, exist_ok=True)
        saved = result if args.alpha_map else result.convert("RGB")
        if args.png8:
            saved = indexed_png(saved)
        saved.save(args.output, format="PNG", **({"bits": 8} if args.png8 else {}))
        if args.preview:
            args.preview.parent.mkdir(parents=True, exist_ok=True)
            saved.resize((result.width * 4, result.height * 4), Image.Resampling.NEAREST).save(args.preview, format="PNG")
    except (OSError, ValueError) as exc:
        parser.exit(1, f"エラー: {exc}\n")
    print(f"保存: {args.output} ({result.width}x{result.height}, {args.columns}列×{args.rows}行, "
          f"フレーム{result.width // args.columns}x{result.height // args.rows}, {saved.mode})")


if __name__ == "__main__":
    main()