| import argparse |
| from pathlib import Path |
| import json |
| import matplotlib.pyplot as plt |
| import re |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("dirs", nargs="+", help="One or more directories to process") |
| args = parser.parse_args() |
|
|
| plt.figure(figsize=(8, 6)) |
|
|
| for dir_path in args.dirs: |
| dir_path = Path(dir_path) |
| name = dir_path.name |
|
|
| step_numbers = [] |
| avg_psnrs = [] |
|
|
| step_dirs = [d for d in dir_path.iterdir() if re.fullmatch(r"iter_\d+", d.name)] |
| for step_dir in sorted(step_dirs, key=lambda x: int(x.name.split("_")[-1])): |
| if len(list(step_dir.iterdir())) == 0: |
| continue |
| if len(list(step_dir.iterdir())) != 501: |
| continue |
| psnrs = [] |
| for f in Path(step_dir).glob("*_metrics.json"): |
| with open(f, "r") as file: |
| data = json.load(file) |
| psnrs.append(data["psnr"]) |
| avg_psnr = sum(psnrs) / len(psnrs) if psnrs else 0.0 |
|
|
| step_num = int(step_dir.name.split("_")[-1]) |
| step_numbers.append(step_num) |
| avg_psnrs.append(avg_psnr) |
|
|
| print(f"Directory: {step_dir}") |
| print(f"Average PSNR: {avg_psnr:.2f} dB ({len(psnrs)} samples)") |
|
|
| plt.plot(step_numbers, avg_psnrs, marker="o", label=name) |
|
|
| plt.xlabel("Step", fontsize=14) |
| plt.ylabel("Average PSNR (dB)", fontsize=14) |
| plt.title("Average PSNR per Step", fontsize=16) |
| plt.grid(True) |
| plt.tight_layout(rect=[0, 0.05, 1, 1]) |
| plt.legend( |
| loc="upper center", |
| bbox_to_anchor=(0.5, -0.1), |
| frameon=False, |
| fontsize=14 |
| ) |
|
|
| plt.savefig("psnr.png", bbox_inches="tight") |
|
|
