mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Small tweak of timm ToTensor for clarity

Browse Source
This commit is contained in:
Ross Wightman 2024-02-10 14:57:40 -08:00
parent 5a58f4d3dc
commit 7d121ac2ef
1 changed files with 2 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
timm/data/transforms.py
View File

@ -32,16 +32,12 @@ class ToNumpy:
class ToTensor:
""" ToTensor with no rescaling of values"""
def __init__(self, dtype=torch.float32):
self.dtype = dtype
def __call__(self, pil_img):
np_img = np.array(pil_img, dtype=np.uint8)
if np_img.ndim < 3:
np_img = np.expand_dims(np_img, axis=-1)
np_img = np.rollaxis(np_img, 2) # HWC to CHW
return torch.from_numpy(np_img).to(dtype=self.dtype)
return F.pil_to_tensor(pil_img).to(dtype=self.dtype)
# Pillow is deprecating the top-level resampling attributes (e.g., Image.BILINEAR) in