Diff of /loss/IoU.py [000000] .. [f77492]

Switch to side-by-side view

--- a
+++ b/loss/IoU.py
@@ -0,0 +1,33 @@
+import torch
+from .utils import *
+import numpy as np
+
+
+def IoU_loss(input, target, threshold=0.5):
+    """
+    2d dice loss
+    :param input: predict tensor
+    :param target: target tensor
+    :return: scalar loss value
+    """
+    
+    input = input > 0.5
+    target = target == torch.max(target)
+
+    input = to_float_and_cuda(input)
+    target = to_float_and_cuda(target)
+    num = input * target
+    num = torch.sum(num, dim=2)
+    num = torch.sum(num, dim=2)
+
+    den1 = input * input
+    den1 = torch.sum(den1, dim=2)
+    den1 = torch.sum(den1, dim=2)
+
+    den2 = target * target
+    den2 = torch.sum(den2, dim=2)
+    den2 = torch.sum(den2, dim=2)
+
+    iou =  num / (den1 + den2 - num) + 1e-6
+    iou_total = 1 - 1 * torch.sum(iou) / iou.size(0)  # divide by batchsize
+    return iou_total