# Source code for archai.datasets.transforms.cutout_custom

    # Copyright (c) Microsoft Corporation.
    # Licensed under the MIT license.

    import numpy as np

    import torch

    class CutoutCustom:
        """
        Reference : https://github.com/quark0/darts/blob/master/cnn/utils.py
        """
        def __init__(self, length):
            self.length = length

        def __call__(self, img):
            h, w = img.size(1), img.size(2)
            mask = np.ones((h, w), np.float32)
            y = np.random.randint(h)
            x = np.random.randint(w)

            y1 = np.clip(y - self.length // 2, 0, h)
            y2 = np.clip(y + self.length // 2, 0, h)
            x1 = np.clip(x - self.length // 2, 0, w)
            x2 = np.clip(x + self.length // 2, 0, w)

            mask[y1: y2, x1: x2] = 0.
            mask = torch.from_numpy(mask)
            mask = mask.expand_as(img)
            img *= mask
            return img