|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +from dataclasses import dataclass |
| 4 | +from pathlib import Path |
| 5 | + |
| 6 | +from PIL import Image, ImageDraw, ImageFont |
| 7 | +from playwright.sync_api import sync_playwright |
| 8 | + |
| 9 | + |
| 10 | +SCRIPTS_DIR = Path(__file__).resolve().parent |
| 11 | +POST_DIR = SCRIPTS_DIR.parent |
| 12 | +ROOT = POST_DIR.parents[3] |
| 13 | +HTML = SCRIPTS_DIR / "sr_rounding_snapshot.html" |
| 14 | +TMP = ROOT / "tmp" |
| 15 | + |
| 16 | + |
| 17 | +@dataclass(frozen=True) |
| 18 | +class PanelSpec: |
| 19 | + title: str |
| 20 | + r_bits: int |
| 21 | + centered: bool = False |
| 22 | + min_samples: int = 200000 |
| 23 | + |
| 24 | + |
| 25 | +def panel_specs(mode: str) -> list[PanelSpec]: |
| 26 | + if mode == "floor": |
| 27 | + return [ |
| 28 | + PanelSpec("Floor SR, R = 5", r_bits=5, centered=False), |
| 29 | + PanelSpec("Floor SR, R = 4", r_bits=4, centered=False), |
| 30 | + PanelSpec("Floor SR, R = 2", r_bits=2, centered=False), |
| 31 | + ] |
| 32 | + if mode == "centered": |
| 33 | + return [ |
| 34 | + PanelSpec("Centered SR, R = 5", r_bits=5, centered=True), |
| 35 | + PanelSpec("Centered SR, R = 4", r_bits=4, centered=True), |
| 36 | + PanelSpec("Centered SR, R = 2", r_bits=2, centered=True), |
| 37 | + ] |
| 38 | + raise ValueError(f"unknown mode: {mode}") |
| 39 | + |
| 40 | + |
| 41 | +def get_font(size: int, bold: bool = False): |
| 42 | + candidates = [ |
| 43 | + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" |
| 44 | + if bold |
| 45 | + else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", |
| 46 | + "/usr/share/fonts/dejavu/DejaVuSans-Bold.ttf" |
| 47 | + if bold |
| 48 | + else "/usr/share/fonts/dejavu/DejaVuSans.ttf", |
| 49 | + ] |
| 50 | + for candidate in candidates: |
| 51 | + path = Path(candidate) |
| 52 | + if path.exists(): |
| 53 | + return ImageFont.truetype(str(path), size=size) |
| 54 | + return ImageFont.load_default() |
| 55 | + |
| 56 | + |
| 57 | +BG = (255, 254, 250) |
| 58 | +INK = (35, 38, 41) |
| 59 | +BOX = (80, 88, 98) |
| 60 | +BOX_FILL = (255, 254, 250) |
| 61 | +GAP = 28 |
| 62 | +PAD_X = 24 |
| 63 | +PAD_Y = 18 |
| 64 | +X_MIN = 2.9 |
| 65 | +X_MAX = 7.0 |
| 66 | +Y_MIN = 2.4 |
| 67 | +Y_MAX = 7.6 |
| 68 | +ZOOM_X0 = 3.7 |
| 69 | +ZOOM_X1 = 4.7 |
| 70 | +ZOOM_Y0 = 3.7 |
| 71 | +ZOOM_Y1 = 4.7 |
| 72 | + |
| 73 | + |
| 74 | +def wait_for_samples(page, threshold: int): |
| 75 | + page.wait_for_function( |
| 76 | + """threshold => { |
| 77 | + const el = document.getElementById("sr-rounding-seed"); |
| 78 | + if (!el) return false; |
| 79 | + const m = /Samples:\\s*([0-9]+)/.exec(el.textContent || ""); |
| 80 | + return m && Number(m[1]) >= threshold; |
| 81 | + }""", |
| 82 | + arg=threshold, |
| 83 | + timeout=20000, |
| 84 | + ) |
| 85 | + |
| 86 | + |
| 87 | +def capture_panel(page, spec: PanelSpec, out_path: Path): |
| 88 | + page.get_by_role("button", name="Reset").click() |
| 89 | + if spec.centered: |
| 90 | + page.locator("#sr-rounding-mode-centered").check() |
| 91 | + else: |
| 92 | + page.locator("#sr-rounding-mode-floor").check() |
| 93 | + page.locator("#sr-rounding-demo").evaluate( |
| 94 | + "(el, label) => { el.dataset.meanLegendLabel = label; }", |
| 95 | + f"SR, {'Centered' if spec.centered else 'Floor'}, R = {spec.r_bits}", |
| 96 | + ) |
| 97 | + r_input = page.locator("#sr-rounding-rbits") |
| 98 | + r_input.fill(str(spec.r_bits)) |
| 99 | + r_input.blur() |
| 100 | + page.get_by_role("button", name="Ffwd").click() |
| 101 | + wait_for_samples(page, spec.min_samples) |
| 102 | + page.get_by_role("button", name="Pause").click() |
| 103 | + page.locator(".sr-rounding-canvas-wrap").screenshot(path=str(out_path)) |
| 104 | + |
| 105 | + |
| 106 | +def sx(image_width: int, x: float) -> float: |
| 107 | + left = 58 |
| 108 | + right = 20 |
| 109 | + plot_width = image_width - left - right |
| 110 | + return left + (x - X_MIN) / (X_MAX - X_MIN) * plot_width |
| 111 | + |
| 112 | + |
| 113 | +def sy(image_height: int, y: float) -> float: |
| 114 | + top = 22 |
| 115 | + bottom = 58 |
| 116 | + plot_height = image_height - top - bottom |
| 117 | + return top + (1.0 - (y - Y_MIN) / (Y_MAX - Y_MIN)) * plot_height |
| 118 | + |
| 119 | + |
| 120 | +def add_zoom_inset(panel: Image.Image) -> Image.Image: |
| 121 | + panel = panel.copy() |
| 122 | + draw = ImageDraw.Draw(panel) |
| 123 | + x0 = sx(panel.width, ZOOM_X0) |
| 124 | + x1 = sx(panel.width, ZOOM_X1) |
| 125 | + y0 = sy(panel.height, ZOOM_Y1) |
| 126 | + y1 = sy(panel.height, ZOOM_Y0) |
| 127 | + crop_box = ( |
| 128 | + int(round(x0)), |
| 129 | + int(round(y0)), |
| 130 | + int(round(x1)), |
| 131 | + int(round(y1)), |
| 132 | + ) |
| 133 | + draw.rectangle(crop_box, outline=BOX, width=3) |
| 134 | + |
| 135 | + crop = panel.crop(crop_box) |
| 136 | + inset_w = int(panel.width * 0.44) |
| 137 | + inset_h = int(inset_w * crop.height / crop.width) |
| 138 | + inset = crop.resize((inset_w, inset_h), Image.Resampling.BICUBIC) |
| 139 | + inset_box = ( |
| 140 | + panel.width - inset_w - 92, |
| 141 | + panel.height - inset_h - 220, |
| 142 | + ) |
| 143 | + outer_box = [ |
| 144 | + inset_box[0], |
| 145 | + inset_box[1], |
| 146 | + inset_box[0] + inset_w, |
| 147 | + inset_box[1] + inset_h, |
| 148 | + ] |
| 149 | + panel.paste(inset, inset_box) |
| 150 | + draw.rectangle(outer_box, outline=BOX, width=3) |
| 151 | + src_tl = (crop_box[0], crop_box[1]) |
| 152 | + src_bl = (crop_box[0], crop_box[3]) |
| 153 | + inset_tl = (outer_box[0], outer_box[1]) |
| 154 | + inset_bl = (outer_box[0], outer_box[3]) |
| 155 | + draw.line([src_tl, inset_tl], fill=BOX, width=3) |
| 156 | + draw.line([src_bl, inset_bl], fill=BOX, width=3) |
| 157 | + return panel |
| 158 | + |
| 159 | + |
| 160 | +def panel_output_path(mode: str, r_bits: int) -> Path: |
| 161 | + return POST_DIR / "img" / f"sr-rbits-{mode}-r{r_bits}.png" |
| 162 | + |
| 163 | + |
| 164 | +def main(): |
| 165 | + import argparse |
| 166 | + |
| 167 | + parser = argparse.ArgumentParser() |
| 168 | + parser.add_argument("--mode", choices=["floor", "centered"], default="floor") |
| 169 | + args = parser.parse_args() |
| 170 | + |
| 171 | + panels_spec = panel_specs(args.mode) |
| 172 | + out_path = POST_DIR / "img" / f"sr-rbits-5-4-2-{args.mode}.png" |
| 173 | + |
| 174 | + TMP.mkdir(parents=True, exist_ok=True) |
| 175 | + out_path.parent.mkdir(parents=True, exist_ok=True) |
| 176 | + shot_paths = [ |
| 177 | + TMP / f"sr-snapshot-{args.mode}-{i}.png" for i in range(len(panels_spec)) |
| 178 | + ] |
| 179 | + |
| 180 | + with sync_playwright() as p: |
| 181 | + browser = p.chromium.launch() |
| 182 | + page = browser.new_page( |
| 183 | + viewport={"width": 1040, "height": 980}, device_scale_factor=2 |
| 184 | + ) |
| 185 | + page.goto(HTML.resolve().as_uri()) |
| 186 | + page.wait_for_selector("#sr-rounding-demo.js-ready", timeout=10000) |
| 187 | + for spec, path in zip(panels_spec, shot_paths): |
| 188 | + capture_panel(page, spec, path) |
| 189 | + browser.close() |
| 190 | + |
| 191 | + panels = [add_zoom_inset(Image.open(path).convert("RGB")) for path in shot_paths] |
| 192 | + for panel, spec in zip(panels, panels_spec): |
| 193 | + panel.save(panel_output_path(args.mode, spec.r_bits), "PNG") |
| 194 | + width = PAD_X * 2 + sum(panel.width for panel in panels) + GAP * (len(panels) - 1) |
| 195 | + height = PAD_Y * 2 + max(panel.height for panel in panels) |
| 196 | + canvas = Image.new("RGB", (width, height), BG) |
| 197 | + |
| 198 | + x = PAD_X |
| 199 | + top_y = PAD_Y |
| 200 | + for panel in panels: |
| 201 | + canvas.paste(panel, (x, top_y)) |
| 202 | + x += panel.width + GAP |
| 203 | + canvas.save(out_path, "PNG") |
| 204 | + |
| 205 | + for panel in panels: |
| 206 | + panel.close() |
| 207 | + for path in shot_paths: |
| 208 | + path.unlink(missing_ok=True) |
| 209 | + |
| 210 | + print(out_path) |
| 211 | + |
| 212 | + |
| 213 | +if __name__ == "__main__": |
| 214 | + main() |
0 commit comments