emilyxuan/accesscontrol
0
1import os2import sys3import csv4import shutil5from collections import OrderedDict6 7import numpy as np8import torch9import torch.nn.functional as F10from PIL import Image, ImageDraw, ImageFont11from torchvision import transforms12from torchvision.utils import save_image13 14SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))15PROJECT_DIR = os.path.dirname(SCRIPT_DIR)16sys.path.insert(0, PROJECT_DIR)17 18from model import UNet19from diffusion import GaussianDiffusionSampler20from trigger_encoder import TriggerEncoder21 22try:23 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity24 from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure25 TORCHMETRICS_OK = True26except Exception as e:27 TORCHMETRICS_OK = False28 TORCHMETRICS_IMPORT_ERROR = repr(e)29 30 31# ============================================================32# 0. 主要修改區33# ============================================================34 35DEVICE = "cuda:1" if torch.cuda.is_available() else "cpu"36 37SAVE_DIR = "./logs/cifar10_cond_acsctrl/cond128_neg0.25_0.25_seed42_Bsize128/cond128_neg0.5_seed42_20260511_000804/eval_/eval_embedding_sensitivity_l2_ckpt500"38 39BASELINE_CKPT_PATH = "./logs/cifar10_cond_acsctrl/cond128_baseline/cond128_neg0.2_seed42_20260506_221046/checkpoints/global_ckpt_round500.pt"40COND_CKPT_PATH = "./logs/cifar10_cond_acsctrl/cond128_neg0.25_0.25_seed42_Bsize128/cond128_neg0.5_seed42_20260511_000804/checkpoints/global_ckpt_round500.pt"41 42CKPT_NAME = "ckpt500"43 44CORRECT_CLIENT_IDS = [0, 1, 2, 3, 4]45 46FALLBACK_TRIGGER_DIR = "./logs/cifar10_cond_acsctrl/cond128_neg0.25_0.25_seed42_Bsize128/cond128_neg0.5_seed42_20260511_000804/client_triggers"47 48# 每個 synthetic embedding sample 幾張圖49NUM_SAMPLES_PER_EMB = 1050XT_SEED = 651 52# L2 sensitivity 設定53L2_VALUES = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]54NUM_DIRECTIONS_PER_L2 = 555DIRECTION_SEED = 2026052156 57# 圖片輸出58SAVE_ROW_IMAGES = True59SAVE_FINAL_IMAGE = True60 61ROW_SCALE = 462FINAL_SCALE = 263LABEL_W = 24064COND_W = 0 # synthetic embedding 沒有 trigger 圖，所以不用留 trigger 欄65PAD = 466 67 68# ============================================================69# 1. default params70# ============================================================71 72DEFAULT_IMG_SIZE = 3273DEFAULT_COND_DIM = 12874DEFAULT_TRIGGER_BASE_CH = 3275 76DEFAULT_T = 100077DEFAULT_BETA_1 = 1e-478DEFAULT_BETA_T = 0.0279 80DEFAULT_CH = 12881DEFAULT_CH_MULT = [1, 2, 2, 2]82DEFAULT_ATTN = [1]83DEFAULT_NUM_RES_BLOCKS = 284DEFAULT_DROPOUT = 0.185 86MODEL_KEYS = [87 "global_ema_model",88 "ema_model",89 "global_model",90 "model",91 "net_model",92 "state_dict",93]94 95ENCODER_KEYS = [96 "trigger_encoder",97 "trigger_encoder_model",98 "global_trigger_encoder",99 "ema_trigger_encoder",100 "trigger_encoder_state_dict",101]102 103 104# ============================================================105# 2. basic utils106# ============================================================107 108def ensure_dir(path):109 os.makedirs(path, exist_ok=True)110 111 112def reset_dir(path):113 if os.path.exists(path):114 shutil.rmtree(path)115 os.makedirs(path, exist_ok=True)116 117 118def denorm_to_01(x):119 return ((x + 1.0) / 2.0).clamp(0, 1)120 121 122def tensor_to_pil_01(t, scale=1):123 t = t.detach().cpu().clamp(0, 1)124 arr = (t.permute(1, 2, 0).numpy() * 255).astype(np.uint8)125 img = Image.fromarray(arr)126 if scale != 1:127 img = img.resize((img.width * scale, img.height * scale), Image.NEAREST)128 return img129 130 131def save_tensor_image_01(img_01, path):132 save_image(img_01.detach().cpu().clamp(0, 1), path)133 134 135def read_image_as_tensor(img_path, img_size, device=None):136 tfm = transforms.Compose([137 transforms.Resize((img_size, img_size)),138 transforms.ToTensor(),139 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),140 ])141 img = Image.open(img_path).convert("RGB")142 x = tfm(img)143 if device is not None:144 x = x.to(device)145 return x146 147 148def get_config(ckpt):149 config = ckpt.get("config", {})150 if config is None:151 config = {}152 return config153 154 155def get_from_config(config, key, default):156 value = config.get(key, default)157 if value is None:158 return default159 return value160 161 162def append_csv_row(csv_path, row):163 file_exists = os.path.exists(csv_path)164 165 with open(csv_path, "a", newline="") as f:166 writer = csv.DictWriter(f, fieldnames=list(row.keys()))167 168 if not file_exists:169 writer.writeheader()170 171 writer.writerow(row)172 173 174# ============================================================175# 3. model / encoder / sampler176# ============================================================177 178def infer_model_params_from_ckpt(ckpt):179 config = get_config(ckpt)180 181 mode = get_from_config(config, "mode", "access_control")182 183 img_size = int(get_from_config(config, "img_size", DEFAULT_IMG_SIZE))184 cond_dim = get_from_config(config, "cond_dim", DEFAULT_COND_DIM)185 186 T = int(get_from_config(config, "T", DEFAULT_T))187 beta_1 = float(get_from_config(config, "beta_1", DEFAULT_BETA_1))188 beta_T = float(get_from_config(config, "beta_T", DEFAULT_BETA_T))189 190 ch = int(get_from_config(config, "ch", DEFAULT_CH))191 ch_mult = list(get_from_config(config, "ch_mult", DEFAULT_CH_MULT))192 attn = list(get_from_config(config, "attn", DEFAULT_ATTN))193 num_res_blocks = int(get_from_config(config, "num_res_blocks", DEFAULT_NUM_RES_BLOCKS))194 dropout = float(get_from_config(config, "dropout", DEFAULT_DROPOUT))195 196 if mode == "unconditional":197 model_cond_dim = None198 else:199 model_cond_dim = int(cond_dim)200 201 return {202 "mode": mode,203 "img_size": img_size,204 "cond_dim": None if model_cond_dim is None else int(model_cond_dim),205 "T": T,206 "beta_1": beta_1,207 "beta_T": beta_T,208 "ch": ch,209 "ch_mult": ch_mult,210 "attn": attn,211 "num_res_blocks": num_res_blocks,212 "dropout": dropout,213 "model_cond_dim": model_cond_dim,214 "trigger_base_ch": int(get_from_config(config, "trigger_base_ch", DEFAULT_TRIGGER_BASE_CH)),215 "trigger_grid_size": int(get_from_config(config, "trigger_grid_size", 8)),216 "trigger_hidden_dim": int(get_from_config(config, "trigger_hidden_dim", 256)),217 "trigger_use_coord_feat": bool(get_from_config(config, "trigger_use_coord_feat", True)),218 }219 220 221def load_state_by_keys(module, ckpt, keys, ckpt_path, what):222 for key in keys:223 if key in ckpt:224 module.load_state_dict(ckpt[key], strict=True)225 return key226 227 try:228 module.load_state_dict(ckpt, strict=True)229 return "<ckpt_is_state_dict>"230 except Exception:231 pass232 233 raise KeyError(f"{ckpt_path} 找不到 {what} 權重，已嘗試 keys={keys}")234 235 236def load_sampler_from_ckpt(ckpt_path, device):237 ckpt = torch.load(ckpt_path, map_location="cpu")238 params = infer_model_params_from_ckpt(ckpt)239 240 net = UNet(241 T=params["T"],242 ch=params["ch"],243 ch_mult=params["ch_mult"],244 attn=params["attn"],245 num_res_blocks=params["num_res_blocks"],246 dropout=params["dropout"],247 cond_dim=params["model_cond_dim"],248 ).to(device)249 250 model_key = load_state_by_keys(251 module=net,252 ckpt=ckpt,253 keys=MODEL_KEYS,254 ckpt_path=ckpt_path,255 what="UNet",256 )257 258 net.eval()259 260 sampler = GaussianDiffusionSampler(261 net,262 params["beta_1"],263 params["beta_T"],264 params["T"],265 params["img_size"],266 "epsilon",267 "fixedlarge",268 ).to(device)269 270 sampler.eval()271 272 return {273 "ckpt": ckpt,274 "params": params,275 "sampler": sampler,276 "model_key": model_key,277 }278 279 280def build_trigger_encoder_from_params(params, device):281 if params["cond_dim"] is None:282 raise ValueError("unconditional model has no trigger encoder")283 284 encoder = TriggerEncoder(285 in_ch=3,286 base_ch=params["trigger_base_ch"],287 cond_dim=params["cond_dim"],288 img_size=params["img_size"],289 grid_size=params["trigger_grid_size"],290 hidden_dim=params["trigger_hidden_dim"],291 use_coord_feat=params["trigger_use_coord_feat"],292 ).to(device)293 294 encoder.eval()295 return encoder296 297 298def load_cond_pack(ckpt_path, device):299 pack = load_sampler_from_ckpt(ckpt_path, device)300 301 if pack["params"]["mode"] == "unconditional":302 raise ValueError("COND_CKPT_PATH should not be unconditional.")303 304 encoder = build_trigger_encoder_from_params(pack["params"], device)305 306 encoder_key = load_state_by_keys(307 module=encoder,308 ckpt=pack["ckpt"],309 keys=ENCODER_KEYS,310 ckpt_path=ckpt_path,311 what="TriggerEncoder",312 )313 314 pack["trigger_encoder"] = encoder315 pack["encoder_key"] = encoder_key316 317 return pack318 319 320def encode_trigger(trigger_encoder, trigger_tensor, batch_size, device):321 x = trigger_tensor322 323 if x.dim() == 3:324 x = x.unsqueeze(0)325 326 if x.size(0) == 1:327 x = x.repeat(batch_size, 1, 1, 1)328 329 with torch.no_grad():330 return trigger_encoder(x.to(device))331 332 333def zero_cond(batch_size, cond_dim, device):334 return torch.zeros(batch_size, cond_dim, device=device)335 336 337def build_xT_batch(num_samples, img_size, device):338 xs = []339 seeds = []340 341 for i in range(num_samples):342 seed = XT_SEED + i343 gen = torch.Generator(device=device)344 gen.manual_seed(seed)345 x = torch.randn((1, 3, img_size, img_size), generator=gen, device=device)346 xs.append(x)347 seeds.append(seed)348 349 return torch.cat(xs, dim=0), seeds350 351 352def sample_batch(sampler, x_T, cond, params):353 with torch.no_grad():354 if params["mode"] == "unconditional":355 out = sampler(x_T, 0, params["T"])356 else:357 out = sampler(x_T, 0, params["T"], cond=cond)358 359 return denorm_to_01(out).detach().cpu()360 361 362def get_correct_triggers(cond_ckpt, img_size, device):363 triggers = []364 sources = []365 366 if "client_triggers" in cond_ckpt:367 client_triggers = cond_ckpt["client_triggers"]368 369 for cid in CORRECT_CLIENT_IDS:370 if cid >= len(client_triggers):371 raise IndexError(f"CLIENT_ID={cid}, but ckpt only has {len(client_triggers)} client_triggers")372 373 trig = client_triggers[cid].clone().float().to(device)374 triggers.append(trig)375 sources.append(f"ckpt.client_triggers[{cid}]")376 377 return triggers, sources378 379 for cid in CORRECT_CLIENT_IDS:380 path = os.path.join(FALLBACK_TRIGGER_DIR, f"client_{cid}.png")381 if not os.path.exists(path):382 raise FileNotFoundError(f"Cannot find fallback trigger: {path}")383 384 triggers.append(read_image_as_tensor(path, img_size, device=device))385 sources.append(path)386 387 return triggers, sources388 389 390# ============================================================391# 4. synthetic embedding generation392# ============================================================393 394def build_correct_embeddings(trigger_encoder, correct_triggers, device):395 correct_embs = OrderedDict()396 397 for cid, trig in zip(CORRECT_CLIENT_IDS, correct_triggers):398 emb = encode_trigger(trigger_encoder, trig, 1, device).detach()399 correct_embs[f"correct_client{cid}"] = emb400 401 return correct_embs402 403 404def normalize_direction(direction, eps=1e-12):405 norm = torch.norm(direction, dim=1, keepdim=True).clamp_min(eps)406 return direction / norm407 408 409def make_synthetic_embedding(center_emb, target_l2, seed, device):410 """411 center_emb: [1, cond_dim]412 target_l2: float413 414 回傳：415 synthetic_emb: [1, cond_dim]416 actual_l2: float417 direction_seed: int418 """419 if target_l2 == 0:420 return center_emb.detach().clone(), 0.0, seed421 422 gen = torch.Generator(device=device)423 gen.manual_seed(seed)424 425 direction = torch.randn(center_emb.shape, generator=gen, device=device)426 direction = normalize_direction(direction)427 428 synthetic_emb = center_emb + direction * float(target_l2)429 actual_l2 = torch.norm(synthetic_emb - center_emb).item()430 431 return synthetic_emb.detach(), float(actual_l2), seed432 433 434def repeat_cond(cond_1, batch_size):435 if cond_1.size(0) == batch_size:436 return cond_1437 return cond_1.repeat(batch_size, 1)438 439 440# ============================================================441# 5. metrics442# ============================================================443 444def build_metrics(device):445 metrics = {}446 447 if not TORCHMETRICS_OK:448 print("[WARN] torchmetrics import failed. SSIM / LPIPS will be nan.")449 print(TORCHMETRICS_IMPORT_ERROR)450 return metrics451 452 metrics["psnr"] = PeakSignalNoiseRatio(data_range=1.0).to(device).eval()453 metrics["ssim"] = StructuralSimilarityIndexMeasure(data_range=1.0).to(device).eval()454 metrics["lpips"] = LearnedPerceptualImagePatchSimilarity(net_type="alex").to(device).eval()455 456 return metrics457 458 459def psnr_from_mse(mse):460 if mse <= 1e-12:461 return float("inf")462 return float(10.0 * np.log10(1.0 / mse))463 464 465def compute_pair_metrics(pred_01, ref_01, metrics, device):466 pred = pred_01.unsqueeze(0).to(device).clamp(0, 1)467 ref = ref_01.unsqueeze(0).to(device).clamp(0, 1)468 469 mse = float(F.mse_loss(pred, ref).item())470 471 out = {472 "mse": mse,473 "psnr": psnr_from_mse(mse),474 "ssim": float("nan"),475 "lpips": float("nan"),476 }477 478 if "psnr" in metrics:479 out["psnr"] = float(metrics["psnr"](pred, ref).item())480 481 if "ssim" in metrics:482 out["ssim"] = float(metrics["ssim"](pred, ref).item())483 484 if "lpips" in metrics:485 pred_lp = (pred * 2.0 - 1.0).clamp(-1, 1)486 ref_lp = (ref * 2.0 - 1.0).clamp(-1, 1)487 out["lpips"] = float(metrics["lpips"](pred_lp, ref_lp).item())488 489 return out490 491 492def compute_list_metrics(pred_list, ref_list, metrics, device):493 assert len(pred_list) == len(ref_list)494 495 detail = []496 497 for i, (pred, ref) in enumerate(zip(pred_list, ref_list)):498 m = compute_pair_metrics(pred, ref, metrics, device)499 m["sample_idx"] = i500 detail.append(m)501 502 summary = {}503 504 for key in ["mse", "psnr", "ssim", "lpips"]:505 vals = np.array([x[key] for x in detail], dtype=np.float64)506 summary[f"mean_{key}"] = float(np.nanmean(vals))507 summary[f"std_{key}"] = float(np.nanstd(vals))508 509 return detail, summary510 511 512# ============================================================513# 6. image rows514# ============================================================515 516def make_labeled_row_image(label, imgs_01, scale=4, label_w=240):517 font = ImageFont.load_default()518 519 img_w = imgs_01[0].shape[-1] * scale520 img_h = imgs_01[0].shape[-2] * scale521 n = len(imgs_01)522 523 canvas_w = label_w + n * img_w + (n + 1) * PAD524 canvas_h = img_h + 2 * PAD525 526 canvas = Image.new("RGB", (canvas_w, canvas_h), "white")527 draw = ImageDraw.Draw(canvas)528 529 draw.text((8, canvas_h // 2 - 6), label, fill="black", font=font)530 531 x = label_w + PAD532 y = PAD533 534 for img_01 in imgs_01:535 pil = tensor_to_pil_01(img_01, scale=scale)536 canvas.paste(pil, (x, y))537 x += img_w + PAD538 539 return canvas540 541 542def save_row_image(label, imgs_01, path, scale=4):543 canvas = make_labeled_row_image(544 label=label,545 imgs_01=imgs_01,546 scale=scale,547 label_w=LABEL_W,548 )549 canvas.save(path)550 551 552def save_final_canvas(rows, path):553 row_images = []554 555 for label, imgs_01 in rows:556 row_img = make_labeled_row_image(557 label=label,558 imgs_01=imgs_01,559 scale=FINAL_SCALE,560 label_w=LABEL_W,561 )562 row_images.append(row_img)563 564 w = max(img.width for img in row_images)565 h = sum(img.height for img in row_images)566 567 canvas = Image.new("RGB", (w, h), "white")568 569 y = 0570 for img in row_images:571 canvas.paste(img, (0, y))572 y += img.height573 574 canvas.save(path)575 576 577# ============================================================578# 7. plots579# ============================================================580 581def save_line_plots(summary_csv_path, save_dir):582 try:583 import pandas as pd584 import matplotlib.pyplot as plt585 except Exception as e:586 print("[WARN] cannot plot because pandas/matplotlib import failed:", repr(e))587 return588 589 df = pd.read_csv(summary_csv_path)590 591 sub = df[592 (df["kind"] == "synthetic_embedding") &593 (df["comparison"] == "vs_own_correct")594 ].copy()595 596 if len(sub) == 0:597 return598 599 # plot per client average600 for metric in ["mean_mse", "mean_psnr", "mean_ssim", "mean_lpips"]:601 plt.figure(figsize=(8, 5))602 603 for center_name in sorted(sub["center_name"].unique()):604 cdf = sub[sub["center_name"] == center_name]605 grouped = cdf.groupby("target_l2")[metric].mean().reset_index()606 plt.plot(grouped["target_l2"], grouped[metric], marker="o", label=center_name)607 608 plt.xlabel("target_l2")609 plt.ylabel(metric)610 plt.title(f"L2 sensitivity | {metric} vs own correct")611 plt.grid(True)612 plt.legend()613 614 save_path = os.path.join(save_dir, f"l2_vs_{metric}_per_client.png")615 plt.savefig(save_path, dpi=200, bbox_inches="tight")616 plt.close()617 print("saved plot:", save_path)618 619 # all-client average620 for metric in ["mean_mse", "mean_psnr", "mean_ssim", "mean_lpips"]:621 grouped = sub.groupby("target_l2")[metric].agg(["mean", "std"]).reset_index()622 623 plt.figure(figsize=(8, 5))624 plt.errorbar(625 grouped["target_l2"],626 grouped["mean"],627 yerr=grouped["std"],628 marker="o",629 capsize=3,630 )631 plt.xlabel("target_l2")632 plt.ylabel(metric)633 plt.title(f"L2 sensitivity | all clients avg | {metric}")634 plt.grid(True)635 636 save_path = os.path.join(save_dir, f"l2_vs_{metric}_all_clients_avg.png")637 plt.savefig(save_path, dpi=200, bbox_inches="tight")638 plt.close()639 print("saved plot:", save_path)640 641 642# ============================================================643# 8. main644# ============================================================645 646def main():647 print("=== eval_embedding_sensitivity_l2.py ===")648 print("DEVICE =", DEVICE)649 print("SAVE_DIR =", SAVE_DIR)650 print("BASELINE_CKPT_PATH =", BASELINE_CKPT_PATH)651 print("COND_CKPT_PATH =", COND_CKPT_PATH)652 print("L2_VALUES =", L2_VALUES)653 print("NUM_DIRECTIONS_PER_L2 =", NUM_DIRECTIONS_PER_L2)654 print("NUM_SAMPLES_PER_EMB =", NUM_SAMPLES_PER_EMB)655 656 reset_dir(SAVE_DIR)657 658 rows_dir = os.path.join(SAVE_DIR, "rows")659 ensure_dir(rows_dir)660 661 plots_dir = os.path.join(SAVE_DIR, "plots")662 ensure_dir(plots_dir)663 664 summary_csv = os.path.join(SAVE_DIR, "metrics_summary.csv")665 detail_csv = os.path.join(SAVE_DIR, "metrics_detail.csv")666 report_path = os.path.join(SAVE_DIR, "report.txt")667 668 baseline_pack = load_sampler_from_ckpt(BASELINE_CKPT_PATH, DEVICE)669 cond_pack = load_cond_pack(COND_CKPT_PATH, DEVICE)670 671 baseline_params = baseline_pack["params"]672 cond_params = cond_pack["params"]673 674 img_size = cond_params["img_size"]675 cond_dim = cond_params["cond_dim"]676 677 if cond_dim is None:678 raise ValueError("conditional model cond_dim is None")679 680 if baseline_params["img_size"] != img_size:681 raise ValueError("baseline img_size != cond img_size")682 683 correct_triggers, correct_sources = get_correct_triggers(684 cond_pack["ckpt"],685 img_size,686 DEVICE,687 )688 689 correct_embs = build_correct_embeddings(690 cond_pack["trigger_encoder"],691 correct_triggers,692 DEVICE,693 )694 695 x_T, seed_list = build_xT_batch(NUM_SAMPLES_PER_EMB, img_size, DEVICE)696 697 if baseline_params["mode"] == "unconditional":698 baseline_cond = None699 else:700 baseline_cond = zero_cond(NUM_SAMPLES_PER_EMB, baseline_params["cond_dim"], DEVICE)701 702 cond_none = zero_cond(NUM_SAMPLES_PER_EMB, cond_dim, DEVICE)703 704 metrics = build_metrics(DEVICE)705 706 # --------------------------------------------------------707 # baseline708 # --------------------------------------------------------709 print("sampling baseline...")710 baseline_imgs = sample_batch(711 baseline_pack["sampler"],712 x_T,713 baseline_cond,714 baseline_params,715 )716 baseline_list = [baseline_imgs[i] for i in range(baseline_imgs.size(0))]717 718 # --------------------------------------------------------719 # correct outputs720 # --------------------------------------------------------721 correct_outputs = OrderedDict()722 final_rows = []723 724 final_rows.append(("baseline", baseline_list))725 if SAVE_ROW_IMAGES:726 save_row_image(727 "baseline",728 baseline_list,729 os.path.join(rows_dir, "000_baseline.png"),730 scale=ROW_SCALE,731 )732 733 row_counter = 1734 735 for center_name, center_emb in correct_embs.items():736 print(f"sampling {center_name}...")737 738 cond = repeat_cond(center_emb, NUM_SAMPLES_PER_EMB)739 imgs = sample_batch(740 cond_pack["sampler"],741 x_T,742 cond,743 cond_params,744 )745 746 img_list = [imgs[i] for i in range(imgs.size(0))]747 correct_outputs[center_name] = img_list748 749 final_rows.append((center_name, img_list))750 751 if SAVE_ROW_IMAGES:752 save_row_image(753 center_name,754 img_list,755 os.path.join(rows_dir, f"{row_counter:03d}_{center_name}.png"),756 scale=ROW_SCALE,757 )758 759 row_counter += 1760 761 # --------------------------------------------------------762 # none763 # --------------------------------------------------------764 print("sampling none...")765 none_imgs = sample_batch(766 cond_pack["sampler"],767 x_T,768 cond_none,769 cond_params,770 )771 none_list = [none_imgs[i] for i in range(none_imgs.size(0))]772 773 final_rows.append(("none", none_list))774 775 if SAVE_ROW_IMAGES:776 save_row_image(777 "none",778 none_list,779 os.path.join(rows_dir, f"{row_counter:03d}_none.png"),780 scale=ROW_SCALE,781 )782 783 row_counter += 1784 785 # --------------------------------------------------------786 # report header787 # --------------------------------------------------------788 with open(report_path, "w", encoding="utf-8") as f:789 f.write("embedding-space L2 sensitivity evaluation report\n")790 f.write(f"CKPT_NAME = {CKPT_NAME}\n")791 f.write(f"BASELINE_CKPT_PATH = {BASELINE_CKPT_PATH}\n")792 f.write(f"COND_CKPT_PATH = {COND_CKPT_PATH}\n")793 f.write(f"baseline_mode = {baseline_params['mode']}\n")794 f.write(f"cond_mode = {cond_params['mode']}\n")795 f.write(f"baseline_model_key = {baseline_pack['model_key']}\n")796 f.write(f"cond_model_key = {cond_pack['model_key']}\n")797 f.write(f"cond_encoder_key = {cond_pack['encoder_key']}\n")798 f.write(f"correct_sources = {correct_sources}\n")799 f.write(f"NUM_SAMPLES_PER_EMB = {NUM_SAMPLES_PER_EMB}\n")800 f.write(f"XT_SEED = {XT_SEED}\n")801 f.write(f"seed_list = {seed_list}\n")802 f.write(f"L2_VALUES = {L2_VALUES}\n")803 f.write(f"NUM_DIRECTIONS_PER_L2 = {NUM_DIRECTIONS_PER_L2}\n")804 f.write(f"DIRECTION_SEED = {DIRECTION_SEED}\n")805 f.write(f"SAVE_ROW_IMAGES = {SAVE_ROW_IMAGES}\n")806 807 if not TORCHMETRICS_OK:808 f.write(f"TORCHMETRICS_IMPORT_ERROR = {TORCHMETRICS_IMPORT_ERROR}\n")809 810 f.write("\n")811 812 # --------------------------------------------------------813 # synthetic embeddings814 # --------------------------------------------------------815 summaries_for_report = []816 817 for center_idx, (center_name, center_emb) in enumerate(correct_embs.items()):818 own_correct_list = correct_outputs[center_name]819 820 for target_l2 in L2_VALUES:821 if target_l2 == 0:822 direction_indices = [0]823 else:824 direction_indices = list(range(NUM_DIRECTIONS_PER_L2))825 826 for dir_idx in direction_indices:827 direction_seed = DIRECTION_SEED + center_idx * 100000 + int(target_l2 * 1000) + dir_idx828 829 synth_emb, actual_l2, used_seed = make_synthetic_embedding(830 center_emb=center_emb,831 target_l2=target_l2,832 seed=direction_seed,833 device=DEVICE,834 )835 836 cond = repeat_cond(synth_emb, NUM_SAMPLES_PER_EMB)837 838 label = f"{center_name}_L2{target_l2}_d{dir_idx}"839 print(f"sampling {label}: actual_l2={actual_l2:.6f}")840 841 imgs = sample_batch(842 cond_pack["sampler"],843 x_T,844 cond,845 cond_params,846 )847 848 img_list = [imgs[i] for i in range(imgs.size(0))]849 final_rows.append((label, img_list))850 851 if SAVE_ROW_IMAGES:852 save_row_image(853 label,854 img_list,855 os.path.join(rows_dir, f"{row_counter:03d}_{label}.png"),856 scale=ROW_SCALE,857 )858 859 row_counter += 1860 861 # -------------------------862 # metrics: vs own correct863 # -------------------------864 detail_rows, summary = compute_list_metrics(865 img_list,866 own_correct_list,867 metrics,868 DEVICE,869 )870 871 append_csv_row(summary_csv, {872 "ckpt": CKPT_NAME,873 "kind": "synthetic_embedding",874 "center_name": center_name,875 "target_l2": target_l2,876 "direction_idx": dir_idx,877 "direction_seed": used_seed,878 "actual_l2": actual_l2,879 "comparison": "vs_own_correct",880 "mean_mse": summary["mean_mse"],881 "std_mse": summary["std_mse"],882 "mean_psnr": summary["mean_psnr"],883 "std_psnr": summary["std_psnr"],884 "mean_ssim": summary["mean_ssim"],885 "std_ssim": summary["std_ssim"],886 "mean_lpips": summary["mean_lpips"],887 "std_lpips": summary["std_lpips"],888 })889 890 summaries_for_report.append({891 "center_name": center_name,892 "target_l2": target_l2,893 "comparison": "vs_own_correct",894 **summary,895 })896 897 for row in detail_rows:898 append_csv_row(detail_csv, {899 "ckpt": CKPT_NAME,900 "kind": "synthetic_embedding",901 "center_name": center_name,902 "target_l2": target_l2,903 "direction_idx": dir_idx,904 "direction_seed": used_seed,905 "actual_l2": actual_l2,906 "comparison": "vs_own_correct",907 "sample_idx": row["sample_idx"],908 "mse": row["mse"],909 "psnr": row["psnr"],910 "ssim": row["ssim"],911 "lpips": row["lpips"],912 })913 914 # -------------------------915 # metrics: vs baseline916 # -------------------------917 detail_rows, summary_base = compute_list_metrics(918 img_list,919 baseline_list,920 metrics,921 DEVICE,922 )923 924 append_csv_row(summary_csv, {925 "ckpt": CKPT_NAME,926 "kind": "synthetic_embedding",927 "center_name": center_name,928 "target_l2": target_l2,929 "direction_idx": dir_idx,930 "direction_seed": used_seed,931 "actual_l2": actual_l2,932 "comparison": "vs_baseline",933 "mean_mse": summary_base["mean_mse"],934 "std_mse": summary_base["std_mse"],935 "mean_psnr": summary_base["mean_psnr"],936 "std_psnr": summary_base["std_psnr"],937 "mean_ssim": summary_base["mean_ssim"],938 "std_ssim": summary_base["std_ssim"],939 "mean_lpips": summary_base["mean_lpips"],940 "std_lpips": summary_base["std_lpips"],941 })942 943 for row in detail_rows:944 append_csv_row(detail_csv, {945 "ckpt": CKPT_NAME,946 "kind": "synthetic_embedding",947 "center_name": center_name,948 "target_l2": target_l2,949 "direction_idx": dir_idx,950 "direction_seed": used_seed,951 "actual_l2": actual_l2,952 "comparison": "vs_baseline",953 "sample_idx": row["sample_idx"],954 "mse": row["mse"],955 "psnr": row["psnr"],956 "ssim": row["ssim"],957 "lpips": row["lpips"],958 })959 960 # -------------------------961 # metrics: vs none962 # -------------------------963 detail_rows, summary_none = compute_list_metrics(964 img_list,965 none_list,966 metrics,967 DEVICE,968 )969 970 append_csv_row(summary_csv, {971 "ckpt": CKPT_NAME,972 "kind": "synthetic_embedding",973 "center_name": center_name,974 "target_l2": target_l2,975 "direction_idx": dir_idx,976 "direction_seed": used_seed,977 "actual_l2": actual_l2,978 "comparison": "vs_none",979 "mean_mse": summary_none["mean_mse"],980 "std_mse": summary_none["std_mse"],981 "mean_psnr": summary_none["mean_psnr"],982 "std_psnr": summary_none["std_psnr"],983 "mean_ssim": summary_none["mean_ssim"],984 "std_ssim": summary_none["std_ssim"],985 "mean_lpips": summary_none["mean_lpips"],986 "std_lpips": summary_none["std_lpips"],987 })988 989 for row in detail_rows:990 append_csv_row(detail_csv, {991 "ckpt": CKPT_NAME,992 "kind": "synthetic_embedding",993 "center_name": center_name,994 "target_l2": target_l2,995 "direction_idx": dir_idx,996 "direction_seed": used_seed,997 "actual_l2": actual_l2,998 "comparison": "vs_none",999 "sample_idx": row["sample_idx"],1000 "mse": row["mse"],1001 "psnr": row["psnr"],1002 "ssim": row["ssim"],1003 "lpips": row["lpips"],1004 })1005 1006 # --------------------------------------------------------1007 # final image1008 # --------------------------------------------------------1009 if SAVE_FINAL_IMAGE:1010 save_final_canvas(1011 final_rows,1012 os.path.join(SAVE_DIR, "final_embedding_sensitivity_l2.png"),1013 )1014 1015 # --------------------------------------------------------1016 # report summary1017 # --------------------------------------------------------1018 with open(report_path, "a", encoding="utf-8") as f:1019 f.write("\nSUMMARY, comparison = vs_own_correct\n")1020 1021 for target_l2 in L2_VALUES:1022 rows = [1023 s for s in summaries_for_report1024 if s["target_l2"] == target_l2 and s["comparison"] == "vs_own_correct"1025 ]1026 1027 if len(rows) == 0:1028 continue1029 1030 f.write(f"[L2={target_l2}]\n")1031 for key in ["mean_mse", "mean_psnr", "mean_ssim", "mean_lpips"]:1032 vals = np.array([r[key] for r in rows], dtype=np.float64)1033 f.write(f" avg {key} = {np.nanmean(vals):.6f} ± {np.nanstd(vals):.6f}\n")1034 1035 save_line_plots(summary_csv, plots_dir)1036 1037 print("\ndone")1038 print("SAVE_DIR:", SAVE_DIR)1039 print("rows dir:", rows_dir)1040 print("summary csv:", summary_csv)1041 print("detail csv:", detail_csv)1042 print("report:", report_path)1043 print("final image:", os.path.join(SAVE_DIR, "final_embedding_sensitivity_l2.png"))1044 1045 1046if __name__ == "__main__":1047 main()