thunderheronight--splat-workshop/scripts/make_gallery.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["pillow"]
# ///
"""Create the preview image (512x512) and name plate (358x86) for the Gaussian XR gallery.

Usage: uv run make_gallery.py <run folder> <pipeline log> <display name> <file stem>
       uv run make_gallery.py --plan <floor plan.png> <run folder> <display name> <file stem>
       uv run make_gallery.py --label <name plate.png> <text>
Reads the initial view from the log ("Gaussian XR initial view: cam0/... Display rotation X°")
and writes <file stem>.gallery.png and <file stem>.sceneinfo.png into the run folder.
The look matches the example scenes that ship with Gaussian XR.
"""
import re
import shutil
import sys
from pathlib import Path

from PIL import Image, ImageDraw, ImageEnhance, ImageFilter, ImageFont, ImageStat

FONTS = [
    "/usr/share/fonts/google-noto/NotoSans-Regular.ttf",
    "/usr/share/fonts/dejavu-sans-fonts/DejaVuSans.ttf",
    "/usr/share/fonts/liberation-sans/LiberationSans-Regular.ttf",
    "/usr/share/fonts/truetype/noto/NotoSans-Regular.ttf",          # Debian/Ubuntu
    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    "/usr/share/fonts/noto/NotoSans-Regular.ttf",                   # Arch
    "/usr/share/fonts/TTF/DejaVuSans.ttf",
]


def initial_view(run: Path, log: Path) -> tuple[Path, float]:
    match = None
    for match in re.finditer(
        r"Gaussian XR initial view: (\S+\.jpg) .*?Display rotation (-?[\d.]+)°",
        log.read_text(errors="replace"),
    ):
        pass  # the last message counts
    if match:
        image = next(run.glob(f"dataset/images/{match.group(1)}"), None)
        if image:
            return image, float(match.group(2))
    # No message: an image from the middle of the recording, not rotated
    images = sorted(run.glob("dataset/images/cam0/*.jpg"))
    if not images:
        sys.exit(f"No images found in {run}/dataset/images/cam0")
    return images[len(images) // 2], 0.0


def gallery(image: Path, rotation_ccw: float, out: Path) -> None:
    # PIL rotates counter-clockwise for positive angles; the pipeline already reports the CCW rotation
    img = Image.open(image).convert("RGB").rotate(rotation_ccw, expand=True)
    w, h = img.size
    side = int(min(w, h) * 0.8)  # cut away the fisheye border
    img = img.crop(((w - side) // 2, (h - side) // 2, (w + side) // 2, (h + side) // 2))
    img = img.resize((512, 512), Image.LANCZOS)
    # Brighten dark indoor shots without blowing out bright ones
    brightness = ImageStat.Stat(img.convert("L")).mean[0]
    if brightness < 110:
        img = ImageEnhance.Brightness(img).enhance(min(2.2, 110 / max(brightness, 1)))
    img = ImageEnhance.Contrast(img).enhance(1.1)
    img.convert("RGBA").save(out)


def label(text: str, out: Path) -> None:
    width, height = 358, 86
    img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    shadow = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    ImageDraw.Draw(shadow).rounded_rectangle(
        (10, 10, width - 10, height - 10), radius=16, fill=(40, 46, 60, 200))
    img.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(4)))
    draw = ImageDraw.Draw(img)
    draw.rounded_rectangle((14, 12, width - 14, height - 14), radius=14, fill=(64, 72, 90, 255))
    font_path = next((p for p in FONTS if Path(p).exists()), None)
    size = 34
    while True:  # shrink long names until they fit on the plate
        font = ImageFont.truetype(font_path, size) if font_path else ImageFont.load_default()
        box = draw.textbbox((0, 0), text, font=font)
        if box[2] - box[0] <= width - 50 or size <= 14 or not font_path:
            break
        size -= 2
    draw.text(
        ((width - (box[2] - box[0])) / 2 - box[0], (height - (box[3] - box[1])) / 2 - box[1] - 1),
        text, font=font, fill=(255, 255, 255, 255))
    img.save(out)


def from_plan(plan: Path, out: Path) -> None:
    """Preview of a combined scene: its floor plan, centred on a square."""
    img = Image.open(plan).convert("RGB")
    side = max(img.size)
    square = Image.new("RGB", (side, side), img.getpixel((5, 5)))
    square.paste(img, ((side - img.width) // 2, (side - img.height) // 2))
    square.resize((512, 512), Image.LANCZOS).convert("RGBA").save(out)


def kept_preview(run: Path) -> Path | None:
    """The newest preview image in the run folder, or – for a version (<scene>/versions/<key>) – the scene's."""
    places = [run] + ([run.parent.parent] if run.parent.name == "versions" else [])
    for place in places:
        found = sorted(place.glob("*.gallery.png"), key=lambda p: p.stat().st_mtime, reverse=True)
        if found:
            return found[0]
    return None


def main() -> None:
    if len(sys.argv) == 4 and sys.argv[1] == "--label":
        label(sys.argv[3], Path(sys.argv[2]))
        return
    if len(sys.argv) == 6 and sys.argv[1] == "--plan":
        plan, run, name, stem = Path(sys.argv[2]), Path(sys.argv[3]), sys.argv[4], sys.argv[5]
        from_plan(plan, run / f"{stem}.gallery.png")
        label(name, run / f"{stem}.sceneinfo.png")
        print(f"Preview image from the floor plan, name plate \"{name}\"")
        return
    if len(sys.argv) != 5:
        sys.exit(__doc__)
    run, log, name, stem = Path(sys.argv[1]), Path(sys.argv[2]), sys.argv[3], sys.argv[4]
    preview = run / f"{stem}.gallery.png"
    if not preview.exists() and not any(run.glob("dataset/images/cam0/*.jpg")):
        # No picture under this stem and nothing to make one from (a version has no images of its own, and a
        # renamed scene's versions still carry the old stem): take the one already there, else the scene's
        if old := kept_preview(run):
            shutil.copyfile(old, preview)
    if preview.exists() and not any(run.glob("dataset/images/cam0/*.jpg")):
        # Training images were cleaned up: keep the existing preview, only renew the name plate
        label(name, run / f"{stem}.sceneinfo.png")
        print(f"Name plate \"{name}\" (preview image kept)")
        return
    image, rotation = initial_view(run, log)
    gallery(image, rotation, preview)
    label(name, run / f"{stem}.sceneinfo.png")
    print(f"Preview image from {image.name} (rotation {rotation:g}°), name plate \"{name}\"")


if __name__ == "__main__":
    main()