mirror of https://github.com/JDAI-CV/fast-reid.git
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
# encoding: utf-8
|
|
"""
|
|
@author: liaoxingyu
|
|
@contact: sherlockliao01@gmail.com
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def to_tensor(pic):
|
|
"""Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
|
|
|
|
See ``ToTensor`` for more details.
|
|
|
|
Args:
|
|
pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
|
|
|
|
Returns:
|
|
Tensor: Converted image.
|
|
"""
|
|
if isinstance(pic, np.ndarray):
|
|
assert len(pic.shape) in (2, 3)
|
|
# handle numpy array
|
|
if pic.ndim == 2:
|
|
pic = pic[:, :, None]
|
|
|
|
img = torch.from_numpy(pic.transpose((2, 0, 1)))
|
|
# backward compatibility
|
|
if isinstance(img, torch.ByteTensor):
|
|
return img.float()
|
|
else:
|
|
return img
|
|
|
|
# handle PIL Image
|
|
if pic.mode == 'I':
|
|
img = torch.from_numpy(np.array(pic, np.int32, copy=False))
|
|
elif pic.mode == 'I;16':
|
|
img = torch.from_numpy(np.array(pic, np.int16, copy=False))
|
|
elif pic.mode == 'F':
|
|
img = torch.from_numpy(np.array(pic, np.float32, copy=False))
|
|
elif pic.mode == '1':
|
|
img = 255 * torch.from_numpy(np.array(pic, np.uint8, copy=False))
|
|
else:
|
|
img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
|
|
# PIL image mode: L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK
|
|
if pic.mode == 'YCbCr':
|
|
nchannel = 3
|
|
elif pic.mode == 'I;16':
|
|
nchannel = 1
|
|
else:
|
|
nchannel = len(pic.mode)
|
|
img = img.view(pic.size[1], pic.size[0], nchannel)
|
|
# put it from HWC to CHW format
|
|
# yikes, this transpose takes 80% of the loading time/CPU
|
|
img = img.transpose(0, 1).transpose(0, 2).contiguous()
|
|
if isinstance(img, torch.ByteTensor):
|
|
return img.float()
|
|
else:
|
|
return img
|