validate.py

import argparse
import os

import torch
from tqdm import tqdm
import matplotlib.pyplot as plt

from models.kan_resdiff_sb import Unet
from checkpoint_utils import load_model_state

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

from datetime import datetime
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
    import numpy as np
    import torch
    import os

    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 IoUIoU(target, predicted_mask):
    """
    Args:
        target: (torch.tensor (batchxCxHxW)) Binary Target Segmentation from training set
        predicted_mask: (torch.tensor (batchxCxHxW)) Predicted Segmentation Mask

    Returns:
        IoU: (Float) Average IoUs over Batch
    """
    target = target.detach()
    predicted_mask = predicted_mask.detach()
    smooth = 1e-8
    true_p = (torch.logical_and(target == 1, predicted_mask == 1)).sum()
    false_p = (torch.logical_and(target == 0, predicted_mask == 1)).sum()
    false_n = (torch.logical_and(target == 1, predicted_mask == 0)).sum()
    sample_IoU = (smooth+float(true_p))/(float(true_p) +
                                         float(false_p)+float(false_n)+smooth)

    return sample_IoU
def Dice(target, pred):
    smooth = 1e-8
    tp = (target * pred).sum(dim=[1,2,3])  # (N,)
    denom = target.sum(dim=[1,2,3]) + pred.sum(dim=[1,2,3])
    return (2 * tp + smooth) / (denom + 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(Pred, Masks):
    lcm = np.lcm(len(Pred), len(Masks))
    len1 = len(Pred)
    len2 = len(Masks)
    for i in range((lcm // len1) - 1):
        for j in range(len1):
            Pred.append(Pred[j])
    for i in range((lcm // len2) - 1):
        for j in range(len2):
            Masks.append(Masks[j])
    cost_matrix = np.zeros((lcm, lcm))
    for i in range(lcm):
        for j in range(lcm):
            cost_matrix[i][j] = 1 - IoUIoU(Pred[i], Masks[j])
    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    HM_IoU = np.mean([(1 - cost_matrix[i][j]) for i, j in zip(row_ind, col_ind)])
    return HM_IoU

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 HM_IoU_batch(Preds, Masks):
    """
    Preds: list of length num_pred_sets
           each tensor shape (BS, num_pred_masks, H, W)

    Masks: list of length num_gt_sets
           each tensor shape (BS, num_gt_masks, H, W)

    Return:
        HM-IoU for each sample -> shape (BS,)
    """

    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 = pred_ext[i]
                gt_j   = gt_ext[j]

                pred_i = torch.as_tensor(pred_i, dtype=torch.float32)
                gt_j   = torch.as_tensor(gt_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)          # (T,)
    alphas = 1.0 - betas                                                   # (T,)
    alphas_bar = torch.cumprod(alphas, dim=0)                               # (T,)
    sqrt_ab = torch.sqrt(alphas_bar)                                        # (T,)
    sqrt_1mab = torch.sqrt(1.0 - alphas_bar)                                # (T,)
    return betas, alphas, alphas_bar, sqrt_ab, sqrt_1mab
@torch.no_grad()
def ddim_sampler(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 sampling for epsilon-pred model.
    eta=0 -> deterministic DDIM (recommended for validation speed/stability)
    """
    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  # this is x_0-ish (after final step)


@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):
    """
    Path-defined RSB sampling.

    DDPM time indices are used only to choose model timesteps. The reverse state
    update is x_next = ba(t_next) * x0 + bb(t_next) * eps, so learned paths
    rather than alpha_bar define the final trajectory.
    """
    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 = (t_int.float() / (T - 1)).repeat(x.shape[0])
        t_next_norm = (t_next.float() / (T - 1)).repeat(x.shape[0])
        eps = model(x, t, img_cond, tr_cond).float()

        a_b = paths["ba"](eps, t, img_cond).view(-1, 1, 1, 1)
        b_b = paths["bb"](eps, t, img_cond).view(-1, 1, 1, 1)
        a_b_next = paths["ba"](eps, t_next_norm, img_cond).view(-1, 1, 1, 1)
        b_b_next = paths["bb"](eps, t_next_norm, img_cond).view(-1, 1, 1, 1)
        x0 = (x - b_b * eps) / a_b.clamp(min=1e-8)

        x = a_b_next * x0 + b_b_next * eps

    return x


def visualize_one_sample_panel(
    image,              # (1,H,W)
    gts,                # list of 4 tensors (1,H,W)
    preds,              # list of S tensors (1,H,W) bool or 0/1
    dice_s_k,           # (S,4) tensor for THIS sample
    save_path,
    sample_idx,
):
    """
    保存一个 2x5 panel:
    Row 1: Image | GT1 | GT2 | GT3 | GT4
    Row 2: Other | P(best for GT1) | P(best for GT2) | P(best for GT3) | P(best for GT4)
    """

    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    best_s_idx = dice_s_k.argmax(dim=0)  # (4,)

    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 validate(tr_dataset, model, val_loader, device, sampling_times=40, num_samples=32,
             print_full=True, visual_save=False, visual_root="/root/outputs/Diffusion_KAN1_PUP_SB_m1s1/vis",
             tr_enabled=True, paths=None, sampling_mode="standard"):
    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_maxdice = []
    all_mdm = []
    total_GED = 0.0
    total_IoU = 0.0
    count_samples = 0

    T = 1000
    H = 128
    W = 128

    pbar = tqdm(val_loader, desc="Validating", 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)]  # each: (BS,1,H,W)

        if tr_enabled:
            TR_list = [tr_dataset.get_by_id(cid) for cid in case_ids]  # each (1,H,W)
            TR = torch.stack(TR_list, dim=0).to(device, non_blocking=True)  # (B,1,H,W)
        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(
                    model=model,
                    img_cond=images,
                    tr_cond=TR,
                    device=device,
                    T=T,
                    sampling_steps=sampling_times,
                    shape=(BS, 1, H, W),
                    eta=0.0
                )
            pred = x0.ge(0.5)
            Preds.append(pred.ge(0.5))  # list of (BS,1,H,W) bool/float

        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]))  # (BS,)
            dice_s_k.append(torch.stack(per_rater, dim=0))        # (4,BS)
        dice_s_k = torch.stack(dice_s_k, dim=0)                   # (S,4,BS)

        if visual_save:
            save_dir_all = os.path.join(visual_root, "all_preds")
            os.makedirs(save_dir_all, exist_ok=True)

            for s, p in enumerate(Preds):
                for b in range(BS):
                    save_path = os.path.join(
                        save_dir_all, f"sample{count_samples-BS+b:05d}_pred{s:02d}.png"
                    )
                    plt.imsave(save_path, p[b,0].cpu().numpy(), cmap="gray")

            
            save_dir_panel = os.path.join(visual_root, "panels")
            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(),  # (S,4)
                    save_path=panel_path,
                    sample_idx=count_samples-BS+b
                )

        best_per_rater = dice_s_k.max(dim=0)[0]   # max over S -> (4,BS)
        mdm = best_per_rater.mean(dim=0)          # mean over 4 raters -> (BS,)
        all_mdm.append(mdm.cpu())

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

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

        pbar.set_postfix({
            "MDM":     f"{mdm.mean().item():.4f}",
            "GED":     f"{GED_sum/BS:.4f}",
            "HM-IoU":  f"{IoU_sum/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("\n===== FINAL VALIDATION RESULTS =====")
        print(f"MDM         (per-sample): {final_MDM:.4f}")
        print(f"GED         (per-sample): {final_GED:.4f}")
        print(f"HM-IoU      (per-sample): {final_IoU:.4f}")

    return final_MDM, final_GED, final_IoU

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("--num_epochs", type=int, default=600)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--output_path", type=str, default="/root/outputs")
    parser.add_argument("--resume_filepath", type=str, default=None)
    parser.add_argument("--seed", type=int, default=4242)
    parser.add_argument("--num_samples", type=int, default=32)
    parser.add_argument("--sampling_times", type=int, default=100)

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

    
    parser.add_argument(
        "--visual",
        action="store_true",
        help="enable visualization (default: False)"
    )
    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",        # 或 "cuda"（显存够就放 cuda 更快）
        return_id=False,
        strict=True
    )

    model = Unet(
        channels=1,
        dim_mults=(1, 2, 4),
        dim=args.image_size,
        resnet_block_groups=1,
    ).to(device)
    boost = True
    if boost:
        model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
    model._stage(2)
    if args.resume_filepath is not None:
        checkpoint = torch.load(args.resume_filepath)

        load_model_state(model, checkpoint)

        resume_epoch = checkpoint["epoch"]
        print(f"Testing on main_model FilePath \"{args.resume_filepath}\", Epoch {resume_epoch}:")
    else:
        resume_epoch = 0
        print("Error: main_model file not found")

    _MDM, _GED, _IoU = validate(tr_dataset_test, model, test_loader, device,
                                 sampling_times=args.sampling_times, num_samples=args.num_samples,
                                visual_save=(True if args.visual else False),
                                visual_root=os.path.join(args.output_path, args.task_name))

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

    save_txt = make_timestamp_txt_file(save_txt, args.task_name)

    with open(save_txt, "w") as f:
        f.write(f"Resume File: {args.resume_filepath}\n")
        f.write(f"Parameters: 采样数量:{args.num_samples}, 采样步长:{args.sampling_times}\n")

        f.write(f"MaxDice: {_MDM:.6f}\n")
        f.write(f"GED: {_GED:.6f}\n")
        f.write(f"HM-IoU: {_IoU:.6f}\n")

    print(f"Results saved to {save_txt}")