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. WARNING! Please use this with care and look into sources before usage.


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.


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