utils.py

import torch
import torch.nn.functional as F

def dice_coef(pred, target, eps=1e-6):
    pred = torch.sigmoid(pred)
    pred = (pred>0.5).float()
    inter = (pred*target).sum(dim=(2,3))
    union = pred.sum(dim=(2,3)) + target.sum(dim=(2,3))
    dice = (2*inter + eps)/(union + eps)
    return dice.mean().item()

def iou_score(pred, target, eps=1e-6):
    pred = torch.sigmoid(pred)
    pred = (pred>0.5).float()
    inter = (pred*target).sum(dim=(2,3))
    union = pred.sum(dim=(2,3)) + target.sum(dim=(2,3)) - inter
    iou = (inter + eps)/(union + eps)
    return iou.mean().item()

class BCEDiceLoss:
    def __init__(self, bce_weight=0.5, dice_weight=0.5):
        self.bce_w = bce_weight
        self.dice_w = dice_weight
    def __call__(self, logits, targets):
        bce = F.binary_cross_entropy_with_logits(logits, targets)
        probs = torch.sigmoid(logits)
        inter = (probs*targets).sum(dim=(2,3))
        union = probs.sum(dim=(2,3)) + targets.sum(dim=(2,3))
        dice = 1 - (2*inter + 1e-6)/(union + 1e-6)
        dice = dice.mean()
        return self.bce_w*bce + self.dice_w*dice