PyTorch helpers (albumentations.pytorch)


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

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.

  • num_classes (int) – only for segmentation
  • sigmoid (bool, optional) – only for segmentation, transform mask to LongTensor or not.
  • normalize (dict, optional) – dict with keys [mean, std] to pass it into torchvision.normalize
class albumentations.pytorch.transforms.ToTensorV2(always_apply=True, p=1.0)[source]

Convert image and mask to torch.Tensor.