avi_lut_gen.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Generate AVI optimized palettes and RGB555->palette-index LUTs.

This script expects six AVI files named 1.avi..6.avi (case-insensitive)
and emits a C header suitable for static linking.
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path

AVI_COUNT = 6
PALETTE_SIZE = 256
LUT_SIZE = 1 << 15
RAW_FRAME_BYTES = PALETTE_SIZE * 4


def find_avi_file(input_dir: Path, index: int) -> Path:
    target_name = f"{index}.avi"
    matches = [p for p in sorted(input_dir.iterdir()) if p.is_file() and p.name.lower() == target_name]
    if not matches:
        raise FileNotFoundError(f"missing AVI file: {target_name} in {input_dir}")
    return matches[0]


def generate_palette(avi_path: Path, ffmpeg: str) -> list[tuple[int, int, int, int]]:
    cmd = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(avi_path),
        "-vf",
        "palettegen=stats_mode=full",
        "-frames:v",
        "1",
        "-f",
        "rawvideo",
        "-pix_fmt",
        "rgba",
        "pipe:1",
    ]
    raw = subprocess.check_output(cmd)
    if len(raw) != RAW_FRAME_BYTES:
        raise RuntimeError(f"unexpected palette size for {avi_path}: {len(raw)} bytes")

    palette = []
    for i in range(PALETTE_SIZE):
        base = i * 4
        palette.append((raw[base], raw[base + 1], raw[base + 2], raw[base + 3]))
    return palette


def expand_5_to_8(v: int) -> int:
    return (v << 3) | (v >> 2)


def palette_index_for_rgb(palette: list[tuple[int, int, int, int]], r: int, g: int, b: int) -> int:
    best_index = 0
    best_distance = 1 << 30
    for i, (pr, pg, pb, pa) in enumerate(palette):
        if pa == 0:
            continue
        dr = r - pr
        dg = g - pg
        db = b - pb
        distance = dr * dr + dg * dg + db * db
        if distance < best_distance:
            best_distance = distance
            best_index = i
    return best_index


def build_lut(palette: list[tuple[int, int, int, int]]) -> bytearray:
    lut = bytearray(LUT_SIZE)
    for rgb555 in range(LUT_SIZE):
        r5 = (rgb555 >> 10) & 0x1F
        g5 = (rgb555 >> 5) & 0x1F
        b5 = rgb555 & 0x1F
        r = expand_5_to_8(r5)
        g = expand_5_to_8(g5)
        b = expand_5_to_8(b5)
        lut[rgb555] = palette_index_for_rgb(palette, r, g, b)
    return lut


def emit_header(output: Path, palettes: list[list[tuple[int, int, int, int]]], luts: list[bytearray]) -> None:
    with output.open("w", encoding="utf-8", newline="\n") as f:
        f.write("/* -*- mode: c; tab-width: 4; c-basic-offset: 4; c-file-style: \"linux\" -*- */\n")
        f.write("/* This file is generated by scripts/avi_lut_gen.py. */\n\n")
        f.write("#ifndef AVI_LUT_GENERATED_H\n")
        f.write("#define AVI_LUT_GENERATED_H\n\n")
        f.write("#include \"common.h\"\n\n")
        f.write(f"#define AVI_LUT_FILE_COUNT {AVI_COUNT}\n\n")

        f.write("static const SDL_Color AVI_PALETTE_TABLE[AVI_LUT_FILE_COUNT][256] = {\n")
        for palette in palettes:
            f.write("    {\n")
            for r, g, b, a in palette:
                f.write(f"        {{{r}, {g}, {b}, {a}}},\n")
            f.write("    },\n")
        f.write("};\n\n")

        f.write("static const uint8_t AVI_RGB555_TO_INDEX_LUT[AVI_LUT_FILE_COUNT][32768] = {\n")
        for lut in luts:
            f.write("    {\n")
            for offset in range(0, LUT_SIZE, 16):
                chunk = lut[offset:offset + 16]
                values = ", ".join(f"0x{value:02X}" for value in chunk)
                f.write(f"        {values},\n")
            f.write("    },\n")
        f.write("};\n\n")
        f.write("#endif /* AVI_LUT_GENERATED_H */\n")


def main() -> int:
    parser = argparse.ArgumentParser(description="Generate SDLPAL AVI palette and LUT header.")
    parser.add_argument("input_dir", help="Directory containing 1.avi..6.avi")
    parser.add_argument("output", help="Output header path")
    parser.add_argument("--ffmpeg", default=os.environ.get("FFMPEG", "ffmpeg"), help="ffmpeg executable path")
    args = parser.parse_args()

    input_dir = Path(args.input_dir).resolve()
    output = Path(args.output).resolve()

    palettes: list[list[tuple[int, int, int, int]]] = []
    luts: list[bytearray] = []

    for index in range(1, AVI_COUNT + 1):
        avi_path = find_avi_file(input_dir, index)
        palette = generate_palette(avi_path, args.ffmpeg)
        lut = build_lut(palette)
        palettes.append(palette)
        luts.append(lut)
        print(f"generated {avi_path.name}", file=sys.stderr)

    emit_header(output, palettes, luts)
    print(f"wrote {output}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())