postprocess.py

import tempfile
from collections import deque
from pathlib import Path

import numpy as np
from perfect_pixel import get_perfect_pixel
from PIL import Image, ImageFilter


COLUMNS = 4
ROWS = 2
SHEET_TASKS = {"Propagate frame 1 appearance", "Dress 4x2 walk sheet"}


def remove_white_background(image, minimum_tolerance=48):
    rgb = np.asarray(image.convert("RGB"), dtype=np.float32)
    border = np.concatenate((rgb[0], rgb[-1], rgb[:, 0], rgb[:, -1]))
    background = np.median(border, axis=0)
    tolerance = max(
        minimum_tolerance,
        float(np.percentile(np.linalg.norm(border - background, axis=1), 50) + 8),
    )
    rough = np.linalg.norm(rgb - background, axis=2) > tolerance
    rough = np.asarray(
        Image.fromarray(rough.astype(np.uint8) * 255)
        .filter(ImageFilter.MaxFilter(3))
        .filter(ImageFilter.MinFilter(3))
    ) > 0
    height, width = rough.shape
    seen = np.zeros_like(rough)
    components = []
    for seed_y, seed_x in zip(*np.nonzero(rough)):
        if seen[seed_y, seed_x]:
            continue
        seen[seed_y, seed_x] = True
        queue = [(int(seed_y), int(seed_x))]
        component = []
        while queue:
            y, x = queue.pop()
            component.append((y, x))
            for next_y, next_x in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
                if (
                    0 <= next_y < height
                    and 0 <= next_x < width
                    and rough[next_y, next_x]
                    and not seen[next_y, next_x]
                ):
                    seen[next_y, next_x] = True
                    queue.append((next_y, next_x))
        components.append(component)
    if not components:
        return Image.new("RGBA", image.size)

    minimum_area = max(4, round(max(map(len, components)) * 0.002))
    solid = np.zeros_like(rough)
    for component in components:
        if len(component) >= minimum_area:
            y, x = zip(*component)
            solid[y, x] = True

    outside = np.zeros_like(solid)
    queue = deque()
    for x in range(width):
        queue.extend(((0, x), (height - 1, x)))
    for y in range(height):
        queue.extend(((y, 0), (y, width - 1)))
    while queue:
        y, x = queue.popleft()
        if outside[y, x] or solid[y, x]:
            continue
        outside[y, x] = True
        for next_y, next_x in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
            if 0 <= next_y < height and 0 <= next_x < width:
                queue.append((next_y, next_x))
    solid |= ~outside
    alpha = solid.astype(np.uint8) * 255
    return Image.fromarray(np.dstack((rgb.astype(np.uint8), alpha)), "RGBA")


def foot_anchor(image):
    alpha = np.asarray(image.getchannel("A"), dtype=np.float64) / 255
    y, x = np.nonzero(alpha > 0.25)
    if not len(x):
        raise ValueError("A 4x2 frame contains no foreground sprite.")
    bottom = int(y.max())
    band = y >= bottom - max(2, round(image.height * 0.06))
    return float(np.average(x[band], weights=alpha[y[band], x[band]])), bottom


def align_4x2(source, reference):
    if source.width % COLUMNS or source.height % ROWS:
        raise ValueError("The generated sheet must be divisible into a 4x2 grid.")
    frame_width = source.width // COLUMNS
    frame_height = source.height // ROWS
    reference = reference.resize(source.size, Image.Resampling.NEAREST)
    result = Image.new("RGBA", source.size)

    for index in range(COLUMNS * ROWS):
        row, column = divmod(index, COLUMNS)
        box = (
            column * frame_width,
            row * frame_height,
            (column + 1) * frame_width,
            (row + 1) * frame_height,
        )
        generated = remove_white_background(source.crop(box))
        sprite_box = generated.getbbox()
        if not sprite_box:
            raise ValueError(f"Generated frame {index + 1} is empty.")
        sprite = generated.crop(sprite_box)
        generated_x, generated_y = foot_anchor(sprite)
        reference_frame = remove_white_background(reference.crop(box), 4)
        reference_x, reference_y = foot_anchor(reference_frame)

        frame = Image.new("RGBA", (frame_width, frame_height))
        frame.alpha_composite(
            sprite,
            (
                round(reference_x - generated_x),
                round(reference_y - generated_y),
            ),
        )
        result.alpha_composite(frame, (box[0], box[1]))

    white = Image.new("RGBA", result.size, "white")
    white.alpha_composite(result)
    return white.convert("RGB")


def adaptive_palette(image, colors=32):
    quantized = image.convert("RGB").quantize(
        colors=colors, method=Image.Quantize.MEDIANCUT, dither=Image.Dither.NONE
    )
    palette = quantized.getpalette()
    used = sorted(quantized.getcolors(), reverse=True)
    result = [
        tuple(palette[index * 3 : index * 3 + 3])
        for _, index in used[:colors]
    ]
    whitest = max(range(len(result)), key=lambda i: sum(result[i]))
    result[whitest] = (255, 255, 255)
    return list(dict.fromkeys(result))


def reference_palette(image, colors=32):
    rgb = np.asarray(image.convert("RGB"), dtype=np.uint8).reshape(-1, 3)
    unique, counts = np.unique(rgb, axis=0, return_counts=True)
    if len(unique) <= colors:
        order = np.argsort(counts)[::-1]
        palette = [tuple(map(int, color)) for color in unique[order]]
    else:
        palette = adaptive_palette(image, colors)
    if (255, 255, 255) not in palette:
        palette = [(255, 255, 255), *palette[: colors - 1]]
    return palette[:colors]


def indexed_image(image, palette):
    palette = palette[:32]
    pixels = np.asarray(image.convert("RGB"), dtype=np.int16)
    colors = np.asarray(palette, dtype=np.int16)
    flat = pixels.reshape(-1, 3)
    indexes = np.empty(len(flat), dtype=np.uint8)
    for start in range(0, len(flat), 65536):
        chunk = flat[start : start + 65536].astype(np.int32)
        delta = chunk[:, None] - colors[None].astype(np.int32)
        distance = (delta**2).sum(axis=2)
        indexes[start : start + len(chunk)] = distance.argmin(axis=1)
    result = Image.fromarray(indexes.reshape(pixels.shape[:2]), "P")
    padded = [channel for color in palette for channel in color]
    padded.extend([channel for _ in range(32 - len(palette)) for channel in palette[-1]])
    result.putpalette(padded + [0] * (768 - len(padded)))
    return result


def native_size(task):
    return (256, 128) if task in SHEET_TASKS else (64, 64)


def perfect_pixel_image(image, task):
    expected = native_size(task)
    if image.size == expected:
        return image.convert("RGB")
    width, height, refined = get_perfect_pixel(
        np.asarray(image.convert("RGB")),
        sample_method="median",
        min_size=4.0,
        peak_width=6,
        refine_intensity=0.25,
        fix_square=True,
    )
    if width is None or height is None:
        raise ValueError("Perfect Pixel이 이미지의 픽셀 격자를 찾지 못했습니다.")
    if (width, height) != expected:
        raise ValueError(
            f"Perfect Pixel 검출 크기는 {width}×{height}이지만 "
            f"이 작업에는 {expected[0]}×{expected[1]} 격자가 필요합니다."
        )
    return Image.fromarray(np.asarray(refined, dtype=np.uint8), "RGB")


def shared_palette(paths, colors=32):
    images = [Image.open(path).convert("RGB") for path in paths]
    if not images:
        raise ValueError("공통 팔레트를 만들 이미지가 없습니다.")
    width = max(image.width for image in images)
    height = sum(image.height for image in images)
    combined = Image.new("RGB", (width, height), "white")
    y = 0
    for image in images:
        combined.paste(image, (0, y))
        y += image.height
    return adaptive_palette(combined, colors)


def apply_shared_palette(paths, palette):
    results = []
    for path in paths:
        image = indexed_image(Image.open(path).convert("RGB"), palette)
        output = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
        image.save(output, bits=5)
        results.append(output)
    return results


def save_gif(sheet, palette, fps):
    frame_width = sheet.width // COLUMNS
    frame_height = sheet.height // ROWS
    frames = []
    for row in range(ROWS):
        for column in range(COLUMNS):
            frame = sheet.crop(
                (
                    column * frame_width,
                    row * frame_height,
                    (column + 1) * frame_width,
                    (row + 1) * frame_height,
                )
            )
            frames.append(indexed_image(frame.convert("RGB"), palette))
    path = tempfile.NamedTemporaryFile(delete=False, suffix=".gif").name
    frames[0].save(
        path,
        save_all=True,
        append_images=frames[1:],
        duration=round(1000 / fps),
        loop=0,
        disposal=2,
    )
    return path


def process_output(
    generated_path,
    input_path,
    task,
    palette_reference_path,
    palette_mode,
    align_frames,
    pixel_snap,
    output_resolution,
    output_format,
    fps,
):
    generated = Image.open(generated_path).convert("RGB")
    reference = Image.open(input_path).convert("RGB")
    sheet_task = task in SHEET_TASKS
    if align_frames and sheet_task:
        generated = align_4x2(generated, reference)

    target_native = native_size(task)
    if pixel_snap:
        working = perfect_pixel_image(generated, task)
    elif output_resolution == "Native LPC":
        working = generated.resize(target_native, Image.Resampling.NEAREST)
    else:
        working = generated

    if palette_mode == "Defer shared palette":
        palette = None
    elif palette_mode == "Lock reference palette":
        palette_source = Image.open(palette_reference_path).convert("RGB") if palette_reference_path else reference
        palette_source = palette_source.resize(target_native, Image.Resampling.NEAREST)
        palette = reference_palette(palette_source)
    else:
        palette = adaptive_palette(working)
    if palette:
        working = indexed_image(working, palette)

    if output_resolution == "Upscaled" and working.size != generated.size:
        working = working.resize(generated.size, Image.Resampling.NEAREST)
    if output_format == "GIF":
        if not sheet_task:
            raise ValueError("GIF output is available for 4x2 sheet tasks.")
        if not palette:
            raise ValueError("공통 팔레트 적용 전에는 GIF를 만들 수 없습니다.")
        return save_gif(working, palette, fps)

    path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
    working.save(path, bits=5)
    return path


def self_check():
    sheet = Image.new("RGB", (256, 128), "white")
    array = np.asarray(sheet).copy()
    for index in range(8):
        row, column = divmod(index, 4)
        array[row * 64 + 20 : row * 64 + 60, column * 64 + 24 : column * 64 + 40] = (
            index * 20,
            80,
            160,
        )
    palette = adaptive_palette(Image.fromarray(array))
    indexed = indexed_image(Image.fromarray(array), palette)
    assert len(indexed.getcolors()) <= 32
    rng = np.random.default_rng(7)
    test_colors = np.asarray(
        [(255, 255, 255), (20, 30, 40), (50, 90, 160), (200, 120, 80)],
        dtype=np.uint8,
    )
    test_grid = test_colors[rng.integers(0, len(test_colors), size=(128, 256))]
    upscaled = Image.fromarray(test_grid).resize(
        (2048, 1024), Image.Resampling.NEAREST
    )
    assert perfect_pixel_image(upscaled, "Propagate frame 1 appearance").size == (
        256,
        128,
    )
    gif = save_gif(indexed, palette, 8)
    with Image.open(gif) as animation:
        assert animation.n_frames == 8
    Path(gif).unlink()
    print("postprocess self-check passed")


if __name__ == "__main__":
    self_check()