apply the color correction Args: x: (*, 3) Returns: (*, 3) corrected x
(
self,
x: torch.Tensor,
)
| 4277 | self.correction_type = state['correction_type'] |
| 4278 | |
| 4279 | def forward( |
| 4280 | self, |
| 4281 | x: torch.Tensor, |
| 4282 | ) -> torch.Tensor: |
| 4283 | """ |
| 4284 | apply the color correction |
| 4285 | Args: |
| 4286 | x: |
| 4287 | (*, 3) |
| 4288 | Returns: |
| 4289 | (*, 3) corrected x |
| 4290 | """ |
| 4291 | if self.correction_type == 'wrgb': |
| 4292 | y = x * self.wrgb.reshape(*([1] * (x.ndim - 1)), -1) |
| 4293 | return y |
| 4294 | elif self.correction_type == 'identify': |
| 4295 | return x |
| 4296 | else: |
| 4297 | raise NotImplementedError |