import torch


def put_gray_mask(imgs, mask):
    """
    fill image with gray pixels wherever mask is one
    - imgs: Tensor of size [num_images, 3, H, W] or [3, H, W]
    - mask: Tensor of size [num_images, 1, H, W] or [3, H, W]
    return: masked image of the same size as input image
    """

    reset = False
    if imgs.dim() == 3:
        imgs = torch.unsqueeze(imgs, dim=0)
        mask = torch.unsqueeze(mask, dim=0)
        reset = True

    imgs_masked = imgs * (1 - mask) + 0.5 * mask

    if reset:
        return imgs_masked[0]
    else:
        return imgs_masked