scripts/extract_star_action_assets.py

from __future__ import annotations

from pathlib import Path

from PIL import Image


ROOT = Path(__file__).resolve().parents[1]
SHEET_PATH = ROOT / "assets" / "generated" / "star_action_sheet.png"
OUT_DIR = ROOT / "assets" / "characters" / "star"
CANVAS_SIZE = (900, 1200)
POSES = [
    "idle",
    "listening",
    "thinking",
    "worried",
    "smile",
    "happy",
    "focus",
    "talk",
]


def _remove_green_screen(image: Image.Image) -> Image.Image:
    image = image.convert("RGBA")
    pixels = image.load()
    width, height = image.size
    for y in range(height):
        for x in range(width):
            red, green, blue, alpha = pixels[x, y]
            is_green_screen = green > 110 and green > red * 1.28 and green > blue * 1.12
            is_green_edge = green > 80 and green > red * 1.45 and green > blue * 1.25
            if is_green_screen or is_green_edge:
                pixels[x, y] = (red, green, blue, 0)
            elif alpha:
                spill = max(0, green - max(red, blue) - 24)
                if spill:
                    pixels[x, y] = (red, max(0, green - spill), blue, alpha)
    return _remove_right_edge_artifacts(image)


def _remove_right_edge_artifacts(image: Image.Image) -> Image.Image:
    alpha = image.getchannel("A")
    width, height = image.size
    seen: set[tuple[int, int]] = set()
    pixels_to_clear: list[tuple[int, int]] = []

    for y in range(height):
        for x in range(width):
            if (x, y) in seen or alpha.getpixel((x, y)) < 12:
                continue
            stack = [(x, y)]
            seen.add((x, y))
            component: list[tuple[int, int]] = []
            min_x = max_x = x

            while stack:
                px, py = stack.pop()
                component.append((px, py))
                min_x = min(min_x, px)
                max_x = max(max_x, px)
                for nx, ny in ((px - 1, py), (px + 1, py), (px, py - 1), (px, py + 1)):
                    if nx < 0 or ny < 0 or nx >= width or ny >= height or (nx, ny) in seen:
                        continue
                    if alpha.getpixel((nx, ny)) >= 12:
                        seen.add((nx, ny))
                        stack.append((nx, ny))

            touches_right_edge = max_x >= width - 2
            starts_on_right_side = min_x > width * 0.55
            if len(component) < 16 or (touches_right_edge and starts_on_right_side):
                pixels_to_clear.extend(component)

    if not pixels_to_clear:
        return image

    pixels = image.load()
    for x, y in pixels_to_clear:
        red, green, blue, _alpha = pixels[x, y]
        pixels[x, y] = (red, green, blue, 0)
    return image


def _fit_to_stage_canvas(image: Image.Image) -> Image.Image:
    alpha = image.getchannel("A")
    bbox = alpha.getbbox()
    if bbox is None:
        return Image.new("RGBA", CANVAS_SIZE, (0, 0, 0, 0))

    crop = image.crop(bbox)
    crop_width, crop_height = crop.size
    max_width = 760
    target_height = 1080
    scale = min(target_height / crop_height, max_width / crop_width)
    new_size = (round(crop_width * scale), round(crop_height * scale))
    crop = crop.resize(new_size, Image.Resampling.LANCZOS)

    canvas = Image.new("RGBA", CANVAS_SIZE, (0, 0, 0, 0))
    x = (CANVAS_SIZE[0] - new_size[0]) // 2
    y = CANVAS_SIZE[1] - new_size[1] - 24
    canvas.alpha_composite(crop, (x, y))
    return canvas


def main() -> None:
    sheet = Image.open(SHEET_PATH)
    cell_width = sheet.width // 4
    cell_height = sheet.height // 2
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    for index, pose in enumerate(POSES):
        row, col = divmod(index, 4)
        cell = sheet.crop(
            (
                col * cell_width,
                row * cell_height,
                (col + 1) * cell_width,
                (row + 1) * cell_height,
            )
        )
        stage_asset = _fit_to_stage_canvas(_remove_green_screen(cell))
        if pose == "focus":
            _clear_box(stage_asset, x_min=460, y_min=360)
        out_path = OUT_DIR / f"{pose}.png"
        stage_asset.save(out_path)
        print(out_path.relative_to(ROOT))


def _clear_box(image: Image.Image, x_min: int, y_min: int) -> None:
    pixels = image.load()
    for y in range(y_min, image.height):
        for x in range(x_min, image.width):
            red, green, blue, _alpha = pixels[x, y]
            pixels[x, y] = (red, green, blue, 0)


if __name__ == "__main__":
    main()