albumentations.pytorch.transforms
Module containing PyTorch-specific transforms for Albumentations. This module provides transforms that convert NumPy arrays to PyTorch tensors in the appropriate format. It handles both 2D image data and 3D volumetric data, ensuring that the tensor dimensions are correctly arranged according to PyTorch's expected format (channels first). These transforms are typically used as the final step in an augmentation pipeline before feeding data to a PyTorch model.
Members
- classToTensor3D
- classToTensorV2
ToTensor3Dclass
ToTensor3D(
p: float = 1.0
)Convert 3D volumes and masks to PyTorch tensors. This transform is designed for 3D medical imaging data. It converts numpy arrays to PyTorch tensors and ensures consistent channel positioning. For all inputs (volumes and masks): - Input: (D, H, W, C) or (D, H, W) - depth, height, width, [channels] - Output: (C, D, H, W) - channels first format for PyTorch For single-channel input, adds C=1 dimension Note: This transform always moves channels to first position as this is the standard PyTorch format. For masks that need to stay in DHWC format, use a different transform or handle the transposition after this transform.
Parameters
| Name | Type | Default | Description |
|---|---|---|---|
| p | float | 1.0 | Probability of applying the transform. Default: 1.0 |
Notes
This transform always moves channels to first position as this is the standard PyTorch format. For masks that need to stay in DHWC format, use a different transform or handle the transposition after this transform.
ToTensorV2class
ToTensorV2(
transpose_mask: bool = False,
p: float = 1.0
)Converts images/masks to PyTorch Tensors, inheriting from BasicTransform. For images: - If input is in `HWC` format, converts to PyTorch `CHW` format - If input is in `HW` format, converts to PyTorch `1HW` format (adds channel dimension)
Parameters
| Name | Type | Default | Description |
|---|---|---|---|
| transpose_mask | bool | False | - |
| p | float | 1.0 | - |