Skip to content

Transforms (pytorch.transforms)

class albumentations.pytorch.transforms.ToTensor (num_classes=1, sigmoid=True, normalize=None) [view source on GitHub]

Convert image and mask to torch.Tensor and divide by 255 if image or mask are uint8 type. This transform is now removed from Albumentations. If you need it downgrade the library to version 0.5.2.


Name Type Description
num_classes int

only for segmentation

sigmoid bool

only for segmentation, transform mask to LongTensor or not.

normalize dict

dict with keys [mean, std] to pass it into torchvision.normalize

class albumentations.pytorch.transforms.ToTensorV2 (transpose_mask=False, always_apply=True, p=1.0) [view source on GitHub]

Convert image and mask to torch.Tensor. The numpy HWC image is converted to pytorch CHW tensor. If the image is in HW format (grayscale image), it will be converted to pytorch HW tensor. This is a simplified and improved version of the old ToTensor transform (ToTensor was deprecated, and now it is not present in Albumentations. You should use ToTensorV2 instead).


Name Type Description
transpose_mask bool

if True and an input mask has three dimensions, this transform will transpose dimensions

standard format for PyTorch Tensors. Default