mirror of
https://github.com/JDAI-CV/fast-reid.git
synced 2025-06-03 14:50:47 +08:00
12957f66aa
1. delete redundant preprocess 2. add data prefetcher to accelerate data loading 3. fix minor bug of triplet sampler when only one image for one id
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
|