safegen-arena / scripts /build_demo_stack.py
Somin-Aggarwal's picture
Refresh demo grid with new before/after pairs
aea64b6
"""
Stack the three before/after pair images in `images/` into a single
vertical grid and write it to `assets/demo_images/figure_4_grid.png`,
replacing the previous demo grid embedded in the README and Blog.
Order top → bottom:
1. violence_01_streetfight_pair
2. violence_02_hostage_pair
3. selfharm_01_bathtub_pair
All three are scaled to a common width (1054 px — the native width of two
of the three; the third is downscaled from 1536). RGBA is flattened over
white. A 16 px white gap separates rows.
"""
from pathlib import Path
from PIL import Image
ROOT = Path(__file__).resolve().parent.parent
SRC = [
ROOT / "images" / "violence_01_streetfight_pair.png",
ROOT / "images" / "violence_02_hostage_pair.png",
ROOT / "images" / "selfharm_01_bathtub_pair.png",
]
OUT = ROOT / "assets" / "demo_images" / "figure_4_grid.png"
TARGET_WIDTH = 1054
GAP = 16
BG = (255, 255, 255)
def load_rgb(path: Path) -> Image.Image:
im = Image.open(path)
if im.mode == "RGBA":
bg = Image.new("RGB", im.size, BG)
bg.paste(im, mask=im.split()[3])
return bg
return im.convert("RGB")
def fit_width(im: Image.Image, w: int) -> Image.Image:
if im.width == w:
return im
h = round(im.height * w / im.width)
return im.resize((w, h), Image.LANCZOS)
def main() -> None:
panels = [fit_width(load_rgb(p), TARGET_WIDTH) for p in SRC]
total_h = sum(p.height for p in panels) + GAP * (len(panels) - 1)
canvas = Image.new("RGB", (TARGET_WIDTH, total_h), BG)
y = 0
for p in panels:
canvas.paste(p, (0, y))
y += p.height + GAP
OUT.parent.mkdir(parents=True, exist_ok=True)
canvas.save(OUT, format="PNG", optimize=True)
print(f"wrote {OUT} ({canvas.size[0]}x{canvas.size[1]}, "
f"{OUT.stat().st_size // 1024} KB)")
if __name__ == "__main__":
main()