| import torch | |
| from einops import einsum, reduce, repeat | |
| from jaxtyping import Float | |
| from torch import Tensor | |
| from ..types import BatchedExample | |
| def inverse_normalize_image(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)): | |
| mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) | |
| std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) | |
| return tensor * std + mean | |
| def normalize_image(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)): | |
| mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) | |
| std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) | |
| return (tensor - mean) / std | |
| def apply_normalize_shim( | |
| batch: BatchedExample, | |
| mean: tuple[float, float, float] = (0.5, 0.5, 0.5), | |
| std: tuple[float, float, float] = (0.5, 0.5, 0.5), | |
| ) -> BatchedExample: | |
| batch["context"]["image"] = normalize_image(batch["context"]["image"], mean, std) | |
| return batch | |