Files
jiapuapp/design-pipeline/scripts/build_g01_backgrounds.py
T
2026-07-17 08:08:33 +08:00

129 lines
5.0 KiB
Python

"""Deterministically build G01 background candidates from saved ImageGen masters."""
import argparse
import hashlib
import json
from pathlib import Path
from typing import Any
from PIL import Image, PngImagePlugin
def process_chroma_image(source: Image.Image, config: dict[str, Any]) -> Image.Image:
"""Convert a green-screen master to clean RGBA with edge despill."""
start = int(config["dominanceStart"])
end = int(config["dominanceEnd"])
allowance = int(config["despillAllowance"])
if end <= start:
raise ValueError("dominanceEnd must be greater than dominanceStart")
output = Image.new("RGBA", source.size, (0, 0, 0, 0))
converted = []
for red, green, blue, source_alpha in source.convert("RGBA").get_flattened_data():
dominance = green - max(red, blue)
if green >= 100 and dominance > start:
removal = min(1.0, (dominance - start) / (end - start))
alpha = round(source_alpha * (1.0 - removal))
if alpha <= 2:
converted.append((0, 0, 0, 0))
continue
green = min(green, max(red, blue) + allowance)
converted.append((red, green, blue, alpha))
else:
converted.append((red, green, blue, source_alpha))
output.putdata(converted)
return output
def process_opaque_image(source: Image.Image, output_size: tuple[int, int] | None = None) -> Image.Image:
"""Normalize a paper-backed master to deterministic opaque RGBA."""
output = source.convert("RGBA")
if output_size is not None and output.size != output_size:
output = output.resize(output_size, Image.Resampling.LANCZOS)
output.putalpha(255)
return output
def save_png(image: Image.Image, path: Path) -> None:
"""Save a deterministic PNG with an explicit standard-sRGB chunk."""
path.parent.mkdir(parents=True, exist_ok=True)
png_info = PngImagePlugin.PngInfo()
png_info.add(b"sRGB", b"\x00")
image.save(path, format="PNG", optimize=True, compress_level=9, pnginfo=png_info)
def _resolve_inside(workspace: Path, relative_path: str) -> Path:
absolute = (workspace / relative_path).resolve()
try:
absolute.relative_to(workspace.resolve())
except ValueError as error:
raise ValueError(f"path escapes workspace: {relative_path}") from error
return absolute
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def build(manifest_path: Path, workspace: Path, report_path: Path) -> list[dict[str, Any]]:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
chroma = manifest["chromaKey"]
reports = []
for candidate in manifest["candidates"]:
source_path = _resolve_inside(workspace, candidate["master"])
output_path = _resolve_inside(workspace, candidate["generatedOutput"])
with Image.open(source_path) as opened:
expected = candidate["sourcePixels"]
if opened.size != (expected["width"], expected["height"]):
raise ValueError(
f"{candidate['id']} expected {expected['width']}x{expected['height']}, got {opened.width}x{opened.height}"
)
if candidate["processingMode"] == "chroma-key":
output = process_chroma_image(opened, chroma)
elif candidate["processingMode"] in {"opaque-paper", "opaque-paper-resize"}:
output_size = None
if candidate["processingMode"] == "opaque-paper-resize":
target = candidate["outputPixels"]
output_size = (target["width"], target["height"])
output = process_opaque_image(opened, output_size)
else:
raise ValueError(f"unsupported processingMode: {candidate['processingMode']}")
save_png(output, output_path)
reports.append(
{
"id": candidate["id"],
"mode": candidate["processingMode"],
"source": candidate["master"],
"output": candidate["generatedOutput"],
"width": output.width,
"height": output.height,
"bytes": output_path.stat().st_size,
"sha256": _sha256(output_path),
}
)
report_path.parent.mkdir(parents=True, exist_ok=True)
report_path.write_text(json.dumps(reports, ensure_ascii=False, indent=2), encoding="utf-8")
return reports
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("manifest", type=Path)
parser.add_argument("--workspace", type=Path, required=True)
parser.add_argument("--report", type=Path, required=True)
args = parser.parse_args()
reports = build(args.manifest, args.workspace.resolve(), args.report)
for report in reports:
print(f"G01-BACKGROUND-BUILD PASS {report['id']} {report['width']}x{report['height']} {report['sha256']}")
if __name__ == "__main__":
main()