losses.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Optional


class CELoss(nn.Module):
    def __init__(self, ignore_label=None, weight=None):
        super().__init__()
        if weight is not None:
            weight = torch.from_numpy(weight).float()
            print(f'----->Using weighted CE Loss weights: {weight}')

        self.loss = nn.CrossEntropyLoss(ignore_index=ignore_label, weight=weight)
        self.ignored_label = ignore_label

    def forward(self, preds, gt):

        loss = self.loss(preds, gt)
        return loss


class SoftCELoss(nn.Module):
    def __init__(self, ignore_label=None):
        super().__init__()

        self.ignore_label = ignore_label

    @staticmethod
    def soft_ce(preds, gt):
        log_probs = F.log_softmax(preds, dim=1)
        loss = -(gt * log_probs).sum() / preds.shape[0]
        return loss

    def forward(self, preds, gt):
        bs, num_pts, num_classes = preds.shape

        preds = preds.view(-1, num_classes)

        gt = gt.view(-1)
        if self.ignore_label is not None:
            valid_idx = torch.logical_not(self.ignore_label == gt)
            preds = preds[valid_idx]
            gt = gt[valid_idx]

        return self.soft_ce(preds, gt)


class DICELoss(nn.Module):

    def __init__(self, ignore_label=None, powerize=True, use_tmask=True):
        super(DICELoss, self).__init__()

        if ignore_label is not None:
            self.ignore_label = torch.tensor(ignore_label)
        else:
            self.ignore_label = ignore_label

        self.powerize = powerize
        self.use_tmask = use_tmask

    def forward(self, output, target):
        input_device = output.device
        # temporal solution to avoid nan
        output = output.cpu()
        target = target.cpu()

        if self.ignore_label is not None:
            valid_idx = torch.logical_not(target == self.ignore_label)
            target = target[valid_idx]
            output = output[valid_idx, :]

        target = F.one_hot(target, num_classes=output.shape[1])
        output = F.softmax(output, dim=-1)

        intersection = (output * target).sum(dim=0)
        if self.powerize:
            union = (output.pow(2).sum(dim=0) + target.sum(dim=0)) + 1e-12
        else:
            union = (output.sum(dim=0) + target.sum(dim=0)) + 1e-12
        if self.use_tmask:
            tmask = (target.sum(dim=0) > 0).int()
        else:
            tmask = torch.ones(target.shape[1]).int()

        iou = (tmask * 2 * intersection / union).sum(dim=0) / (tmask.sum(dim=0) + 1e-12)

        dice_loss = 1 - iou.mean()

        return dice_loss.to(input_device)


def get_soft(t_vector, eps=0.25):

    max_val = 1 - eps
    min_val = eps / (t_vector.shape[-1] - 1)

    t_soft = torch.empty(t_vector.shape)
    t_soft[t_vector == 0] = min_val
    t_soft[t_vector == 1] = max_val

    return t_soft

def get_soft_new(t_vector, score, eps=0.25):
    eps = eps * (1 - score)
    max_val = 1 - eps
    min_val = eps / (t_vector.shape[-1] - 1)

    min_val = min_val.unsqueeze(1).expand_as(t_vector)
    max_val = max_val.unsqueeze(1).expand_as(t_vector)

    t_soft = torch.empty(t_vector.shape).cuda()
    t_soft[t_vector == 0] = min_val[t_vector == 0]
    t_soft[t_vector == 1] = max_val[t_vector == 1]

    return t_soft


class SoftDICELoss(nn.Module):

    def __init__(self, ignore_label=None, powerize=True, use_tmask=True,
                 neg_range=False, eps=0.):
        super(SoftDICELoss, self).__init__()

        if ignore_label is not None:
            self.ignore_label = torch.tensor(ignore_label)
        else:
            self.ignore_label = ignore_label
        self.powerize = powerize
        self.use_tmask = use_tmask
        self.neg_range = neg_range
        self.eps = eps

    def forward(self, output, target, return_class=False, score=None, loss_method_num=0):
        input_device = output.device
        # temporal solution to avoid nan
        output = output
        target = target

        if self.ignore_label is not None:
            valid_idx = torch.logical_not(target == self.ignore_label)
            target = target[valid_idx]
            output = output[valid_idx, :]

            if score is not None:
                score = score.squeeze()
                score = score[valid_idx]

        target_onehot = F.one_hot(target, num_classes=output.shape[1])

        if score is not None and (loss_method_num==1 or loss_method_num==3):
            target_soft = get_soft_new(target_onehot, score, eps=self.eps).cuda()
        else:
            target_soft = get_soft(target_onehot, eps=self.eps).cuda()

        # target_soft = get_soft(target_onehot, eps=self.eps).cuda()

        output = F.softmax(output, dim=-1)
        
        if score is not None and (loss_method_num==2 or loss_method_num==3):
            intersection = (output * target_soft * score.unsqueeze(1).expand_as(output)).sum(dim=0)
        else:
            intersection = (output * target_soft).sum(dim=0)

        # intersection = (output * target_soft).sum(dim=0)

        if self.powerize:
            union = (output.pow(2).sum(dim=0) + target_soft.sum(dim=0)) + 1e-12
        else:
            union = (output.sum(dim=0) + target_soft.sum(dim=0)) + 1e-12
        if self.use_tmask:
            tmask = (target_onehot.sum(dim=0) > 0).int()
        else:
            tmask = torch.ones(target_onehot.shape[1]).int()

        iou = (tmask * 2 * intersection / union).sum(dim=0) / (tmask.sum(dim=0) + 1e-12)
        iou_class = tmask * 2 * intersection / union

        if self.neg_range:
            dice_loss = -iou.mean()
            dice_class = -iou_class
        else:
            dice_loss = 1 - iou.mean()
            dice_class = 1 - iou_class
        if return_class:
            return dice_loss.to(input_device), dice_class
        else:
            return dice_loss.to(input_device)


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

    def forward(self, x):
        b = F.softmax(x, dim=-1) * F.log_softmax(x, dim=-1)
        b = -1.0 * b.sum(dim=-1)
        return b


class SCELoss(torch.nn.Module):
    def __init__(self, alpha, beta, num_classes=10, reduction='mean', ignore_label=None):
        super(SCELoss, self).__init__()
        self.device = 'cpu'
        self.alpha = alpha
        self.beta = beta
        self.num_classes = num_classes
        self.reduction = reduction
        self.ignore_label = ignore_label
        self.cross_entropy = torch.nn.CrossEntropyLoss(reduction=reduction)

    def forward(self, pred, labels):

        pred = pred.cpu()
        labels = labels.cpu()
        if self.ignore_label is not None:
            valid_idx = torch.logical_not(labels == self.ignore_label)
            pred = pred[valid_idx]
            labels = labels[valid_idx]

        # CCE
        ce = self.cross_entropy(pred, labels)

        # RCE
        pred = F.softmax(pred, dim=-1)
        pred = torch.clamp(pred, min=1e-4, max=1.0)
        label_one_hot = torch.nn.functional.one_hot(labels, self.num_classes).float()
        label_one_hot = torch.clamp(label_one_hot, min=1e-4, max=1.0)
        rce = (-1*torch.sum(pred * torch.log(label_one_hot), dim=1))

        if self.reduction == 'mean':
            rce = rce.mean()
        # Loss
        loss = self.alpha * ce + self.beta * rce
        return loss