Class ColorCorrector

pointersect/inference/structures.py:4249–4297

class ColorCorrector(torch.nn.Module):
    def __init__(
            self,
            correction_type: str = 'wrgb',
    ):
        """
        Apply the color correction to an rgbd_image

        Args:
            correction_type:
                'wrgb': the correction is 3 scalars \in [0, 1] that multiply to RGB channels separately
                'identify': do nothing
        """
        super().__init__()
        self.correction_type = correction_type
        if self.correction_type == 'wrgb':
            self.wrgb = torch.nn.parameter.Parameter(torch.ones(3))
        elif self.correction_type == 'identify':
            self.register_buffer('wrgb', torch.ones(3))
        else:
            raise NotImplementedError

    def get_extra_state(self):
        return dict(
            correction_type=self.correction_type,
        )

    def set_extra_state(self, state):
        self.correction_type = state['correction_type']

    def forward(
            self,
            x: torch.Tensor,
    ) -> torch.Tensor:
        """
        apply the color correction
        Args:
            x:
                (*, 3)
        Returns:
            (*, 3) corrected x
        """
        if self.correction_type == 'wrgb':
            y = x * self.wrgb.reshape(*([1] * (x.ndim - 1)), -1)
            return y
        elif self.correction_type == 'identify':
            return x
        else:
            raise NotImplementedError