File size: 3,555 Bytes
da43581 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | # E-OS v6 — eval: PSNR on fresh val scenes @ app-res + preview grids.
import os, sys, json, math
import numpy as np
import torch
sys.path.insert(0, "/home/z/my-project/scripts/os6")
from os6_models import FrameGenerator
from os6_draw import render_frame, rasterize_hint, build_state
from os6_spec import lay_home, lay_random, CARD, CIRCLE
from PIL import Image
CKPT = sys.argv[1] if len(sys.argv) > 1 else \
"/home/z/my-project/checkpoints6/gen_ckpt_final.pt"
OUT = "/home/z/my-project/download/e-os-v6/eval"
os.makedirs(OUT, exist_ok=True)
m = FrameGenerator()
ck = torch.load(CKPT, map_location="cpu", weights_only=False)
m.load_state_dict(ck["m"])
m.eval()
print("ckpt step %s best %s" % (ck.get("step"), ck.get("best")))
rng = np.random.default_rng(999)
def run(shapes, theme, day, shimmer, ripples, pressed=0):
hint = rasterize_hint(shapes, ripples, 256, 544)
st = build_state(theme, shimmer, day, pressed)
with torch.no_grad():
p = m(torch.from_numpy(hint).unsqueeze(0),
torch.from_numpy(st).unsqueeze(0))
return p[0].permute(1, 2, 0).numpy()
def gt(shapes, theme, day, shimmer, ripples, pressed=0):
return render_frame(shapes, theme, day, shimmer, ripples, pressed) / 255.0
# ---- 40 fresh val scenes ----
ps = []
samples = []
for k in range(40):
shapes = (lay_home if k % 3 == 0 else
lay_random if k % 3 == 1 else
(lay_home if k % 2 else lay_random))(rng)
shapes.sort(key=lambda s: s.get("z", 0))
theme = int(rng.integers(0, 4))
day = float(rng.uniform(0.05, 0.95))
shimmer = float(rng.uniform(0, 1))
rip = (float(rng.uniform(0.15, 0.85)), float(rng.uniform(0.15, 0.85)),
float(rng.uniform(0.05, 0.9))) if k % 4 != 3 else None
args = (shapes, theme, day, shimmer, [rip] if rip else [])
g = gt(*args)
p = np.clip(run(*args), 0, 1)
mse = ((g - p) ** 2).mean()
ps.append(10 * math.log10(1.0 / max(1e-8, mse)))
if k < 6:
samples.append((shapes, theme, rip, g, p))
ps = np.array(ps)
print("PSNR mean %.2f median %.2f (40 scenes @ 512x1088)" %
(ps.mean(), np.median(ps)))
# ---- preview grid: GT | NPU-Output (6 rows) ----
rows = []
for (shapes, theme, rip, g, p) in samples:
rows.append(np.concatenate([g, p], axis=1))
grid = (np.concatenate(rows, axis=0) * 255).astype(np.uint8)
Image.fromarray(grid).save(os.path.join(OUT, "eval_gt_vs_npu.png"))
# ---- GENERALIZATION PROOF: edited layout the net never saw ----
sh = lay_home(rng)
sh[0] = dict(sh[0], x0=100, x1=430, y0=90, y1=330) # card moved/resized
sh.append(dict(t=CIRCLE, x0=60, y0=430, x1=150, y1=520, z=9)) # extra circle
sh.append(dict(t=CARD, x0=210, y0=560, x1=470, y1=700, z=8)) # extra card
sh.sort(key=lambda s: s.get("z", 0))
theme = 1
rip = (0.4, 0.5, 0.35)
ge = gt(sh, theme, 0.5, 0.4, [rip])
pe = np.clip(run(sh, theme, 0.5, 0.4, [rip]), 0, 1)
mse = ((ge - pe) ** 2).mean()
psnr_edit = 10 * math.log10(1.0 / max(1e-8, mse))
hint_e = rasterize_hint(sh, [rip], 256, 544)
wire = np.stack([hint_e[0]] * 3, -1)
proof = (np.concatenate([ge, pe], axis=1) * 255).astype(np.uint8)
Image.fromarray(proof).save(os.path.join(OUT, "eval_generalization.png"))
print("EDITED-LAYOUT PSNR %.2f dB (proof of generalization)" % psnr_edit)
json.dump(dict(psnr_mean=float(ps.mean()), psnr_median=float(np.median(ps)),
psnr_edited_layout=float(psnr_edit), n=40, app_res=[512, 1088],
ckpt_step=ck.get("step")),
open(os.path.join(OUT, "eval_metrics.json"), "w"), indent=1)
print("EVAL DONE")
|