File size: 1,525 Bytes
abd2a81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import torch
def iou(y_pred, y_true, threshold):
assert len(y_pred.shape) == len(y_true.shape) == 2, "Input tensor shapes should be (N, .)"
mask_pred = threshold < y_pred
mask_true = threshold < y_true
intersection = torch.sum(mask_pred * mask_true, dim=-1)
union = torch.sum(mask_pred + mask_true, dim=-1)
r = intersection.float() / union.float()
r[union == 0] = 1
return r
def dice_loss(y_pred, y_true, smooth=1, eps=1e-7):
"""
@param y_pred: (N, C, H, W)
@param y_true: (N, C, H, W)
@param smooth:
@param eps:
@return: (N, C)
"""
numerator = 2 * torch.sum(y_true * y_pred, dim=(-1, -2))
denominator = torch.sum(y_true, dim=(-1, -2)) + torch.sum(y_pred, dim=(-1, -2))
return 1 - (numerator + smooth) / (denominator + smooth + eps)
def main():
# import kornia
# spatial_gradient_function = kornia.filters.SpatialGradient()
#
# image = torch.zeros((7, 7))
# image[2:5, 2:5] = 1
# print(image)
#
# grads = spatial_gradient_function(image[None, None, ...])[0, 0, ...] / 4
# print(grads[0])
# print(grads[1])
y_true = torch.tensor([
[0, 0, 0, 0, 1, 1, 1],
[0, 0, 0, 0, 1, 1, 1],
[0, 0, 0, 0, 0, 0, 0]
])
y_pred = torch.tensor([
[0, 0, 0, 0, 1, 1, 1],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
])
print(y_true.shape)
print(y_pred.shape)
r = iou(y_pred, y_true, threshold=0.5)
print(r)
if __name__ == "__main__":
main()
|