Repository: Project-MONAI/MONAI
File: monai/transforms/regularization/array.py (Python)

Code before:
            idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)]
            mask[s][idx] = 0

        return mask * data + (1 - mask) * data[perm, ...]

    def apply_on_labels(self, labels: torch.Tensor):
        weights, perm, _ = self._params
        nsamples, *dims = labels.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mixweight = weights[(Ellipsis,) + (None,) * len(dims)]
        return mixweight * labels + (1 - mixweight) * labels[perm, ...]

    def __call__(self, data: torch.Tensor, labels: torch.Tensor | None = None, randomize=True):
        data = convert_to_tensor(data, track_meta=get_track_meta())
        data_t = convert_to_tensor(data, track_meta=False)
        if labels is not None:
            labels_t = convert_to_tensor(labels, track_meta=get_track_meta())
            labels_t = convert_to_tensor(labels, track_meta=False)
        if randomize:
            self.randomize(data)
        augmented = convert_to_dst_type(self.apply(data_t), dst=data)[0]
        if labels is not None:
            augmented_label = convert_to_dst_type(self.apply(labels_t), dst=labels)[0]
        return (augmented, augmented_label) if labels is not None else augmented


class CutOut(Mixer):
    """Cutout as described in the paper:
    Terrance DeVries, Graham W. Taylor.
    Improved Regularization of Convolutional Neural Networks with Cutout,
    arXiv:1708.04552

    Class derived from :py:class:`monai.transforms.Mixer`. See corresponding
    documentation for details on the constructor parameters. Here, alpha not only determines
    the mixing weight but also the size of the random rectangles being cut put.
    Please refer to the paper for details.
    """

    def apply(self, data: torch.Tensor):
        weights, _, coords = self._params
        nsamples, _, *dims = data.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mask = torch.ones_like(data)
        for s, weight in enumerate(weights):
            lengths = [d * sqrt(1 - weight) for d in dims]
            idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)]
            mask[s][idx] = 0

Review comment (suggestion):
Hello, I'm a researcher currently working with CutMix and other similar method. My group and I noticed an inconsistent behavior in CutMix on the segmentation labels on our different machines. After some digging, we think it may come from this PR: before this change, `self.apply_on_labels(labels)` is used for the label, but the modified change use `self.apply(labels_t)` which uses the same cropping mechanism as for the image.

This PR does not mention the reason for this change, could you enlighten me please ?

Thank you in advance for your time.

Code after:
        for s, weight in enumerate(weights):
            lengths = [d * sqrt(1 - weight) for d in dims]
            idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)]
            mask[s][idx] = 0

        return mask * data + (1 - mask) * data[perm, ...]

    def apply_on_labels(self, labels: torch.Tensor):
        weights, perm, _ = self._params
        nsamples, *dims = labels.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mixweight = weights[(Ellipsis,) + (None,) * len(dims)]
        return mixweight * labels + (1 - mixweight) * labels[perm, ...]

    def __call__(self, data: torch.Tensor, labels: torch.Tensor | None = None, randomize=True):
        data_t = convert_to_tensor(data, track_meta=get_track_meta())
        if labels is not None:
            labels_t = convert_to_tensor(labels, track_meta=get_track_meta())
        if randomize:
            self.randomize(data)
        augmented = convert_to_dst_type(self.apply(data_t), dst=data)[0]
        if labels is not None:
            augmented_label = convert_to_dst_type(self.apply(labels_t), dst=labels)[0]
        return (augmented, augmented_label) if labels is not None else augmented


class CutOut(Mixer):
    """Cutout as described in the paper:
    Terrance DeVries, Graham W. Taylor.
    Improved Regularization of Convolutional Neural Networks with Cutout,
    arXiv:1708.04552

    Class derived from :py:class:`monai.transforms.Mixer`. See corresponding
    documentation for details on the constructor parameters. Here, alpha not only determines
    the mixing weight but also the size of the random rectangles being cut put.
    Please refer to the paper for details.
    """

    def apply(self, data: torch.Tensor):
        weights, _, coords = self._params
        nsamples, _, *dims = data.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mask = torch.ones_like(data)
        for s, weight in enumerate(weights):
            lengths = [d * sqrt(1 - weight) for d in dims]
            idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)]
            mask[s][idx] = 0