| 100 | return output_tensor.float() |
| 101 | |
| 102 | class DiceLoss(nn.Module): |
| 103 | def __init__(self, n_classes): |
| 104 | super(DiceLoss, self).__init__() |
| 105 | self.n_classes = n_classes |
| 106 | |
| 107 | def _one_hot_encoder(self, input_tensor): |
| 108 | tensor_list = [] |
| 109 | for i in range(self.n_classes): |
| 110 | temp_prob = input_tensor == i # * torch.ones_like(input_tensor) |
| 111 | tensor_list.append(temp_prob.unsqueeze(1)) |
| 112 | output_tensor = torch.cat(tensor_list, dim=1) |
| 113 | return output_tensor.float() |
| 114 | |
| 115 | def _dice_loss(self, score, target): |
| 116 | target = target.float() |
| 117 | smooth = 1e-5 |
| 118 | intersect = torch.sum(score * target) |
| 119 | y_sum = torch.sum(target * target) |
| 120 | z_sum = torch.sum(score * score) |
| 121 | loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth) |
| 122 | loss = 1 - loss |
| 123 | return loss |
| 124 | |
| 125 | def forward(self, inputs, target, weight=None, softmax=False): |
| 126 | if softmax: |
| 127 | inputs = torch.softmax(inputs, dim=1) |
| 128 | target = self._one_hot_encoder(target) |
| 129 | if weight is None: |
| 130 | weight = [1] * self.n_classes |
| 131 | assert inputs.size() == target.size(), 'predict {} & target {} shape do not match'.format(inputs.size(), target.size()) |
| 132 | class_wise_dice = [] |
| 133 | loss = 0.0 |
| 134 | for i in range(0, self.n_classes): |
| 135 | dice = self._dice_loss(inputs[:, i], target[:, i]) |
| 136 | class_wise_dice.append(1.0 - dice.item()) |
| 137 | loss += dice * weight[i] |
| 138 | return loss / self.n_classes |
| 139 | |
| 140 | def calculate_metric_percase(pred, gt): |
| 141 | pred[pred > 0] = 1 |