File size: 2,227 Bytes
c978bcc | 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 | # E-OS v4 — evaluation: full-res (app) PSNR on fresh renders + preview sheet.
import os, sys, math, json, random
import numpy as np
import torch
sys.path.insert(0, "/home/z/my-project/scripts/os4")
from os4_models import FrameGenerator
import os4_spec as S
import os4_draw as D
from gen_os4_dataset import build_vec
from PIL import Image
CKPT = "/home/z/my-project/checkpoints4/gen_ckpt_final.pt"
OUT = "/home/z/my-project/download/e-os-v4"
os.makedirs(OUT, exist_ok=True)
def main():
m = FrameGenerator()
m.load_state_dict(torch.load(CKPT, map_location="cpu", weights_only=False)["m"])
m.eval()
rng = random.Random(4242)
ps, previews = [], []
N = 40
with torch.no_grad():
for i in range(N):
st = __import__("gen_os4_dataset").gen_scene(900000 + i)
hint = np.asarray(D.draw_scene(st, "hint"), np.float32) / 255.0
gt = np.asarray(D.draw_scene(st, "styled"), np.float32) / 255.0
h = torch.from_numpy(hint).permute(2, 0, 1).unsqueeze(0)
s = torch.from_numpy(build_vec(st)).unsqueeze(0)
pred = m(h, s)[0].permute(1, 2, 0).numpy()
mse = ((pred - gt) ** 2).mean()
ps.append(10 * math.log10(1.0 / max(1e-9, mse)))
if i % 8 == 0:
h2 = np.asarray(Image.fromarray((hint * 255).astype(np.uint8))
.resize((S.FULL_W, S.FULL_H)), np.float32) / 255.0
row = np.concatenate([h2, gt, np.clip(pred, 0, 1)], axis=1)
previews.append((row * 255).astype(np.uint8))
ps = np.array(ps)
print("app-res PSNR: mean %.2f dB | median %.2f | p10 %.2f | min %.2f (N=%d)" %
(ps.mean(), np.median(ps), np.percentile(ps, 10), ps.min(), N))
sheet = np.concatenate(previews, axis=0)
Image.fromarray(sheet).save(os.path.join(OUT, "eval_preview_grid.png"))
json.dump({"psnr_mean": float(ps.mean()), "psnr_median": float(np.median(ps)),
"psnr_p10": float(np.percentile(ps, 10)), "psnr_min": float(ps.min()),
"n": N}, open(os.path.join(OUT, "eval_metrics.json"), "w"), indent=1)
print("saved", os.path.join(OUT, "eval_preview_grid.png"))
if __name__ == "__main__":
main()
|