samihs_src/utils/evaluation.py

# -*- coding: utf-8 -*-
# this file is utilized to evaluate the models from mode: 2D-slice level
from torch.autograd import Variable
from torch.utils.data import DataLoader
import os
import numpy as np
import torch
import torch.nn.functional as F
import utils.metrics as metrics
from hausdorff import hausdorff_distance
from utils.generate_prompts import get_click_prompt, get_click_prompt_eval
import time
import pandas as pd
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as Func
from tqdm import tqdm
import cv2
import nibabel as nib                       # ★ 新

# ------------------------------- 可视化辅助函数 ------------------------------- #
def _overlay(gray, pred_mask, gt_mask=None, alpha=0.4):
    gray_uint8 = (gray * 255).astype(np.uint8)
    gray3      = cv2.cvtColor(gray_uint8, cv2.COLOR_GRAY2BGR)

    over = gray3.copy()
    if gt_mask is not None:                       # 绿色通道
        over[gt_mask > 0, 1] = 255
    if pred_mask is not None:                     # 红色通道
        over[pred_mask > 0, 2] = 255

    return cv2.addWeighted(gray3, 1 - alpha, over, alpha, 0)

def _color_overlay(gray, mask, color, alpha=0.8):
    g3 = cv2.cvtColor((gray * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR)
    color_layer = np.zeros_like(g3)
    color_layer[mask > 0] = color
    return cv2.addWeighted(color_layer, alpha, g3, 1 - alpha, 0)

def _two_panel(img_left, img_right):
    return cv2.hconcat([img_left, img_right])

def _hstack_many(panels):
    """横向拼接多张同尺寸 BGR 图"""
    assert len(panels) >= 2
    h, w = panels[0].shape[:2]
    for p in panels:
        assert p.shape[:2] == (h, w), f"shape mismatch: {p.shape} vs {(h,w)}"
    return cv2.hconcat(panels)
# ----------------------------------------------------------------------------- #

def eval_mask_slice2(valloader, model, criterion, opt, args):
    if getattr(args, 'vis_dir', None):
        os.makedirs(args.vis_dir, exist_ok=True)

    model.eval()
    val_losses, mean_dice = 0, 0
    max_slice_number = opt.batch_size * (len(valloader) + 1)

    # 初始化指标容器
    fore_dice, hds_fore = [], []
    dices = np.zeros((max_slice_number, opt.classes))
    hds = np.zeros((max_slice_number, opt.classes))
    ious, accs, ses, sps = np.zeros_like(dices), np.zeros_like(dices), np.zeros_like(dices), np.zeros_like(dices)
    eval_number = 0
    sum_time = 0

    with tqdm(total=len(valloader), desc='Validation round', unit='batch', leave=False) as pbar:
        for batch_idx, datapack in enumerate(valloader):
            imgs = Variable(datapack['image'].to(dtype=torch.float32, device=opt.device))
            image_filename = datapack['image_name']

            pt = get_click_prompt_eval(datapack, opt)

            # ============== DEBUG: prompt / bbox 体检 =================
            if isinstance(pt, (tuple, list)):
                _coords, _labels = pt
                if not torch.isfinite(_coords).all():
                    print(f"[DEBUG][VAL] pt coords NaN/Inf at batch={batch_idx}")
                if not torch.isfinite(_labels.float()).all():
                    print(f"[DEBUG][VAL] pt labels NaN/Inf at batch={batch_idx}")
            else:
                print(f"[DEBUG][VAL] pt type unexpected: {type(pt)} at batch={batch_idx}")

            bbox = None  # 你的原逻辑保持不变；若你传 bbox，可在此做 clamp 并打印
            # =========================================================

            with torch.no_grad():
                start_time = time.time()
                pred = model(imgs, pt, bbox=None)
                print(f"pred mean = {pred['masks'].mean().item()}, std = {pred['masks'].std().item()}")
                sum_time += (time.time() - start_time)

            # ============== DEBUG: logits 体检 ========================
            if not isinstance(pred, dict) or 'masks' not in pred:
                print(f"[DEBUG][VAL] pred format unexpected at batch={batch_idx}: keys={list(pred.keys()) if isinstance(pred, dict) else type(pred)}")
            else:
                _logits = pred['masks']
                if not torch.isfinite(_logits).all():
                    print(f"[DEBUG][VAL] logits NaN/Inf at batch={batch_idx}",
                        " min/max:", float(_logits.min()), float(_logits.max()))
            # =========================================================

            predict = torch.sigmoid(pred['masks'])
            predict = Func.resize(predict, (512, 512), InterpolationMode.BILINEAR)
            predict = predict.detach().cpu().numpy()[:, 0, :, :] > 0.5

            b, h, w = predict.shape

            has_gt = 'low_mask' in datapack and 'label' in datapack
            if has_gt:
                masks = Variable(datapack['low_mask'].to(dtype=torch.float32, device=opt.device))
                label = Variable(datapack['label'].to(dtype=torch.float32, device=opt.device))

                # ============== DEBUG: GT 体检 & 空掩膜比例 =============
                dims = tuple(range(1, masks.ndim))
                per_sample_sum = masks.sum(dim=dims)
                empty_ratio = (per_sample_sum == 0).float().mean().item()
                if empty_ratio > 0:
                    # print(f"[DEBUG][VAL] empty GT ratio={empty_ratio:.3f} at batch={batch_idx}")
                    pass
                if not torch.isfinite(masks).all():
                    print(f"[DEBUG][VAL] masks NaN/Inf at batch={batch_idx}")
                # =======================================================

                gt = label.detach().cpu().numpy()[:, 0, :, :]
                val_loss = criterion(pred, masks)

                # ============== DEBUG: val loss 体检 ====================
                if not torch.isfinite(val_loss):
                    print(f"[DEBUG][VAL] loss NaN/Inf at batch={batch_idx} "
                        f"(loss={val_loss})")
                    # 打印更多上下文
                    print("  imgs finite:", torch.isfinite(imgs).all().item(),
                        " min/max:", float(imgs.min()), float(imgs.max()))
                    print("  logits finite:", torch.isfinite(pred['masks']).all().item(),
                        " min/max:", float(pred['masks'].min()), float(pred['masks'].max()))
                    print("  masks finite:", torch.isfinite(masks).all().item(),
                        " sum per-sample:", per_sample_sum.tolist())
                # =======================================================

                val_losses += val_loss.item()
                pbar.set_postfix(**{'loss (batch)': val_loss.item()})
            else:
                gt = [None] * b


            for j in range(b):
                pred_i = predict[j].astype(np.uint8)
                if gt[j] is not None:
                    gt_i = (gt[j] > 0).astype(np.uint8)

                    # ---- 先计算指标（不依赖是否可视化） ----
                    # both-empty 保护
                    if pred_i.sum() + gt_i.sum() == 0:
                        dice_i = 1.0
                        hd_i   = 0.0
                        iou, acc, se, sp = 1.0, 1.0, 1.0, 1.0
                    else:
                        dice_i = metrics.dice_coefficient(pred_i[None], gt_i[None])
                        iou, acc, se, sp = metrics.sespiou_coefficient2(pred_i[None], gt_i[None], all=False)
                        hd_i   = hausdorff_distance(pred_i, gt_i)

                    # 写入容器
                    dices[eval_number + j, 1] = dice_i
                    ious [eval_number + j, 1] = iou
                    accs [eval_number + j, 1] = acc
                    ses  [eval_number + j, 1] = se
                    sps  [eval_number + j, 1] = sp
                    hds  [eval_number + j, 1] = hd_i
                    if gt_i.sum() > 0:
                        fore_dice.append(dice_i)
                        hds_fore.append(hd_i)

                # ---- 只有这里保留可视化 ----
                if getattr(args, 'vis_dir', None):
                    base = os.path.splitext(image_filename[j])[0]
                    gray = imgs[j, 0].cpu().numpy()
                    gray = (gray - gray.min()) / (gray.ptp() + 1e-6)
                    if gray.shape != (h, w):
                        gray = cv2.resize(gray, (w, h), interpolation=cv2.INTER_LINEAR)

                    gray_bgr  = cv2.cvtColor((gray * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR)
                    panel_pred = _color_overlay(gray, pred_i, color=(0, 0, 255), alpha=0.8)

                    if gt[j] is not None:
                        panel_gt = _color_overlay(gray, (gt[j] > 0).astype(np.uint8), color=(0, 255, 0), alpha=0.8)
                        # 三联图：原图 | GT 叠加 | Pred 叠加
                        combined = _hstack_many([gray_bgr, panel_gt, panel_pred])
                    else:
                        # 无 GT 时仍旧两联图：原图 | Pred 叠加
                        combined = _two_panel(gray_bgr, panel_pred)

                    cv2.imwrite(os.path.join(args.vis_dir, f"{base}_compare.png"), combined[:, :, ::-1])

            eval_number += b
            pbar.update()

    if not has_gt:
        print("⚠️ Warning: No ground-truth masks detected. Only predictions are visualized.")
        return

    # 指标统计
    dices = dices[:eval_number, :]
    hds   = hds[:eval_number, :]
    ious, accs, ses, sps = (ious[:eval_number, :], accs[:eval_number, :],
                            ses[:eval_number, :],  sps[:eval_number, :])
    val_losses /= (batch_idx + 1)

    fore_dice_mean = np.mean(fore_dice)
    print("fore_dice_mean", fore_dice_mean)
    hds_fore_mean  = np.mean(hds_fore)
    dice_mean      = np.mean(dices, axis=0)
    print("dice_mean", dice_mean)
    dices_std      = np.std(dices,  axis=0)
    hd_mean        = np.mean(hds,   axis=0)
    hd_std         = np.std(hds,    axis=0)

    mean_dice  = np.mean(dice_mean[1:])
    mean_hdis  = np.mean(hd_mean[1:])
    print("test speed", eval_number / sum_time)

    if opt.mode == "train":
        return dices, fore_dice_mean, hds_fore_mean, val_losses
    else:
        iou_mean, iou_std = np.mean(ious, axis=0), np.std(ious, axis=0)
        acc_mean, acc_std = np.mean(accs, axis=0), np.std(accs, axis=0)
        se_mean, se_std   = np.mean(ses,  axis=0), np.std(ses,  axis=0)
        sp_mean, sp_std   = np.mean(sps,  axis=0), np.std(sps,  axis=0)

        return fore_dice_mean, hds_fore_mean, iou_mean, acc_mean, se_mean, sp_mean, dices_std[1],   hd_std[1],     iou_std,   acc_std,  se_std,  sp_std


# ------------------------------------------------------------ #
def get_eval(valloader, model, criterion, opt, args):
    if opt.eval_mode == "mask_slice":
        return eval_mask_slice2(valloader, model, criterion, opt, args)
    else:
        raise RuntimeError("Could not find the eval mode:", opt.eval_mode)