Source code for draugr.torch_utilities.opencv.torch_conversion
from typing import Tuple
import numpy
import torch
__all__ = [
"CV2ToImage",
"CV2ToTensor",
]
[docs]class CV2ToImage(object):
def __call__(
self,
tensor: torch.Tensor,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
return (
tensor.cpu().numpy().astype(numpy.float32).transpose((1, 2, 0)),
boxes,
labels,
)
[docs]class CV2ToTensor(object):
def __call__(
self,
cvimage: numpy.ndarray,
boxes: numpy.ndarray = None,
labels: numpy.ndarray = None,
) -> Tuple:
return (
torch.from_numpy(cvimage.astype(numpy.float32)).permute(2, 0, 1),
boxes,
labels,
)