def _tensorize()

in src/datasets/formatting/torch_formatter.py [0:0]

• 29 lines of code
• 21 McCabe index (conditional complexity)

    def _tensorize(self, value):
        import torch

        if isinstance(value, (str, bytes, type(None))):
            return value
        elif isinstance(value, (np.character, np.ndarray)) and np.issubdtype(value.dtype, np.character):
            return value.tolist()

        default_dtype = {}

        if isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.integer):
            default_dtype = {"dtype": torch.int64}

            # Convert dtype to np.int64 if it's either np.uint16 or np.uint32 to ensure compatibility.
            # np.uint64 is excluded from this conversion as there is no compatible PyTorch dtype that can handle it without loss.
            if value.dtype in [np.uint16, np.uint32]:
                value = value.astype(np.int64)

        elif isinstance(value, (np.number, np.ndarray)) and np.issubdtype(value.dtype, np.floating):
            default_dtype = {"dtype": torch.float32}

        if config.PIL_AVAILABLE and "PIL" in sys.modules:
            import PIL.Image

            if isinstance(value, PIL.Image.Image):
                value = np.asarray(value)
                if value.ndim == 2:
                    value = value[:, :, np.newaxis]

                value = value.transpose((2, 0, 1))
        if config.TORCHVISION_AVAILABLE and "torchvision" in sys.modules:
            from torchvision.io import VideoReader

            if isinstance(value, VideoReader):
                return value  # TODO(QL): set output to torch tensors ?
        if config.TORCHCODEC_AVAILABLE and "torchcodec" in sys.modules:
            from torchcodec.decoders import AudioDecoder, VideoDecoder

            if isinstance(value, (VideoDecoder, AudioDecoder)):
                return value  # TODO(QL): set output to jax arrays ?

        return torch.tensor(value, **{**default_dtype, **self.torch_tensor_kwargs})