utils.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torchvision.transforms.functional as TF
import numpy as np
import os
import math
import random
import logging
import logging.handlers
from matplotlib import pyplot as plt
from scipy.ndimage import zoom
import SimpleITK as sitk
from medpy import metric
from thop import profile

def set_seed(seed):
    # for hash
    os.environ['PYTHONHASHSEED'] = str(seed)
    # for python and numpy
    random.seed(seed)
    np.random.seed(seed)
    # for cpu gpu
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # for cudnn
    cudnn.benchmark = False
    cudnn.deterministic = True


def get_logger(name, log_dir):
    '''
    Args:
        name(str): name of logger
        log_dir(str): path of log
    '''

    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)

    info_name = os.path.join(log_dir, '{}.info.log'.format(name))
    info_handler = logging.handlers.TimedRotatingFileHandler(info_name,
                                                             when='D',
                                                             encoding='utf-8')
    info_handler.setLevel(logging.INFO)

    formatter = logging.Formatter('%(asctime)s - %(message)s',
                                  datefmt='%Y-%m-%d %H:%M:%S')

    info_handler.setFormatter(formatter)

    logger.addHandler(info_handler)

    return logger


def log_config_info(config, logger):
    config_dict = config.__dict__
    log_info = f'#----------Config info----------#'
    logger.info(log_info)
    for k, v in config_dict.items():
        if k[0] == '_':
            continue
        else:
            log_info = f'{k}: {v},'
            logger.info(log_info)



def get_optimizer(config, model):
    assert config.opt in ['Adadelta', 'Adagrad', 'Adam', 'AdamW', 'Adamax', 'ASGD', 'RMSprop', 'Rprop', 'SGD'], 'Unsupported optimizer!'

    if config.opt == 'Adadelta':
        return torch.optim.Adadelta(
            model.parameters(),
            lr = config.lr,
            rho = config.rho,
            eps = config.eps,
            weight_decay = config.weight_decay
        )
    elif config.opt == 'Adagrad':
        return torch.optim.Adagrad(
            model.parameters(),
            lr = config.lr,
            lr_decay = config.lr_decay,
            eps = config.eps,
            weight_decay = config.weight_decay
        )
    elif config.opt == 'Adam':
        return torch.optim.Adam(
            model.parameters(),
            lr = config.lr,
            betas = config.betas,
            eps = config.eps,
            weight_decay = config.weight_decay,
            amsgrad = config.amsgrad
        )
    elif config.opt == 'AdamW':
        return torch.optim.AdamW(
            model.parameters(),
            lr = config.lr,
            betas = config.betas,
            eps = config.eps,
            weight_decay = config.weight_decay,
            amsgrad = config.amsgrad
        )
    elif config.opt == 'Adamax':
        return torch.optim.Adamax(
            model.parameters(),
            lr = config.lr,
            betas = config.betas,
            eps = config.eps,
            weight_decay = config.weight_decay
        )
    elif config.opt == 'ASGD':
        return torch.optim.ASGD(
            model.parameters(),
            lr = config.lr,
            lambd = config.lambd,
            alpha  = config.alpha,
            t0 = config.t0,
            weight_decay = config.weight_decay
        )
    elif config.opt == 'RMSprop':
        return torch.optim.RMSprop(
            model.parameters(),
            lr = config.lr,
            momentum = config.momentum,
            alpha = config.alpha,
            eps = config.eps,
            centered = config.centered,
            weight_decay = config.weight_decay
        )
    elif config.opt == 'Rprop':
        return torch.optim.Rprop(
            model.parameters(),
            lr = config.lr,
            etas = config.etas,
            step_sizes = config.step_sizes,
        )
    elif config.opt == 'SGD':
        return torch.optim.SGD(
            model.parameters(),
            lr = config.lr,
            momentum = config.momentum,
            weight_decay = config.weight_decay,
            dampening = config.dampening,
            nesterov = config.nesterov
        )
    else: # default opt is SGD
        return torch.optim.SGD(
            model.parameters(),
            lr = 0.01,
            momentum = 0.9,
            weight_decay = 0.05,
        )


def get_scheduler(config, optimizer):
    assert config.sch in ['StepLR', 'MultiStepLR', 'ExponentialLR', 'CosineAnnealingLR', 'ReduceLROnPlateau',
                        'CosineAnnealingWarmRestarts', 'WP_MultiStepLR', 'WP_CosineLR'], 'Unsupported scheduler!'
    if config.sch == 'StepLR':
        scheduler = torch.optim.lr_scheduler.StepLR(
            optimizer,
            step_size = config.step_size,
            gamma = config.gamma,
            last_epoch = config.last_epoch
        )
    elif config.sch == 'MultiStepLR':
        scheduler = torch.optim.lr_scheduler.MultiStepLR(
            optimizer,
            milestones = config.milestones,
            gamma = config.gamma,
            last_epoch = config.last_epoch
        )
    elif config.sch == 'ExponentialLR':
        scheduler = torch.optim.lr_scheduler.ExponentialLR(
            optimizer,
            gamma = config.gamma,
            last_epoch = config.last_epoch
        )
    elif config.sch == 'CosineAnnealingLR':
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max = config.T_max,
            eta_min = config.eta_min,
            last_epoch = config.last_epoch
        )
    elif config.sch == 'ReduceLROnPlateau':
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, 
            mode = config.mode, 
            factor = config.factor, 
            patience = config.patience, 
            threshold = config.threshold, 
            threshold_mode = config.threshold_mode, 
            cooldown = config.cooldown, 
            min_lr = config.min_lr, 
            eps = config.eps
        )
    elif config.sch == 'CosineAnnealingWarmRestarts':
        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer,
            T_0 = config.T_0,
            T_mult = config.T_mult,
            eta_min = config.eta_min,
            last_epoch = config.last_epoch
        )
    elif config.sch == 'WP_MultiStepLR':
        lr_func = lambda epoch: epoch / config.warm_up_epochs if epoch <= config.warm_up_epochs else config.gamma**len(
                [m for m in config.milestones if m <= epoch])
        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_func)
    elif config.sch == 'WP_CosineLR':
        lr_func = lambda epoch: epoch / config.warm_up_epochs if epoch <= config.warm_up_epochs else 0.5 * (
                math.cos((epoch - config.warm_up_epochs) / (config.epochs - config.warm_up_epochs) * math.pi) + 1)
        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_func)

    return scheduler



def save_imgs(img, msk, msk_pred, i, save_path, datasets, threshold=0.5, test_data_name=None):
    img = img.squeeze(0).permute(1,2,0).detach().cpu().numpy()
    img = img / 255. if img.max() > 1.1 else img
    if datasets == 'retinal':
        msk = np.squeeze(msk, axis=0)
        msk_pred = np.squeeze(msk_pred, axis=0)
    else:
        msk = np.where(np.squeeze(msk, axis=0) > 0.5, 1, 0)
        msk_pred = np.where(np.squeeze(msk_pred, axis=0) > threshold, 1, 0) 

    plt.figure(figsize=(7,15))

    plt.subplot(3,1,1)
    plt.imshow(img)
    plt.axis('off')

    plt.subplot(3,1,2)
    plt.imshow(msk, cmap= 'gray')
    plt.axis('off')

    plt.subplot(3,1,3)
    plt.imshow(msk_pred, cmap = 'gray')
    plt.axis('off')

    if test_data_name is not None:
        save_path = save_path + test_data_name + '_'
    plt.savefig(save_path + str(i) +'.png')
    plt.close()
    


class DiceLoss(nn.Module):
    def __init__(self):
        super(DiceLoss, self).__init__()

    def forward(self, pred, target):
        smooth = 1
        size = pred.size(0)

        pred_ = pred.view(size, -1)
        target_ = target.view(size, -1)
        intersection = pred_ * target_
        dice_score = (2 * intersection.sum(1) + smooth)/(pred_.sum(1) + target_.sum(1) + smooth)
        dice_loss = 1 - dice_score.sum()/size

        return dice_loss
    

class nDiceLoss(nn.Module):
    def __init__(self, n_classes):
        super(nDiceLoss, self).__init__()
        self.n_classes = n_classes

    def _one_hot_encoder(self, input_tensor):
        tensor_list = []
        for i in range(self.n_classes):
            temp_prob = input_tensor == i  
            tensor_list.append(temp_prob.unsqueeze(1))
        output_tensor = torch.cat(tensor_list, dim=1)
        return output_tensor.float()

    def _dice_loss(self, score, target):
        target = target.float()
        smooth = 1e-5

        intersect = torch.sum(score * target)
        y_sum = torch.sum(target)
        z_sum = torch.sum(score)
    
        loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
        loss = 1 - loss
        return loss

    def forward(self, inputs, target, weight=None, softmax=False):
        if softmax:
            inputs = torch.softmax(inputs, dim=1)
        target = self._one_hot_encoder(target)
        if weight is None:
            weight = [1] * self.n_classes
        assert len(weight) == self.n_classes, "Weight length must equal n_classes"
        assert inputs.size() == target.size(), f'Predict {inputs.size()} & target {target.size()} shape mismatch'
        loss = 0.0
        for i in range(0, self.n_classes):
            dice = self._dice_loss(inputs[:, i], target[:, i])
            loss += dice * weight[i]
        weight_sum = sum(weight)
        if weight_sum == 0:
            return torch.tensor(0.0, device=inputs.device)
        return loss / weight_sum


class CeDiceLoss(nn.Module):
    def __init__(self, num_classes, loss_weight=[0.4, 0.6]):
        super(CeDiceLoss, self).__init__()
        self.celoss = nn.CrossEntropyLoss()
        self.diceloss = nDiceLoss(num_classes)
        self.loss_weight = loss_weight
    
    def forward(self, pred, target):
        loss_ce = self.celoss(pred, target.long())
        loss_dice = self.diceloss(pred, target, softmax=True, weight=[0,1])
        loss = self.loss_weight[0] * loss_ce + self.loss_weight[1] * loss_dice
        return loss
    
class BCELoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.bce = nn.BCELoss()

    def forward(self, pred, target):
        assert target.dtype == torch.float32, "Target must be float32 with values 0.0 or 1.0"
        return self.bce(pred, target)

    
class BceDiceLoss(nn.Module):
    def __init__(self, weight_bce=1.0, weight_dice=1.0, smooth=1e-6):
        super(BceDiceLoss, self).__init__()
        self.weight_bce = weight_bce
        self.weight_dice = weight_dice
        self.smooth = smooth

    def forward(self, pred, target):
        _, C, _, _ = pred.size()
        if target.dtype == torch.float32:
            target = target.long()
        bce_loss = F.cross_entropy(pred, target)
        prob = F.softmax(pred, dim=1)
        target_one_hot = F.one_hot(target, num_classes=C).permute(0, 3, 1, 2).float()
        intersection = (prob * target_one_hot).sum(dim=(2, 3))  # [B, C]
        union = prob.sum(dim=(2, 3)) + target_one_hot.sum(dim=(2, 3))
        dice_coeff = (2. * intersection + self.smooth) / (union + self.smooth)
        dice_loss = 1 - dice_coeff.mean() 

        return self.weight_bce * bce_loss + self.weight_dice * dice_loss
 


class AdaptiveHierarchicalLoss(nn.Module):
    def __init__(self, num_layers=4, tau=0.2, alpha=0.8, grad_clip=1.0, device=torch.device('cuda')):
        super().__init__()
        self.num_layers = num_layers
        self.tau = tau
        self.alpha = alpha
        self.grad_clip = grad_clip
        self.bce_dice_loss = BceDiceLoss()
        self.register_buffer('weights', torch.ones(num_layers) / num_layers)
        self.device = device or torch.device('cpu')  # 默认CPU
        self.weights = torch.ones(num_layers, device=self.device) / num_layers
        self.grad_buffer = []

    def forward(self, final_output, target, layer_outputs):

        assert len(layer_outputs) == self.num_layers
        losses = []
        self.grad_buffer.clear()  
        
        for l in range(self.num_layers):
            _, H, W, _= layer_outputs[l].shape
            target_resized = F.interpolate(target.unsqueeze(1).float(), size=(H, W), mode='bilinear').float()
            loss_l = self.bce_dice_loss(layer_outputs[l].permute(0, 3, 1, 2), target_resized.squeeze(1)).mean()
            loss_l.register_hook(self._grad_hook(l))
            losses.append(loss_l)
            if self.grad_clip > 0:
                torch.nn.utils.clip_grad_norm_(
                    parameters=self.parameters(),  
                    max_norm=self.grad_clip,
                    error_if_nonfinite=True
                )
        

        loss_final = self.bce_dice_loss(final_output, target)
        sum(losses).backward(retain_graph=True)
        
        with torch.no_grad():
            if len(self.grad_buffer) == self.num_layers:
                grads = torch.stack(self.grad_buffer)          
                grads = (grads - grads.mean()) / (grads.std() + 1e-8)  
                weights_new = torch.softmax(grads / self.tau, dim=0)
                self.weights = self.alpha * self.weights + (1 - self.alpha) * weights_new
                self.weights /= self.weights.sum()
        total_loss = sum(w * loss for w, loss in zip(self.weights, losses))
        total_loss += 0.75 * loss_final     
        return total_loss

    def _grad_hook(self, l):
        def hook(grad):
            grad_mag = grad.abs().mean()
            self.grad_buffer.append(grad_mag)
        return hook
    def get_weights(self):
        return self.weights.detach().cpu().numpy()





def calculate_tp_fp_tn_fn(pred, gt):
    tp = np.sum((pred == 1) & (gt == 1))
    fp = np.sum((pred == 1) & (gt == 0))
    tn = np.sum((pred == 0) & (gt == 0))
    fn = np.sum((pred == 0) & (gt == 1))
    return tp, fp, tn, fn

def calculate_metric_percase(pred, gt, threshold):
    if len(pred.shape) == 3 and pred.shape[0] > 1:
        assert pred.shape[0] == gt.shape[0]
        sum_dice, sum_hd95, sum_recall, sum_iou, sum_acc, sum_spe, sum_ignore = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0
        for i in range(pred.shape[0]):
            metrics = calculate_metric_percase(pred[i], gt[i], threshold)
            sum_dice += metrics[0]
            sum_hd95 += metrics[1]
            sum_recall += metrics[2]
            sum_iou += metrics[3]
            sum_acc += metrics[4]
            sum_spe += metrics[5]
            sum_ignore += metrics[6]
        return sum_dice, sum_hd95, sum_recall, sum_iou, sum_acc, sum_spe, sum_ignore
    pred[pred > 0] = 1
    gt[gt > 0] = 1
    tp, fp, tn, fn = calculate_tp_fp_tn_fn(pred, gt)
    total = tp + fp + tn + fn
    recall = tp / (tp + fn + 1e-10) 
    iou = tp / (tp + fp + fn + 1e-10)
    accuracy = (tp + tn) / (total + 1e-10)
    spe = tn / (tn + fp + 1e-10)
    

    if pred.sum() > 0 and gt.sum() > 0:
        hd95 = metric.binary.hd95(pred, gt)
        sum_pred = np.sum(pred)
        sum_true = np.sum(gt)
        denominator = sum_pred + sum_true
        if denominator == 0:
            denominator = 1.0  
        return 2 * tp / (2*tp+fp+fn), hd95, recall, iou, accuracy, spe, 1
    elif pred.sum() > 0 and gt.sum() == 0:
        return 0.0, 0.0, 0.0, 0.0, tn/(tn+fp), tn/(tn+fp+1e-10), 0 
    elif pred.sum() == 0 and gt.sum() > 0:
        return 0.0, threshold, 0.0, 0.0, tn/(tn+fn), 1.0, 1
    else:
        return 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1  



def test_score(image, label, net, classes, patch_size=[224, 224]):
    image, label = image.squeeze(0).cpu().detach().numpy(), label.squeeze(0).cpu().detach().numpy()
    if len(image.shape) == 3:
        prediction = np.zeros_like(label)
        for ind in range(image.shape[0]):
            slice = image[ind, :, :]
            x, y = slice.shape[0], slice.shape[1]
            if x != patch_size[0] or y != patch_size[1]:
                slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y), order=3) 
            input = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0).float().cuda()
            
            net.eval()
            with torch.no_grad():
                outputs, _ = net(input)
                out = torch.argmax(torch.softmax(outputs, dim=1), dim=1).squeeze(0)
                out = out.cpu().detach().numpy()
                if x != patch_size[0] or y != patch_size[1]:
                    pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
                else:
                    pred = out
                prediction[ind] = pred
    elif len(image.shape) == 4:
        prediction = np.zeros_like(label)
        for ind in range(image.shape[0]):
            slice = image[ind, :, :, :]
            x, y = slice.shape[0], slice.shape[1]

            if x != patch_size[0] or y != patch_size[1]:
                slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y, 1), order=0)  
            input = torch.from_numpy(slice).unsqueeze(0).float().cuda().permute(0,3,1,2)
            net.eval()
            with torch.no_grad():
                outputs, _ = net(input)
                out = torch.argmax(torch.softmax(outputs, dim=1), dim=1).squeeze(0)
                out = out.cpu().detach().numpy()
                if x != patch_size[0] or y != patch_size[1]:
                    pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
                else:
                    pred = out
                prediction[ind] = pred
    else:
        input = torch.from_numpy(image).unsqueeze(
            0).unsqueeze(0).float().cuda()
        net.eval()
        with torch.no_grad():
            xxx, _ = net(input)
            out = torch.argmax(torch.softmax(xxx, dim=1), dim=1).squeeze(0)
            prediction = out.cpu().detach().numpy()
    metric_list = []
    threshold = max(patch_size) * 1.414
    metric_list.append(calculate_metric_percase(prediction, label, threshold))
    return metric_list

class Early_stop:
    def __init__(self, patience, tolerance, save_dir):
        self.cnt = 0
        self.patience = 15
        self.best_loss = None
        self.tolerance = tolerance
        self.stop = False
        self.save_dir = save_dir

    def __call__(self, test_loss, epoch, model):
        if self.best_loss is None:
            self.best_loss = test_loss
            torch.save(model.state_dict(), self.save_dir + f'epoch_{epoch} test_loss {test_loss}.pth')
        elif test_loss > self.best_loss + self.tolerance:
            self.cnt += 1
            if self.cnt >= self.patience:
                self.stop = True
        elif test_loss < self.best_loss:
            self.best_loss = test_loss
            torch.save(model.state_dict(), self.save_dir + f'epoch_{epoch} test_loss {test_loss}.pth')
            self.cnt = 0
        else:
            self.cnt = 0