File size: 1,517 Bytes
f71ac1d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """ToTensor transformation."""
import numpy as np
import torch
from vis4d.data.const import CommonKeys as K
from vis4d.data.typing import DictData
from .base import Transform
def _replace_arrays(data: DictData) -> None:
"""Replace numpy arrays with tensors."""
for key in data.keys():
if key in [K.images, K.original_images]:
if not data[key].flags.c_contiguous:
data[key] = np.ascontiguousarray(
data[key].transpose(0, 3, 1, 2)
)
data[key] = torch.from_numpy(data[key])
else:
data[key] = (
torch.from_numpy(data[key])
.permute(0, 3, 1, 2)
.contiguous()
)
elif isinstance(data[key], np.ndarray):
data[key] = torch.from_numpy(data[key])
elif isinstance(data[key], dict):
_replace_arrays(data[key])
elif isinstance(data[key], list):
for i, entry in enumerate(data[key]):
if isinstance(entry, np.ndarray):
data[key][i] = torch.from_numpy(entry)
@Transform("data", "data")
class ToTensor:
"""Transform all entries in a list of DataDict from numpy to torch.
Note that we reshape K.images from NHWC to NCHW.
"""
def __call__(self, batch: list[DictData]) -> list[DictData]:
"""Transform all entries to tensor."""
for data in batch:
_replace_arrays(data)
return batch
|