# Source code for archai.datasets.transforms.lighting

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

    from typing import List

    import torch

    class Lighting:
        """Lighting noise(AlexNet - style PCA - based noise)"""

        def __init__(self, alphastd, eigval:List[float], eigvec:List[float]):
            self.alphastd = alphastd
            self.eigval = torch.Tensor(eigval)
            self.eigvec = torch.Tensor(eigvec)

        def __call__(self, img):
            if self.alphastd == 0:
                return img

            alpha = img.new().resize_(3).normal_(0, self.alphastd)
            rgb = self.eigvec.type_as(img).clone() \
                .mul(alpha.view(1, 3).expand(3, 3)) \
                .mul(self.eigval.view(1, 3).expand(3, 3)) \
                .sum(1).squeeze()

            return img.add(rgb.view(3, 1, 1).expand_as(img))