Codex
79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
|
|
|
|
BACKGROUND_COLORS: dict[str, tuple[int, int, int]] = {
|
|
"black": (0, 0, 0),
|
|
"white": (255, 255, 255),
|
|
"gray": (128, 128, 128),
|
|
"red": (220, 32, 32),
|
|
"blue": (32, 96, 220),
|
|
}
|
|
|
|
|
|
def composite(rgb: np.ndarray, alpha: np.ndarray, color: tuple[int, int, int]) -> np.ndarray:
|
|
fg = rgb.astype(np.float32)
|
|
bg = np.zeros_like(fg)
|
|
bg[..., 0] = color[0]
|
|
bg[..., 1] = color[1]
|
|
bg[..., 2] = color[2]
|
|
a = alpha[..., None].astype(np.float32)
|
|
return np.clip(fg * a + bg * (1.0 - a), 0, 255).astype(np.uint8)
|
|
|
|
|
|
def save_previews(debug_dir: str | Path, rgb: np.ndarray, alpha: np.ndarray) -> dict[str, Path]:
|
|
debug = Path(debug_dir)
|
|
debug.mkdir(parents=True, exist_ok=True)
|
|
paths: dict[str, Path] = {}
|
|
for name, color in BACKGROUND_COLORS.items():
|
|
path = debug / f"preview_{name}.png"
|
|
Image.fromarray(composite(rgb, alpha, color), mode="RGB").save(path)
|
|
paths[name] = path
|
|
return paths
|
|
|
|
|
|
def make_qa_grid(debug_dir: str | Path, tiles: dict[str, np.ndarray]) -> Path:
|
|
debug = Path(debug_dir)
|
|
debug.mkdir(parents=True, exist_ok=True)
|
|
tile_items = list(tiles.items())
|
|
if not tile_items:
|
|
raise ValueError("No QA tiles provided")
|
|
|
|
thumb_w = 320
|
|
label_h = 28
|
|
padding = 10
|
|
cols = min(3, len(tile_items))
|
|
rows = int(np.ceil(len(tile_items) / cols))
|
|
first_h, first_w = tile_items[0][1].shape[:2]
|
|
thumb_h = max(180, int(round(thumb_w * first_h / max(first_w, 1))))
|
|
|
|
canvas = Image.new(
|
|
"RGB",
|
|
(
|
|
cols * thumb_w + (cols + 1) * padding,
|
|
rows * (thumb_h + label_h) + (rows + 1) * padding,
|
|
),
|
|
(32, 32, 32),
|
|
)
|
|
draw = ImageDraw.Draw(canvas)
|
|
|
|
for idx, (name, arr) in enumerate(tile_items):
|
|
row = idx // cols
|
|
col = idx % cols
|
|
x = padding + col * (thumb_w + padding)
|
|
y = padding + row * (thumb_h + label_h + padding)
|
|
img = Image.fromarray(arr.astype(np.uint8), mode="RGB")
|
|
img.thumbnail((thumb_w, thumb_h), Image.Resampling.LANCZOS)
|
|
frame = Image.new("RGB", (thumb_w, thumb_h), (240, 240, 240))
|
|
frame.paste(img, ((thumb_w - img.width) // 2, (thumb_h - img.height) // 2))
|
|
canvas.paste(frame, (x, y + label_h))
|
|
draw.text((x, y), name, fill=(245, 245, 245))
|
|
|
|
path = debug / "qa_grid.png"
|
|
canvas.save(path)
|
|
return path
|