mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-11-14 09:08:18 +00:00
81 lines
2.4 KiB
Python
81 lines
2.4 KiB
Python
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
def dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon=1e-6) -> float:
|
|
"""Average of Dice coefficient for all batches, or for a single mask.
|
|
|
|
Args:
|
|
input (Tensor): _description_
|
|
target (Tensor): _description_
|
|
reduce_batch_first (bool, optional): _description_. Defaults to False.
|
|
epsilon (_type_, optional): _description_. Defaults to 1e-6.
|
|
|
|
Raises:
|
|
ValueError: _description_
|
|
|
|
Returns:
|
|
float: _description_
|
|
"""
|
|
assert input.size() == target.size()
|
|
|
|
if input.dim() == 2 and reduce_batch_first:
|
|
raise ValueError(f"Dice: asked to reduce batch but got tensor without batch dimension (shape {input.shape})")
|
|
|
|
if input.dim() == 2 or reduce_batch_first:
|
|
inter = torch.dot(input.reshape(-1), target.reshape(-1))
|
|
sets_sum = torch.sum(input) + torch.sum(target)
|
|
|
|
if sets_sum.item() == 0:
|
|
sets_sum = 2 * inter
|
|
|
|
return (2 * inter + epsilon) / (sets_sum + epsilon)
|
|
else:
|
|
# compute and average metric for each batch element
|
|
dice = 0
|
|
|
|
for i in range(input.shape[0]):
|
|
dice += dice_coeff(input[i, ...], target[i, ...])
|
|
|
|
return dice / input.shape[0]
|
|
|
|
|
|
def multiclass_dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon=1e-6) -> float:
|
|
"""Average of Dice coefficient for all classes.
|
|
|
|
Args:
|
|
input (Tensor): _description_
|
|
target (Tensor): _description_
|
|
reduce_batch_first (bool, optional): _description_. Defaults to False.
|
|
epsilon (_type_, optional): _description_. Defaults to 1e-6.
|
|
|
|
Returns:
|
|
float: _description_
|
|
"""
|
|
assert input.size() == target.size()
|
|
|
|
dice = 0
|
|
|
|
for channel in range(input.shape[1]):
|
|
dice += dice_coeff(input[:, channel, ...], target[:, channel, ...], reduce_batch_first, epsilon)
|
|
|
|
return dice / input.shape[1]
|
|
|
|
|
|
def dice_loss(input: Tensor, target: Tensor, multiclass: bool = False) -> float:
|
|
"""Dice loss (objective to minimize) between 0 and 1.
|
|
|
|
Args:
|
|
input (Tensor): _description_
|
|
target (Tensor): _description_
|
|
multiclass (bool, optional): _description_. Defaults to False.
|
|
|
|
Returns:
|
|
float: _description_
|
|
"""
|
|
assert input.size() == target.size()
|
|
|
|
fn = multiclass_dice_coeff if multiclass else dice_coeff
|
|
|
|
return 1 - fn(input, target, reduce_batch_first=True)
|