Source code for kornia.augmentation._2d.intensity.planckian_jitter

from typing import Any, Dict, List, Optional, Union

from kornia.augmentation import random_generator as rg
from kornia.augmentation._2d.intensity.base import IntensityAugmentationBase2D
from kornia.core import Tensor, stack

_planckian_coeffs = {
    'blackbody': Tensor(
        [
            [0.6743, 0.4029, 0.0013],
            [0.6281, 0.4241, 0.1665],
            [0.5919, 0.4372, 0.2513],
            [0.5623, 0.4457, 0.3154],
            [0.5376, 0.4515, 0.3672],
            [0.5163, 0.4555, 0.4103],
            [0.4979, 0.4584, 0.4468],
            [0.4816, 0.4604, 0.4782],
            [0.4672, 0.4619, 0.5053],
            [0.4542, 0.4630, 0.5289],
            [0.4426, 0.4638, 0.5497],
            [0.4320, 0.4644, 0.5681],
            [0.4223, 0.4648, 0.5844],
            [0.4135, 0.4651, 0.5990],
            [0.4054, 0.4653, 0.6121],
            [0.3980, 0.4654, 0.6239],
            [0.3911, 0.4655, 0.6346],
            [0.3847, 0.4656, 0.6444],
            [0.3787, 0.4656, 0.6532],
            [0.3732, 0.4656, 0.6613],
            [0.3680, 0.4655, 0.6688],
            [0.3632, 0.4655, 0.6756],
            [0.3586, 0.4655, 0.6820],
            [0.3544, 0.4654, 0.6878],
            [0.3503, 0.4653, 0.6933],
        ]
    ),
    'CIED': Tensor(
        [
            [0.5829, 0.4421, 0.2288],
            [0.5510, 0.4514, 0.2948],
            [0.5246, 0.4576, 0.3488],
            [0.5021, 0.4618, 0.3941],
            [0.4826, 0.4646, 0.4325],
            [0.4654, 0.4667, 0.4654],
            [0.4502, 0.4681, 0.4938],
            [0.4364, 0.4692, 0.5186],
            [0.4240, 0.4700, 0.5403],
            [0.4127, 0.4705, 0.5594],
            [0.4023, 0.4709, 0.5763],
            [0.3928, 0.4713, 0.5914],
            [0.3839, 0.4715, 0.6049],
            [0.3757, 0.4716, 0.6171],
            [0.3681, 0.4717, 0.6281],
            [0.3609, 0.4718, 0.6380],
            [0.3543, 0.4719, 0.6472],
            [0.3480, 0.4719, 0.6555],
            [0.3421, 0.4719, 0.6631],
            [0.3365, 0.4719, 0.6702],
            [0.3313, 0.4719, 0.6766],
            [0.3263, 0.4719, 0.6826],
            [0.3217, 0.4719, 0.6882],
        ]
    ),
}

_planckian_coeffs_ratio = {
    'blackbody': stack(
        (
            _planckian_coeffs['blackbody'][:, 0] / _planckian_coeffs['blackbody'][:, 1],
            _planckian_coeffs['blackbody'][:, 2] / _planckian_coeffs['blackbody'][:, 1],
        ),
        1,
    ),
    'CIED': stack(
        (
            _planckian_coeffs['CIED'][:, 0] / _planckian_coeffs['CIED'][:, 1],
            _planckian_coeffs['CIED'][:, 2] / _planckian_coeffs['CIED'][:, 1],
        ),
        1,
    ),
}


[docs]class RandomPlanckianJitter(IntensityAugmentationBase2D): r"""Apply planckian jitter transformation to input tensor. .. image:: _static/img/RandomPlanckianJitter.png This is physics based color augmentation, that creates realistic variations in chromaticity, this can simulate the illumination changes in the scene. See :cite:`zini2022planckian` for more details. Args: mode: 'blackbody' or 'CIED'. select_from: choose a list of jitters to apply from. `blackbody` range [0-24], `CIED` range [0-22] same_on_batch: apply the same transformation across the batch. p: probability that the random erasing operation will be performed. keepdim: whether to keep the output shape the same as input (True) or broadcast it to the batch form (False). Shape: - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)` - Output: :math:`(B, C, H, W)` .. note:: Input tensor must be float and normalized into [0, 1]. Examples: To apply planckian jitter based on mode >>> rng = torch.manual_seed(0) >>> input = torch.randn(1, 3, 2, 2) >>> aug = RandomPlanckianJitter(mode='CIED') >>> aug(input) tensor([[[[ 1.0000, -0.2389], [-1.7740, 0.4628]], <BLANKLINE> [[-1.0845, -1.3986], [ 0.4033, 0.8380]], <BLANKLINE> [[-0.9228, -0.5175], [-0.7654, 0.2335]]]]) To apply planckian jitter on image(s) from list of interested jitters >>> rng = torch.manual_seed(0) >>> input = torch.randn(2, 3, 2, 2) >>> aug = RandomPlanckianJitter(mode='blackbody', select_from=[23, 24, 1, 2]) >>> aug(input) tensor([[[[-1.1258, -1.1524], [-0.2506, -0.4339]], <BLANKLINE> [[ 0.8487, 0.6920], [-0.3160, -2.1152]], <BLANKLINE> [[ 0.4681, -0.1577], [ 1.4437, 0.2660]]], <BLANKLINE> <BLANKLINE> [[[ 0.2465, 1.0000], [-0.2125, -0.1653]], <BLANKLINE> [[ 0.9318, 1.0000], [ 1.0000, 0.0537]], <BLANKLINE> [[ 0.2426, -0.1621], [-0.3302, -0.9093]]]]) """ def __init__( self, mode: str = "blackbody", select_from: Optional[Union[int, List[int]]] = None, same_on_batch: bool = False, p: float = 0.5, keepdim: bool = False, return_transform: Optional[bool] = None, ) -> None: super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim) if isinstance(select_from, int): select_from = [select_from] if select_from is not None: self.register_buffer('pl', _planckian_coeffs_ratio[mode][select_from]) else: self.register_buffer('pl', _planckian_coeffs_ratio[mode]) if not isinstance(self.pl, Tensor): raise TypeError(f'Expected the `pl` property be a Tensor. Gotcha {type(self.pl)}') # the range of the sampling parameters _param_min: float = 0.0 _param_max: float = float(self.pl.shape[0]) self._param_generator = rg.PlanckianJitterGenerator([_param_min, _param_max]) def apply_transform( self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None ) -> Tensor: list_idx = params['idx'].tolist() if not isinstance(self.pl, Tensor): raise TypeError(f'Expected the `pl` property be a Tensor. Gotcha {type(self.pl)}') coeffs = self.pl[list_idx] r_w = coeffs[:, 0][..., None, None] b_w = coeffs[:, 1][..., None, None] r = input[..., 0, :, :] * r_w g = input[..., 1, :, :] b = input[..., 2, :, :] * b_w output = stack([r, g, b], -3) return output.clamp(max=1.0)