mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 23:43:44 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
95 lines
3.4 KiB
Python
95 lines
3.4 KiB
Python
# https://github.com/mlcommons/training/blob/master/image_segmentation/pytorch/model/losses.py
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class Dice:
|
|
def __init__(self,
|
|
to_onehot_y: bool = True,
|
|
to_onehot_x: bool = False,
|
|
use_softmax: bool = True,
|
|
use_argmax: bool = False,
|
|
include_background: bool = False,
|
|
layout: str = "NCDHW"):
|
|
self.include_background = include_background
|
|
self.to_onehot_y = to_onehot_y
|
|
self.to_onehot_x = to_onehot_x
|
|
self.use_softmax = use_softmax
|
|
self.use_argmax = use_argmax
|
|
self.smooth_nr = 1e-6
|
|
self.smooth_dr = 1e-6
|
|
self.layout = layout
|
|
|
|
def __call__(self, prediction, target):
|
|
if self.layout == "NCDHW":
|
|
channel_axis = 1
|
|
reduce_axis = list(range(2, len(prediction.shape)))
|
|
else:
|
|
channel_axis = -1
|
|
reduce_axis = list(range(1, len(prediction.shape) - 1))
|
|
num_pred_ch = prediction.shape[channel_axis]
|
|
|
|
if self.use_softmax:
|
|
prediction = torch.softmax(prediction, dim=channel_axis)
|
|
elif self.use_argmax:
|
|
prediction = torch.argmax(prediction, dim=channel_axis)
|
|
|
|
if self.to_onehot_y:
|
|
target = to_one_hot(target, self.layout, channel_axis)
|
|
|
|
if self.to_onehot_x:
|
|
prediction = to_one_hot(prediction, self.layout, channel_axis)
|
|
|
|
if not self.include_background:
|
|
assert num_pred_ch > 1, \
|
|
f"To exclude background the prediction needs more than one channel. Got {num_pred_ch}."
|
|
if self.layout == "NCDHW":
|
|
target = target[:, 1:]
|
|
prediction = prediction[:, 1:]
|
|
else:
|
|
target = target[..., 1:]
|
|
prediction = prediction[..., 1:]
|
|
|
|
assert (target.shape == prediction.shape), \
|
|
f"Target and prediction shape do not match. Target: ({target.shape}), prediction: ({prediction.shape})."
|
|
|
|
intersection = torch.sum(target * prediction, dim=reduce_axis)
|
|
target_sum = torch.sum(target, dim=reduce_axis)
|
|
prediction_sum = torch.sum(prediction, dim=reduce_axis)
|
|
|
|
return (2.0 * intersection + self.smooth_nr) / (target_sum + prediction_sum + self.smooth_dr)
|
|
|
|
|
|
def to_one_hot(array, layout, channel_axis):
|
|
if len(array.shape) >= 5:
|
|
array = torch.squeeze(array, dim=channel_axis)
|
|
array = F.one_hot(array.long(), num_classes=3)
|
|
if layout == "NCDHW":
|
|
array = array.permute(0, 4, 1, 2, 3).float()
|
|
return array
|
|
|
|
|
|
class DiceCELoss(nn.Module):
|
|
def __init__(self, to_onehot_y, use_softmax, layout, include_background):
|
|
super(DiceCELoss, self).__init__()
|
|
self.dice = Dice(to_onehot_y=to_onehot_y, use_softmax=use_softmax, layout=layout,
|
|
include_background=include_background)
|
|
self.cross_entropy = nn.CrossEntropyLoss()
|
|
|
|
def forward(self, y_pred, y_true):
|
|
cross_entropy = self.cross_entropy(y_pred, torch.squeeze(y_true, dim=1).long())
|
|
dice = torch.mean(1.0 - self.dice(y_pred, y_true))
|
|
return (dice + cross_entropy) / 2
|
|
|
|
|
|
class DiceScore:
|
|
def __init__(self, to_onehot_y: bool = True, use_argmax: bool = True, layout: str = "NCDHW",
|
|
include_background: bool = False):
|
|
self.dice = Dice(to_onehot_y=to_onehot_y, to_onehot_x=True, use_softmax=False,
|
|
use_argmax=use_argmax, layout=layout, include_background=include_background)
|
|
|
|
def __call__(self, y_pred, y_true):
|
|
return torch.mean(self.dice(y_pred, y_true), dim=0)
|