diff --git a/mu_map/dataset/transform.py b/mu_map/dataset/transform.py index e2fad43adf45ea4788c97e7e9ae54d3f7337ef68..1d4161492ccf95c1af9fc60a260ac3eeced096b5 100644 --- a/mu_map/dataset/transform.py +++ b/mu_map/dataset/transform.py @@ -1,6 +1,7 @@ import math from typing import List, Tuple +import numpy as np import torch from torch import Tensor @@ -126,6 +127,20 @@ class PadCropTranform(SequenceTransform): ) +class ToNumpy(): + """ + Special transformer that converts torch tensors to numpy array. + """ + def __call__(self, *tensors: Tensor) -> Tuple[np.ndarray, ...]: + return tuple(map(lambda tensor: tensor.numpy(), tensors)) + +class ToTensor(): + """ + Special transformer that converts numpy arrays to torch tensors. + """ + def __call__(self, *tensors: np.ndarray) -> Tuple[Tensor, ...]: + return tuple(map(lambda tensor: torch.from_numpy(tensor), tensors)) + if __name__ == "__main__": transform = PadCropTranform(dim=3, size=32)