Files
lucky-cat/tools/optimize_lcd_assets.py

288 lines
11 KiB
Python

#!/usr/bin/env python3
"""Generate and verify display-lossless PNG assets for the RGB565 LCD."""
from __future__ import annotations
import argparse
import hashlib
import json
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Any, Iterable
PROJECT_ROOT = Path(__file__).resolve().parent.parent
SOURCE_DIR = PROJECT_ROOT / "components" / "platform" / "assets" / "lcd_source"
OPTIMIZED_DIR = PROJECT_ROOT / "components" / "platform" / "assets" / "lcd_optimized"
MANIFEST_NAME = "manifest.json"
MANIFEST_VERSION = 1
class AssetError(RuntimeError):
pass
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for chunk in iter(lambda: source.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def png_files(directory: Path) -> list[Path]:
if not directory.is_dir():
raise AssetError(f"Asset directory does not exist: {directory}")
return sorted(path for path in directory.iterdir() if path.is_file() and path.suffix.lower() == ".png")
def load_pillow() -> tuple[Any, str]:
try:
import PIL
from PIL import Image
except ImportError as exc:
raise AssetError(
"Pillow is required for this command. Install it with:\n"
f" {sys.executable} -m pip install -r {PROJECT_ROOT / 'tools' / 'requirements-assets.txt'}"
) from exc
return Image, PIL.__version__
def rgb565_value(pixel: tuple[int, int, int]) -> int:
red, green, blue = pixel
return ((red & 0xF8) << 8) | ((green & 0xFC) << 3) | (blue >> 3)
def expand_rgb565(pixel: tuple[int, int, int]) -> tuple[int, int, int]:
red, green, blue = pixel
red5 = red >> 3
green6 = green >> 2
blue5 = blue >> 3
return (
(red5 << 3) | (red5 >> 2),
(green6 << 2) | (green6 >> 4),
(blue5 << 3) | (blue5 >> 2),
)
def save_indexed_png(image_module: Any, size: tuple[int, int], pixels: list[tuple[int, int, int]], output: Path) -> int:
palette_colors = sorted(set(pixels))
if len(palette_colors) > 256:
raise AssetError("Indexed PNG requested with more than 256 colors")
palette_index = {color: index for index, color in enumerate(palette_colors)}
indices = bytes(palette_index[pixel] for pixel in pixels)
indexed = image_module.frombytes("P", size, indices)
flat_palette = [channel for color in palette_colors for channel in color]
flat_palette.extend([0] * (768 - len(flat_palette)))
indexed.putpalette(flat_palette)
indexed.save(output, format="PNG", optimize=True, compress_level=9)
return len(palette_colors)
def assert_display_equivalent(source_rgb: Any, optimized_path: Path, image_module: Any) -> None:
with image_module.open(optimized_path) as optimized_image:
optimized_rgb = optimized_image.convert("RGB")
if optimized_rgb.size != source_rgb.size:
raise AssetError(
f"Image dimensions changed for {optimized_path.name}: "
f"{source_rgb.size} -> {optimized_rgb.size}"
)
for index, (source_pixel, optimized_pixel) in enumerate(
zip(source_rgb.get_flattened_data(), optimized_rgb.get_flattened_data())
):
if rgb565_value(source_pixel) != rgb565_value(optimized_pixel):
x = index % source_rgb.width
y = index // source_rgb.width
raise AssetError(
f"RGB565 mismatch in {optimized_path.name} at ({x}, {y}): "
f"{source_pixel} -> {optimized_pixel}"
)
def optimize_one(source_path: Path, output_path: Path, image_module: Any) -> dict[str, Any]:
with image_module.open(source_path) as source_image:
source_mode = source_image.mode
rgba = source_image.convert("RGBA")
alpha_min, alpha_max = rgba.getchannel("A").getextrema()
if (alpha_min, alpha_max) != (255, 255):
raise AssetError(
f"Transparency is not supported by the RGB565 optimizer: {source_path.name} "
f"has alpha range ({alpha_min}, {alpha_max})"
)
source_rgb = rgba.convert("RGB")
source_pixels = list(source_rgb.get_flattened_data())
rgb565_pixels = [expand_rgb565(pixel) for pixel in source_pixels]
color_count = len(set(rgb565_pixels))
if color_count <= 256:
output_mode = "indexed"
save_indexed_png(image_module, source_rgb.size, rgb565_pixels, output_path)
else:
output_mode = "rgb"
optimized_rgb = image_module.new("RGB", source_rgb.size)
optimized_rgb.putdata(rgb565_pixels)
optimized_rgb.save(output_path, format="PNG", optimize=True, compress_level=9)
assert_display_equivalent(source_rgb, output_path, image_module)
return {
"source_sha256": sha256_file(source_path),
"optimized_sha256": sha256_file(output_path),
"width": source_rgb.width,
"height": source_rgb.height,
"source_mode": source_mode,
"alpha_range": [alpha_min, alpha_max],
"rgb565_colors": color_count,
"output_mode": output_mode,
"source_bytes": source_path.stat().st_size,
"optimized_bytes": output_path.stat().st_size,
}
def generate() -> None:
image_module, pillow_version = load_pillow()
sources = png_files(SOURCE_DIR)
if not sources:
raise AssetError(f"No PNG assets found in {SOURCE_DIR}")
OPTIMIZED_DIR.parent.mkdir(parents=True, exist_ok=True)
temporary_dir = Path(tempfile.mkdtemp(prefix="lcd_optimized.", dir=OPTIMIZED_DIR.parent))
try:
assets: dict[str, Any] = {}
for source_path in sources:
output_path = temporary_dir / source_path.name
assets[source_path.name] = optimize_one(source_path, output_path, image_module)
manifest = {
"version": MANIFEST_VERSION,
"generator": "tools/optimize_lcd_assets.py",
"pillow_version": pillow_version,
"assets": assets,
"totals": {
"asset_count": len(assets),
"source_bytes": sum(item["source_bytes"] for item in assets.values()),
"optimized_bytes": sum(item["optimized_bytes"] for item in assets.values()),
},
}
manifest_path = temporary_dir / MANIFEST_NAME
manifest_path.write_text(
json.dumps(manifest, ensure_ascii=True, indent=2, sort_keys=True) + "\n",
encoding="utf-8",
)
if OPTIMIZED_DIR.exists():
shutil.rmtree(OPTIMIZED_DIR)
temporary_dir.replace(OPTIMIZED_DIR)
except Exception:
shutil.rmtree(temporary_dir, ignore_errors=True)
raise
reduction = 100.0 * (1.0 - manifest["totals"]["optimized_bytes"] / manifest["totals"]["source_bytes"])
print(
f"Generated {len(assets)} LCD assets: "
f"{manifest['totals']['source_bytes']} -> {manifest['totals']['optimized_bytes']} bytes "
f"({reduction:.1f}% reduction)"
)
def load_manifest() -> dict[str, Any]:
manifest_path = OPTIMIZED_DIR / MANIFEST_NAME
if not manifest_path.is_file():
raise AssetError(f"Asset manifest does not exist: {manifest_path}")
try:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
raise AssetError(f"Cannot read asset manifest: {manifest_path}: {exc}") from exc
if manifest.get("version") != MANIFEST_VERSION or not isinstance(manifest.get("assets"), dict):
raise AssetError(f"Unsupported or invalid asset manifest: {manifest_path}")
return manifest
def check_current() -> dict[str, Any]:
manifest = load_manifest()
assets = manifest["assets"]
source_names = {path.name for path in png_files(SOURCE_DIR)}
optimized_names = {path.name for path in png_files(OPTIMIZED_DIR)}
manifest_names = set(assets)
errors: list[str] = []
if source_names != manifest_names:
errors.append(
f"Source/manifest file set differs: missing={sorted(source_names - manifest_names)}, "
f"stale={sorted(manifest_names - source_names)}"
)
if optimized_names != manifest_names:
errors.append(
f"Optimized/manifest file set differs: missing={sorted(manifest_names - optimized_names)}, "
f"extra={sorted(optimized_names - manifest_names)}"
)
for name in sorted(source_names & manifest_names):
source_path = SOURCE_DIR / name
if sha256_file(source_path) != assets[name].get("source_sha256"):
errors.append(f"Source asset changed after optimization: {name}")
for name in sorted(optimized_names & manifest_names):
optimized_path = OPTIMIZED_DIR / name
if sha256_file(optimized_path) != assets[name].get("optimized_sha256"):
errors.append(f"Optimized asset changed after generation: {name}")
if errors:
detail = "\n - ".join(errors)
raise AssetError(
"LCD optimized assets are out of date:\n"
f" - {detail}\n"
"Regenerate them with:\n"
f" {sys.executable} {Path(__file__).resolve()} generate"
)
print(f"LCD assets are current: {len(manifest_names)} files")
return manifest
def verify() -> None:
manifest = check_current()
image_module, _ = load_pillow()
for name in sorted(manifest["assets"]):
with image_module.open(SOURCE_DIR / name) as source_image:
source_rgb = source_image.convert("RGB")
assert_display_equivalent(source_rgb, OPTIMIZED_DIR / name, image_module)
print(f"RGB565 display equivalence verified: {len(manifest['assets'])} files")
def check_env() -> None:
_, pillow_version = load_pillow()
print(f"Asset environment is ready: Python {sys.version.split()[0]}, Pillow {pillow_version}")
def parse_args(arguments: Iterable[str]) -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"command",
choices=("check-env", "generate", "check-current", "verify"),
help="operation to perform",
)
return parser.parse_args(arguments)
def main(arguments: Iterable[str] | None = None) -> int:
args = parse_args(sys.argv[1:] if arguments is None else arguments)
try:
{
"check-env": check_env,
"generate": generate,
"check-current": check_current,
"verify": verify,
}[args.command]()
except AssetError as exc:
print(f"error: {exc}", file=sys.stderr)
return 1
return 0
if __name__ == "__main__":
raise SystemExit(main())