import torch


def batch_center_mask(img_size, num_pixels, batch_size):
    """Masks all the output except the num_pixels by num_pixels central square
    of the image.

    Parameters
    ----------
    img_size : see single_random_mask

    num_pixels : int
        Should be even. If not even, num_pixels will be replaced with
        num_pixels - 1.

    batch_size : int
        Number of masks to create.
    """
    mask = torch.zeros(batch_size, 1, *img_size[1:])
    _, height, width = img_size
    lower_height = int(height / 2 - num_pixels / 2)
    upper_height = int(height / 2 + num_pixels / 2)
    lower_width = int(width / 2 - num_pixels / 2)
    upper_width = int(width / 2 + num_pixels / 2)
    mask[:, :, lower_height:upper_height, lower_width:upper_width] = 1.
    return mask