# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/transforms/01_resize.ipynb.

# %% auto 0
__all__ = ['ResizeWithSpacing']

# %% ../../nbs/transforms/01_resize.ipynb 2
import torch
from monai.data import MetaTensor
from monai.transforms.spatial.array import Resize, TraceKeys

# %% ../../nbs/transforms/01_resize.ipynb 4
class ResizeWithSpacing(Resize):  # Turns out that meta_tensor.pixdim already tracks this and so this class is unhelpful
    def __call__(
        self,
        img: MetaTensor,
        *args,
        **kwargs,
    ) -> MetaTensor:
        old_shape = torch.tensor(img.shape[1:])  # Channel is first dim
        old_spacing = torch.tensor(img.meta["spacing"])
        new_shape = torch.tensor(self.spatial_size)
        new_spacing = old_spacing * old_shape / new_shape
        img.meta["spacing"] = new_spacing

        return super().__call__(img, *args, **kwargs)

    def inverse_transform(self, data: MetaTensor, transform) -> MetaTensor:
        new_shape = data.shape[1:]
        new_spacing = data.meta["spacing"]
        old_shape = transform[TraceKeys.ORIG_SIZE]
        old_spacing = new_spacing * new_shape / old_shape
        data.meta["spacing"] = old_spacing
        return super().inverse_transform(data, transform)
