| |
| 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() |
|
|
