Seg example

fuse/data/augmentor/augmentor_toolbox.py

@@ -27,12 +27,16 @@
 from scipy.ndimage.filters import gaussian_filter
 from scipy.ndimage.interpolation import map_coordinates
 from torch import Tensor
+import elasticdeform as ed
 
 from fuse.utils.utils_param_sampler import FuseUtilsParamSamplerGaussianPatch as Gaussian
 from fuse.utils.utils_param_sampler import FuseUtilsParamSamplerRandBool as RandBool
 from fuse.utils.utils_param_sampler import FuseUtilsParamSamplerRandInt as RandInt
 from fuse.utils.utils_param_sampler import FuseUtilsParamSamplerUniform as Uniform
 
+import PIL
+import torch.nn.functional as F
+
 
 ######## Affine augmentation
 def aug_op_affine(aug_input: Tensor, rotate: float = 0.0, translate: Tuple[float, float] = (0.0, 0.0),
@@ -63,7 +67,12 @@ def aug_op_affine(aug_input: Tensor, rotate: float = 0.0, translate: Tuple[float
     for channel in channels:
         aug_channel_tensor = aug_input[channel].numpy()
         aug_channel_tensor = Image.fromarray(aug_channel_tensor)
-        aug_channel_tensor = TTF.affine(aug_channel_tensor, angle=rotate, scale=scale, translate=translate, shear=shear)
+        aug_channel_tensor = TTF.affine(aug_channel_tensor, 
+                                        angle=rotate, 
+                                        scale=scale, 
+                                        resample=PIL.Image.BILINEAR,
+                                        translate=translate, 
+                                        shear=shear)
         if flip[0]:
             aug_channel_tensor = TTF.vflip(aug_channel_tensor)
         if flip[1]:
@@ -265,23 +274,13 @@ def aug_op_elastic_transform(aug_input: Tensor, alpha: float = 1, sigma: float =
        :param channels: which channels to apply the augmentation
        :return distorted image
     """
-    random_state = numpy.random.RandomState(None)
-    if channels is None:
-        channels = list(range(aug_input.shape[0]))
-    aug_tensor = aug_input.numpy()
-    for channel in channels:
-        aug_channel_tensor = aug_input[channel].numpy()
-        shape = aug_channel_tensor.shape
-        dx1 = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
-        dx2 = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
+    # convert back to torch tensor
+    aug_input = [numpy.array(t) for t in aug_input]
+    aug_input_d = ed.deform_random_grid(aug_input, sigma=7, points=3, axis=[(1, 2), (1,2)])
 
-        x1, x2 = numpy.meshgrid(numpy.arange(shape[0]), numpy.arange(shape[1]))
-        indices = numpy.reshape(x2 + dx2, (-1, 1)), numpy.reshape(x1 + dx1, (-1, 1))
+    aug_output = [torch.from_numpy(t) for t in aug_input_d]
 
-        distored_image = map_coordinates(aug_channel_tensor, indices, order=1, mode='reflect')
-        distored_image = distored_image.reshape(aug_channel_tensor.shape)
-        aug_tensor[channel] = distored_image
-    return torch.from_numpy(aug_tensor)
+    return aug_output
 
 
 ######### Default / Example augmentation pipline for a 2D image
@@ -452,3 +451,39 @@ def aug_op_batch_mix_up(aug_input: Tuple[Tensor, Tensor], factor: float) -> Tupl
     img = img * (1.0 - factor) + factor * img_mix_up
     labels = labels * (1.0 - factor) + factor * labels_mix_up
     return img, labels
+
+
+def aug_op_random_crop_and_resize(aug_input: Tensor,
+                                  out_size,
+                                  crop_size: float = 1.0,  # or optional - Tuple[float, float]
+                                  x_off: float = 1.0,
+                                  y_off: float = 1.0, 
+                                  z_off: float = 1.0) -> Tensor:
+    """
+    random crop a (3d) tensor and resize it to a given size
+    :param crop_size: float <= 1.0 - the fraction to crop from the original tensor for each dim
+    :param x_off: float <= 1.0 - the x-offset to take 
+    :param y_off:
+    :param z_off:
+    :param out_size: the size of the output tensor
+    :return: the output tensor
+    """
+    in_shape = aug_input.shape
+
+    if len(aug_input.shape) == 4:
+        ch, z, y, x = in_shape
+
+        x_width = int(crop_size * x)
+        x_off = int(x_off * (x - x_width))
+
+        y_width = int(crop_size * y)
+        y_off = int(y_off * (y - y_width))
+
+        z_width = int(crop_size * z)
+        z_off = int(z_off * (z - z_width))
+
+        aug_tensor = aug_input[:, z_off:z_off+z_width, y_off:y_off+y_width, x_off:x_off+x_width]
+
+        aug_tensor = F.interpolate(aug_tensor, out_size)
+
+    return aug_tensor

fuse/losses/segmentation/loss_dice.py

@@ -65,7 +65,7 @@ def __call__(self, predict, target):
         predict = predict.contiguous().view(predict.shape[0], -1)
         target = target.contiguous().view(target.shape[0], -1)
 
-        if target.dtype == torch.int64:
+        if target.dtype == torch.int64 or target.dtype == torch.int32:
             target = target.type(torch.float32).to(target.device)
         num = 2*torch.sum(torch.mul(predict, target), dim=1) + self.eps
         den = torch.sum(predict.pow(self.p) + target.pow(self.p), dim=1) + self.eps
@@ -82,6 +82,80 @@ def __call__(self, predict, target):
             raise Exception('Unexpected reduction {}'.format(self.reduction))
 
 
+class DiceBCELoss(FuseLossBase):
+
+    def __init__(self, 
+                 pred_name,
+                 target_name,
+                 filter_func: Optional[Callable]=None,
+                 class_weights=None,
+                 bce_weight: float=1.0,
+                 power: int=1, 
+                 eps: float=1.,
+                 reduction: str='mean'):
+        '''
+        Compute a weighted sum of dice-loss and cross entropy loss.
+
+        :param pred_name:        batch_dict key for predicted output (e.g., class probabilities after softmax).
+                                 Expected Tensor shape = [batch, num_classes, height, width]
+        :param target_name:      batch_dict key for target (e.g., ground truth label). Expected Tensor shape = [batch, height, width]
+        :param filter_func:      function that filters batch_dict/ The function gets ans input batch_dict and returns filtered batch_dict
+        :param class_weights:    An array of shape [num_classes,]
+        :param bce_weight:       weight to attach to the bce loss, default : 1.0
+        :param power:            Denominator value: \sum{x^p} + \sum{y^p}, default: 1
+        :param eps:              A float number to smooth loss, and avoid NaN error, default: 1
+        :param reduction:        Reduction method to apply, return mean over batch if 'mean',
+                                 return sum if 'sum', return a tensor of shape [N,] if 'none'
+
+        Returns:            Loss tensor according to arg reduction
+        Raise:              Exception if unexpected reduction
+        '''
+
+        super().__init__(pred_name, target_name, 1.0)
+        self.class_weights = class_weights
+        self.bce_weight = bce_weight
+        self.filter_func = filter_func
+        self.dice = BinaryDiceLoss(power, eps, reduction)
+
+    def __call__(self, batch_dict):
+
+        if self.filter_func is not None:
+            batch_dict = self.filter_func(batch_dict)
+        predict = FuseUtilsHierarchicalDict.get(batch_dict, self.pred_name).float()
+        target = FuseUtilsHierarchicalDict.get(batch_dict, self.target_name).long()
+
+        target = target.type(torch.float32).to(target.device)
+
+        total_loss = 0
+        n_classes = predict.shape[1]
+
+        # Convert target to one hot encoding
+        if n_classes > 1 and target.shape[1] != n_classes:
+            target = make_one_hot(target, n_classes)
+
+        assert predict.shape == target.shape, 'predict & target shape do not match'
+
+        # import ipdb; ipdb.set_trace(context=7) # BREAKPOINT
+        total_class_weights = sum(self.class_weights) if self.class_weights is not None else n_classes
+        for cls_index in range(n_classes):
+            dice_loss = self.dice(predict[:, cls_index, :, :], target[:, cls_index, :, :])
+            if self.bce_weight > 0.0:
+                bce_loss = F.binary_cross_entropy(predict[:, cls_index, :, :].view(-1), 
+                                                  target[:, cls_index, :, :].view(-1), 
+                                                  reduction='mean')
+                dice_loss += self.bce_weight * bce_loss
+
+            if self.class_weights is not None:
+                assert self.class_weights.shape[0] == n_classes, \
+                    'Expect weight shape [{}], got[{}]'.format(n_classes, self.class_weights.shape[0])
+                dice_loss *= self.class_weights[cls_index]
+
+            total_loss += dice_loss
+        total_loss /= total_class_weights
+
+        return self.weight*total_loss
+
+
 class FuseDiceLoss(FuseLossBase):
 
     def __init__(self, pred_name,