test_rsb.py

"""
RSB (Residual Schrödinger Bridge) 测试脚本。

支持两种采样模式:
  - standard: 标准 DDIM（使用固定 DDPM schedule）
  - rsb:      最终方法；每一步使用学习到的反向路径系数 ba, bb 更新状态

用法:
  python test_rsb.py \
      --resume_filepath /path/to/Best_xxx.pth \
      --tr_filepath_test /path/to/tr_cache.h5 \
      --num_samples 32 --sampling_times 40 \
      --sampling_mode rsb

对比测试:
  python test_rsb.py \
      --resume_filepath /path/to/Best_xxx.pth \
      --tr_filepath_test /path/to/tr_cache.h5 \
      --compare  # 同时运行 standard 和 rsb 模式并对比
"""

import argparse
import os

import torch
import torch.nn.functional as F
from tqdm import tqdm
import matplotlib.pyplot as plt
import numpy as np

from models.kan_resdiff_sb import Unet
from models.path_sb.diffusion import PathUNet
from models.path_sb.rsb import bridge_loss, rsb_diffuse

from dataloaders import *
from metrics import *
from scipy.optimize import linear_sum_assignment
from models.pretrain.TRDataset import TRCacheByIdDataset
from datetime import datetime
from checkpoint_utils import load_model_state

torch.set_float32_matmul_precision('high')

def make_timestamp_txt_file(save_path, task_name, time_format="%Y-%m-%d_%H-%M-%S"):
    timestamp = datetime.now().strftime(time_format)
    final_path = os.path.join(save_path, f"{task_name}_{timestamp}.txt")
    return final_path

def set_seed(seed):
    import random
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ["PYTHONHASHSEED"] = str(seed)
    print(f"[INFO] Global seed set to {seed}")

def Dice_batch(pred, target):
    smooth = 1e-8
    tp = (pred * target).sum(dim=[1, 2, 3])
    denom = pred.sum(dim=[1, 2, 3]) + target.sum(dim=[1, 2, 3])
    return (2 * tp + smooth) / (denom + smooth)

def IoU_batch(pred, target):
    smooth = 1e-8
    tp = (pred * target).sum(dim=[1, 2, 3])
    fp = ((1 - target) * pred).sum(dim=[1, 2, 3])
    fn = (target * (1 - pred)).sum(dim=[1, 2, 3])
    return (tp + smooth) / (tp + fp + fn + smooth)

def IoU(target, pred):
    smooth = 1e-8
    tp = (target * pred).sum(dim=[1, 2, 3])
    fp = ((1 - target) * pred).sum(dim=[1, 2, 3])
    fn = (target * (1 - pred)).sum(dim=[1, 2, 3])
    return (tp + smooth) / (tp + fp + fn + smooth)

def HM_IoU_batch(Preds, Masks):
    BS = Preds[0].shape[0]
    num_pred_sets = len(Preds)
    num_gt_sets = len(Masks)
    hm_results = []
    for b in range(BS):
        pred_list = []
        for i in range(num_pred_sets):
            masks_i = Preds[i][b].cpu()
            for m in masks_i:
                pred_list.append(m)
        gt_list = []
        for j in range(num_gt_sets):
            masks_j = Masks[j][b].cpu()
            for m in masks_j:
                gt_list.append(m)
        p = len(pred_list)
        g = len(gt_list)
        k = np.lcm(p, g)
        pred_ext = pred_list * (k // p)
        gt_ext = gt_list * (k // g)
        cost = np.zeros((k, k))
        for i in range(k):
            for j in range(k):
                pred_i = torch.as_tensor(pred_ext[i], dtype=torch.float32)
                gt_j = torch.as_tensor(gt_ext[j], dtype=torch.float32)
                pred_i = pred_i.unsqueeze(0).unsqueeze(0)
                gt_j = gt_j.unsqueeze(0).unsqueeze(0)
                val = IoU(pred_i, gt_j)
                cost[i, j] = 1 - (val.item() if torch.is_tensor(val) else float(val))
        row_ind, col_ind = linear_sum_assignment(cost)
        hm_iou = np.mean([1 - cost[i, j] for i, j in zip(row_ind, col_ind)])
        hm_results.append(hm_iou)
    return torch.tensor(hm_results, dtype=torch.float32)

def make_ddpm_schedule(T=1000, beta_start=1e-4, beta_end=2e-2, device="cuda"):
    betas = torch.linspace(beta_start, beta_end, T, device=device)
    alphas = 1.0 - betas
    alphas_bar = torch.cumprod(alphas, dim=0)
    sqrt_ab = torch.sqrt(alphas_bar)
    sqrt_1mab = torch.sqrt(1.0 - alphas_bar)
    return betas, alphas, alphas_bar, sqrt_ab, sqrt_1mab

def visualize_one_sample_panel(
    image, gts, preds, dice_s_k, save_path, sample_idx,
):
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    best_s_idx = dice_s_k.argmax(dim=0)
    best_preds = [preds[s] for s in best_s_idx.tolist()]
    all_idx = set(range(len(preds)))
    used_idx = set(best_s_idx.tolist())
    other_idx = list(all_idx - used_idx)
    if len(other_idx) > 0:
        other_pred = preds[other_idx[0]]
        other_title = f"Pred #{other_idx[0]}"
    else:
        other_pred = preds[best_s_idx[0]]
        other_title = f"Pred #{best_s_idx[0]} (dup)"
    fig, axes = plt.subplots(2, 5, figsize=(15, 6))
    axes[0, 0].imshow(image.squeeze(), cmap="gray")
    axes[0, 0].set_title("Image")
    for k in range(4):
        axes[0, k + 1].imshow(gts[k].squeeze(), cmap="gray")
        axes[0, k + 1].set_title(f"GT {k+1}")
    axes[1, 0].imshow(other_pred.squeeze(), cmap="gray")
    axes[1, 0].set_title(other_title)
    for k in range(4):
        s_idx = best_s_idx[k].item()
        axes[1, k + 1].imshow(best_preds[k].squeeze(), cmap="gray")
        axes[1, k + 1].set_title(f"Pred #{s_idx} (GT{k+1})")
    for ax in axes.flatten():
        ax.axis("off")
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()

@torch.no_grad()
def ddim_sampler_rsb(
    model, img_cond, tr_cond, paths, device,
    T=1000, sampling_steps=40, shape=(1, 1, 128, 128),
    beta_start=1e-4, beta_end=2e-2, eta=0.0,
):
    """
    纯 RSB/path 采样。
    TRNet 内部处理 alpha(t) 加权；状态更新只使用学习到的 ba/bb。

    Args:
        model:   主 U-Net（TRNet 内部包含 alpha head）
        img_cond: (B, 1, H, W) 输入图像
        tr_cond:  (B, 1, H, W) TR 先验掩码 (可为 None)
        paths:    dict with 'ba', 'bb' (学习到的反向路径模型)
        device:   设备
    Returns:
        x: 去噪后的分割掩码 (B, 1, H, W)
    """
    if eta != 0.0:
        raise ValueError("RSB path sampling currently supports eta=0.0 only.")

    times = torch.linspace(T - 1, 0, steps=sampling_steps, device=device).long()
    times_next = torch.cat([times[1:], torch.tensor([0], device=device, dtype=torch.long)])

    x = torch.randn(shape, device=device)

    for t_int, t_next in zip(times, times_next):
        t_norm = (t_int.float() / (T - 1)).repeat(x.shape[0])
        t_next_norm = (t_next.float() / (T - 1)).repeat(x.shape[0])

        eps_pred = model(x, t_norm, img_cond, tr_cond).float()

        a_b = paths['ba'](eps_pred, t_norm, img_cond).view(-1, 1, 1, 1)
        b_b = paths['bb'](eps_pred, t_norm, img_cond).view(-1, 1, 1, 1)
        a_b_next = paths['ba'](eps_pred, t_next_norm, img_cond).view(-1, 1, 1, 1)
        b_b_next = paths['bb'](eps_pred, t_next_norm, img_cond).view(-1, 1, 1, 1)

        x0 = (x - b_b * eps_pred) / (a_b.clamp(min=1e-8))

        eps = eps_pred
        x = a_b_next * x0 + b_b_next * eps

    return x

@torch.no_grad()
def ddim_sampler_standard(
    model, img_cond, tr_cond, device,
    T=1000, sampling_steps=40, shape=(1, 1, 128, 128),
    beta_start=1e-4, beta_end=2e-2, eta=0.0,
):
    """
    标准 DDIM/DDPM schedule 采样。
    """
    betas, alphas, alphas_bar, _, _ = make_ddpm_schedule(
        T=T, beta_start=beta_start, beta_end=beta_end, device=device
    )

    times = torch.linspace(T - 1, 0, steps=sampling_steps, device=device).long()
    times_next = torch.cat([times[1:], torch.tensor([0], device=device, dtype=torch.long)])

    x = torch.randn(shape, device=device)

    for t_int, t_next in zip(times, times_next):
        t = (t_int.float() / (T - 1)).repeat(x.shape[0])
        eps = model(x, t, img_cond, tr_cond).float()

        ab_t = alphas_bar[t_int].view(1, 1, 1, 1)
        ab_next = alphas_bar[t_next].view(1, 1, 1, 1)

        x0 = (x - torch.sqrt(1 - ab_t) * eps) / (torch.sqrt(ab_t) + 1e-8)

        if eta == 0.0:
            x = torch.sqrt(ab_next) * x0 + torch.sqrt(1 - ab_next) * eps
        else:
            a_t = alphas[t_int].view(1, 1, 1, 1)
            sigma = eta * torch.sqrt((1 - ab_next) / (1 - ab_t) * (1 - a_t))
            noise = torch.randn_like(x)
            x = torch.sqrt(ab_next) * x0 + torch.sqrt(1 - ab_next - sigma**2) * eps + sigma * noise

    return x


@torch.no_grad()
def validate_rsb(
    tr_dataset, model, val_loader, device,
    sampling_times=40, num_samples=32,
    print_full=True, tr_enabled=True,
    visual_save=False, visual_root="/tmp/rsb_vis",
    paths=None, sampling_mode="standard",
):
    """
    推理验证函数。

    Args:
        sampling_mode: "standard" | "rsb"
            - "standard": 标准 DDIM
            - "rsb": 最终方法；每一步使用学习到的 ba, bb 系数更新状态
    """
    model.eval()
    if sampling_mode not in {"standard", "rsb"}:
        raise ValueError(f"Unknown sampling_mode: {sampling_mode}")
    if sampling_mode == "rsb" and paths is None:
        raise ValueError(f"{sampling_mode} sampling requires learned paths, but paths=None was passed.")

    if paths is not None:
        for k in ["fa", "fb", "ba", "bb"]:
            paths[k].eval()

    all_mdm = []
    total_GED = 0.0
    total_IoU = 0.0
    count_samples = 0

    T = 1000
    H = 128
    W = 128

    pbar = tqdm(val_loader, desc=f"Validating ({sampling_mode})", ncols=100, leave=False)

    for images, masks, _, case_ids in pbar:
        BS = images.size(0)
        count_samples += BS

        images = images.to(device)
        Masks = [masks[i].to(device) for i in range(4)]

        if tr_enabled:
            TR_list = [tr_dataset.get_by_id(cid) for cid in case_ids]
            TR = torch.stack(TR_list, dim=0).to(device, non_blocking=True)
        else:
            TR = None

        Preds = []
        for _ in range(num_samples):
            if sampling_mode == "rsb" and paths is not None:
                x0 = ddim_sampler_rsb(
                    model=model, img_cond=images, tr_cond=TR,
                    paths=paths, device=device,
                    T=T, sampling_steps=sampling_times,
                    shape=(BS, 1, H, W), eta=0.0,
                )
            else:
                x0 = ddim_sampler_standard(
                    model=model, img_cond=images, tr_cond=TR,
                    device=device, T=T, sampling_steps=sampling_times,
                    shape=(BS, 1, H, W), eta=0.0,
                )
            Preds.append(x0.ge(0.5))

        dice_s_k = []
        for s in range(num_samples):
            per_rater = []
            for k in range(4):
                per_rater.append(Dice_batch(Preds[s], Masks[k]))
            dice_s_k.append(torch.stack(per_rater, dim=0))
        dice_s_k = torch.stack(dice_s_k, dim=0)  # (S, 4, BS)

        if visual_save:
            save_dir_panel = os.path.join(visual_root, "panels", sampling_mode)
            os.makedirs(save_dir_panel, exist_ok=True)
            for b in range(BS):
                panel_path = os.path.join(
                    save_dir_panel,
                    f"sample{count_samples - BS + b:05d}_panel.png"
                )
                visualize_one_sample_panel(
                    image=images[b].detach().cpu(),
                    gts=[Masks[k][b].detach().cpu() for k in range(4)],
                    preds=[Preds[s][b].detach().cpu() for s in range(num_samples)],
                    dice_s_k=dice_s_k[:, :, b].detach().cpu(),
                    save_path=panel_path,
                    sample_idx=count_samples - BS + b,
                )

        best_per_rater = dice_s_k.max(dim=0)[0]
        mdm = best_per_rater.mean(dim=0)
        all_mdm.append(mdm.cpu())

        GED_now, _ = ged(Masks, Preds)
        total_GED += GED_now.sum().item()

        IoU_now = HM_IoU_batch(Preds, Masks)
        total_IoU += IoU_now.sum().item()

        pbar.set_postfix({
            "MDM": f"{mdm.mean().item():.4f}",
            "GED": f"{GED_now.sum().item() / BS:.4f}",
            "HM-IoU": f"{IoU_now.sum().item() / BS:.4f}",
        })

    final_MDM = torch.cat(all_mdm).mean().item()
    final_GED = total_GED / count_samples
    final_IoU = total_IoU / count_samples

    if print_full:
        print(f"\n===== [{sampling_mode.upper()}] FINAL RESULTS =====")
        print(f"MDM:    {final_MDM:.4f}")
        print(f"GED:    {final_GED:.4f}")
        print(f"HM-IoU: {final_IoU:.4f}")

    return final_MDM, final_GED, final_IoU

@torch.no_grad()
def visualize_alpha_curve(model, device, save_path):
    """
    绘制 TRNet 内部 alpha(t) 随 t 变化的曲线。
    alpha 由 tr_encoder.alpha_head 输出。
    """
    t_vals = torch.linspace(0, 1, 100, device=device)
    alpha_vals = []
    dummy_x = torch.randn(1, 1, 128, 128, device=device)
    dummy_img = torch.randn(1, 1, 128, 128, device=device)
    dummy_tr = torch.randn(1, 1, 128, 128, device=device)
    for t in t_vals:
        _ = model(dummy_x, t.unsqueeze(0), dummy_img, dummy_tr)
        alpha_vals.append(model.alpha.item() if model.alpha is not None else 0)
    alpha_vals = np.array(alpha_vals)
    t_vals = t_vals.cpu().numpy()

    plt.figure(figsize=(8, 5))
    plt.plot(t_vals, alpha_vals, 'b-', linewidth=2, label='alpha(t) from TRNet')
    plt.plot(t_vals, 1 - t_vals, 'r--', linewidth=1, label='target: 1-t')
    plt.xlabel('Normalized time t')
    plt.ylabel('alpha(t)')
    plt.title('Alpha(t) — TRNet internal stage-aware weight')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    print(f"[INFO] Alpha curve saved to {save_path}")

@torch.no_grad()
def visualize_path_coefficients(paths, device, save_path):
    """
    可视化前向和反向路径系数随 t 的变化。
    """
    t_vals = torch.linspace(0, 1, 50, device=device)
    eps_dummy = torch.randn(1, 1, 128, 128, device=device)
    img_dummy = torch.randn(1, 1, 128, 128, device=device)

    a_f_vals = torch.cat([paths['fa'](eps_dummy, t.expand(1), img_dummy) for t in t_vals]).cpu().numpy()
    b_f_vals = torch.cat([paths['fb'](eps_dummy, t.expand(1), img_dummy) for t in t_vals]).cpu().numpy()
    a_b_vals = torch.cat([paths['ba'](eps_dummy, t.expand(1), img_dummy) for t in t_vals]).cpu().numpy()
    b_b_vals = torch.cat([paths['bb'](eps_dummy, t.expand(1), img_dummy) for t in t_vals]).cpu().numpy()

    t_np = t_vals.cpu().numpy()

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    ax1.plot(t_np, a_f_vals, 'b-', linewidth=2, label='a_f (forward)')
    ax1.plot(t_np, a_b_vals, 'r--', linewidth=2, label='a_b (backward)')
    ax1.set_xlabel('t')
    ax1.set_ylabel('a(t)')
    ax1.set_title('Path coefficient a(t)')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    ax2.plot(t_np, b_f_vals, 'b-', linewidth=2, label='b_f (forward)')
    ax2.plot(t_np, b_b_vals, 'r--', linewidth=2, label='b_b (backward)')
    ax2.set_xlabel('t')
    ax2.set_ylabel('b(t)')
    ax2.set_title('Path coefficient b(t)')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.suptitle('RSB Path Coefficients (smaller gap = better bridge)')
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    print(f"[INFO] Path coefficients saved to {save_path}")

    bridge_a = np.mean((a_f_vals - a_b_vals) ** 2)
    bridge_b = np.mean((b_f_vals - b_b_vals) ** 2)
    print(f"[INFO] Bridge gap: a={bridge_a:.6f}, b={bridge_b:.6f}, total={bridge_a + bridge_b:.6f}")

def init_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image_size", type=int, default=128)
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--output_path", type=str, default="/root/outputs")
    parser.add_argument("--resume_filepath", type=str, default=None,
                        help="包含 model_state_dict 和 paths 的 checkpoint 路径")
    parser.add_argument("--seed", type=int, default=4242)
    parser.add_argument("--num_samples", type=int, default=32,
                        help="每张图采样的分割数量")
    parser.add_argument("--sampling_times", type=int, default=40,
                        help="反向采样步数")

    parser.add_argument("--tr_filepath_test", type=str, default="")
    parser.add_argument("--task_name", type=str, default="K1PUPD_SB_RSB_m2s1")

    parser.add_argument("--sampling_mode", type=str, default="rsb",
                        choices=["standard", "rsb"],
                        help="standard=标准DDIM, rsb=path-defined RSB采样")
    parser.add_argument("--compare", action="store_true", default=False,
                        help="同时运行 standard 和 rsb 模式并对比")
    parser.add_argument("--visual", action="store_true",
                        help="保存可视化结果（面板图和曲线图）")

    return parser

if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    parser = init_parser()
    args = parser.parse_args()

    set_seed(args.seed)

    test_loader, _ = get_dataloader_2(
        task="LIDC", split="test", batch_size=args.batch_size,
        shuffle=False, splitratio=[0.8, 0.015, 0.185], randomsplit=True,
    )

    tr_dataset_test = TRCacheByIdDataset(
        args.tr_filepath_test, device="cpu", return_id=False, strict=True,
    )

    model = Unet(
        channels=1, dim_mults=(1, 2, 4), dim=args.image_size,
        resnet_block_groups=1,
    ).to(device)

    model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
    model._stage(2)  # eval

    def path_model_klass():
        return PathUNet(dim=args.image_size).to(device=device)

    paths = {
        "fa": path_model_klass(),
        "fb": path_model_klass(),
        "ba": path_model_klass(),
        "bb": path_model_klass(),
    }

    if args.resume_filepath is not None:
        checkpoint = torch.load(args.resume_filepath, map_location=device)

        load_model_state(model, checkpoint, strict=False)

        ckpt_paths = checkpoint.get("paths")
        if ckpt_paths is not None:
            paths["fa"].load_state_dict(ckpt_paths["fa"])
            paths["fb"].load_state_dict(ckpt_paths["fb"])
            paths["ba"].load_state_dict(ckpt_paths["ba"])
            paths["bb"].load_state_dict(ckpt_paths["bb"])
            print("[INFO] Path models loaded from checkpoint")
        elif args.compare or args.sampling_mode == "rsb":
            raise ValueError("RSB sampling requires checkpoint['paths'], but this checkpoint does not contain paths.")

        resume_epoch = checkpoint.get("epoch", "?")
        print(f"Loaded checkpoint from epoch {resume_epoch}")
    else:
        print("[WARN] No checkpoint specified. Path models are randomly initialized.")

    save_txt = os.path.join(args.output_path, "test_results")
    os.makedirs(save_txt, exist_ok=True)
    vis_root = os.path.join(args.output_path, args.task_name, "vis")

    if args.visual:
        os.makedirs(vis_root, exist_ok=True)
        visualize_alpha_curve(model, device, os.path.join(vis_root, "alpha_curve.png"))
        visualize_path_coefficients(paths, device, os.path.join(vis_root, "path_coefficients.png"))

    if args.compare:
        results = {}
        for mode in ["standard", "rsb"]:
            print(f"\n{'='*60}")
            print(f"Running {mode.upper()} sampling...")
            print(f"{'='*60}")
            _MDM, _GED, _IoU = validate_rsb(
                tr_dataset_test, model, test_loader, device,
                sampling_times=args.sampling_times, num_samples=args.num_samples,
                tr_enabled=True, paths=paths,
                sampling_mode=mode,
                visual_save=args.visual, visual_root=vis_root,
            )
            results[mode] = (_MDM, _GED, _IoU)

        print(f"\n{'='*60}")
        print("COMPARISON: standard vs RSB")
        print(f"{'='*60}")
        print(f"{'Metric':<12} {'Standard':>10} {'RSB':>10} {'Delta':>10}")
        print(f"{'-'*42}")
        for i, metric in enumerate(["MDM", "GED", "HM-IoU"]):
            std_val = results["standard"][i]
            rsb_val = results["rsb"][i]
            delta = rsb_val - std_val
            direction = "better" if (
                (metric == "GED" and delta < 0) or
                (metric != "GED" and delta > 0)
            ) else "worse"
            print(f"{metric:<12} {std_val:>10.4f} {rsb_val:>10.4f} {delta:>+10.4f} ({direction})")

        save_path = make_timestamp_txt_file(save_txt, f"{args.task_name}_compare")
        with open(save_path, "w") as f:
            f.write(f"Checkpoint: {args.resume_filepath}\n")
            f.write(f"Samples: {args.num_samples}, Steps: {args.sampling_times}\n\n")
            f.write(f"{'Metric':<12} {'Standard':>10} {'RSB':>10} {'Delta':>10}\n")
            f.write(f"{'-'*42}\n")
            for i, metric in enumerate(["MDM", "GED", "HM-IoU"]):
                f.write(f"{metric:<12} {results['standard'][i]:>10.4f} {results['rsb'][i]:>10.4f} {results['rsb'][i] - results['standard'][i]:>+10.4f}\n")

    else:
        print(f"\n{'='*60}")
        print(f"Running {args.sampling_mode.upper()} sampling...")
        print(f"{'='*60}")
        _MDM, _GED, _IoU = validate_rsb(
            tr_dataset_test, model, test_loader, device,
            sampling_times=args.sampling_times, num_samples=args.num_samples,
            tr_enabled=True, paths=paths,
            sampling_mode=args.sampling_mode,
            visual_save=args.visual, visual_root=vis_root,
        )

        save_path = make_timestamp_txt_file(save_txt, args.task_name)
        with open(save_path, "w") as f:
            f.write(f"Checkpoint: {args.resume_filepath}\n")
            f.write(f"Samples: {args.num_samples}, Steps: {args.sampling_times}\n")
            f.write(f"Sampling mode: {args.sampling_mode}\n\n")
            f.write(f"MDM:    {_MDM:.6f}\n")
            f.write(f"GED:    {_GED:.6f}\n")
            f.write(f"HM-IoU: {_IoU:.6f}\n")

    print(f"\nResults saved to {save_txt}")