import torch


class Cutout:
    def __init__(self, size=16, p=0.5):
        self.size = size
        self.half_size = size // 2
        self.p = p

    def __call__(self, image):
        assert len(image.shape) == 3  # It should be non-batch image.
        if torch.rand([1]).item() > self.p:
            return image

        left = torch.randint(
            -self.half_size, image.size(1) - self.half_size, [1]
        ).item()
        top = torch.randint(-self.half_size, image.size(2) - self.half_size, [1]).item()
        right = min(image.size(1), left + self.size)
        bottom = min(image.size(2), top + self.size)

        image[:, max(0, left) : right, max(0, top) : bottom] = 0
        return image
