Data Loading
To train a model, we have to load the datasets firstly. In Flemme, Data loader can also be specified through configuration file as follows.
1# loader config
2loader:
3 ## dataset config
4 dataset:
5 name: ImgSegDataset
6 data_dir: images
7 label_dir: masks
8 data_suffix: jpg
9 ### you can specify data path here.
10 # data_path: path/to/single_data_set
11 ## you can specify a list of data paths
12 ## we would read all the datasets and concatenate them into one dataset
13 ## note: at least one of data_path and data_path_list need to be specified.
14 data_path_list:
15 - path/to/datasets/CVC-ClinicDB/fold1/
16 - path/to/datasets/CVC-ClinicDB/fold2/
17 - path/to/datasets/CVC-ClinicDB/fold3/
18 ## augmentations or transforms
19 data_transforms:
20 - name: Resize
21 size: [320, 256]
22 - name: ToTensor
23 ### other parameters related to dataloader
24 ### refer to https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
25 batch_size: 16
26 num_workers: 8
27 shuffle: true
We use torch.Dataloader as our data loader.
Parameters like batch_size and num_workers are directly from torch.Dataloader.
Such parameters not recognized by Flemme will be passed to the torch’s loader initialization.
In specific, Flemme will parse these parameters: dataset, data_path_list, data_suffix_list, label_suffix_list, data_transforms and label_transforms.
Datasets
We provide several dataset classes for easy and flexible data loading.
Image Datasets
ImgDataset
ImgDataset defines a dataset contains only images without any label or condition information.
1ImgDataset.__init__(self,
2 ## data_path, can be specified in dataset[data_path] or loader[data_path_list]
3 data_path,
4 ## dimension of image, can be 2 or 3
5 dim = 2,
6 ## data transforms, specified in loader[data_transforms]
7 data_transform = None,
8 ## sub dir in data_path, we only read data in data_dir
9 data_dir = '',
10 ## data_suffix, we only read data with specified data_suffix
11 ## can be specified in dataset[data_suffix] or loader[data_suffix_list] (if different dataset has different suffix)
12 data_suffix = '.png',
13 **kwargs)
For the above parameters, data_path can be specified in dataset config if there is only one dataset.
It can also be specified through data_path_list in loader config. data_transform need to be specified in loader block.
Other parameters should be specified in dataset config.
ImgClsDataset
ImgClsDataset defines a classification dataset contains images and the corresponding class labels.
1ImgClsDataset.__init__(self,
2 data_path,
3 dim = 2,
4 data_transform = None,
5 label_transform = None,
6 mode = 'train',
7 data_suffix = '.png',
8 ## pre-shuffle the samples, because the samples are load by categories.
9 ## you can also do the shuffle in data-loader
10 pre_shuffle = True,
11 ## map class names to class labels
12 ## also, samples should be store in the subdirs of data_path whose names are the corresponding class names
13 cls_label = {},
14 **kwargs)
ImgSegDataset
ImgSegDataset defines a dataset contains images and segmentation maps. Each image corresponds to one target with the same shape.
1ImgSegDataset.__init__(self,
2 ## data_path, can be specified in dataset[data_path] or loader[data_path_list]
3 data_path,
4 ## dimension of image, can be 2 or 3
5 dim = 2,
6 ## data transforms, specified in loader[data_transforms]
7 data_transform = None,
8 ## label transforms, specified in loader[label_transforms]
9 ## if loader[label_transforms] was not specified,
10 ## we would perform necessary transforms based on data transforms.
11 ## necessary transforms indicate those transforms should be performed on data and label simultaneously, such as resize and crop
12 label_transform = None,
13 ## sub dir for image data in data_path, we only read image in data_dir
14 data_dir = 'raw',
15 ## sub dir for image label in data_path, we only read label in label_dir
16 label_dir = 'label',
17 ## data_suffix, we only read data with specified data_suffix
18 ## can be specified in dataset[data_suffix] or loader[data_suffix_list]
19 data_suffix='.png',
20 ## label_suffix, we only read label with specified label_suffix
21 ## can be specified in dataset[label_suffix] or loader[label_suffix_list]
22 ## if not specified, use the same setting as data_suffix
23 label_suffix = None,
24 ## crop the non-zero region by image or label while keeping some margin
25 ## crop_nonzero should be a dict like: {'crop_by': label, 'margin':10}
26 crop_nonzero = None,
27 **kwargs)
MultiModalityImgSegDataset
MultiModalityImgSegDataset defines a dataset contains multiple modalities (such as BraTS),
which means it may contains images from different imaging devices and labels of different organs, tissues or tumors.
It has a same initialization function with ImgSegDataset, but some of the parameters can be list.
MultiModalityImgSegDataset will load images and labels from all the listed sub directories and combine them to a dataset.
1MultiModalityImgSegDataset.__init__(self,
2 data_path,
3 dim = 2,
4 data_transform = None,
5 label_transform = None,
6 ## the following dir and suffix can be a list
7 data_dir = 'raw',
8 label_dir = 'label',
9 data_suffix='.png',
10 label_suffix = None,
11 crop_nonzero = None,
12 ## how to combine data of different modalities, can be mean, sum or cat
13 data_combine = 'mean',
14 ## how to combine label of different modalities, can be mean, sum or cat
15 ## if not specified, use the same combining method as data_combine
16 label_combine = None,
17 **kwargs)
The following configuration define a loader for BraTS21 dataset:
1loader:
2 dataset:
3 name: MultiModalityImgSegDataset
4 dim: 3
5 data_dir: [flair, t1, t1ce, t2]
6 data_suffix: [flair.nii.gz, t1.nii.gz, t1ce.nii.gz, t2.nii.gz]
7 label_dir: seg
8 label_suffix: seg.nii.gz
9 data_combine: cat
10 crop_nonzero:
11 margin: [2,2,2]
12 crop_by: raw
13 data_path_list:
14 - /work/guoqingzhang/datasets/biomed_3d_dataset/BraTS2021/fold1
15 - /work/guoqingzhang/datasets/biomed_3d_dataset/BraTS2021/fold2
16 batch_size: 4
17 num_workers: 8
18 shuffle: false
19 data_transforms:
20 - name: Resize
21 size: [120, 192, 120]
22 - name: ToTensor
23 label_transforms:
24 - name: Resize
25 size: [120, 192, 120]
26 - name: Relabel
27 map:
28 - [4, 3]
29 - name: ToOneHot
30 num_classes: 4
31 ignore_background: False
32 - name: ToTensor
Point Cloud Datasets
PcdDataset
PcdDataset defines a dataset contains only point clouds without any label or condition information.
1PcdDataset.__init__(self,
2 data_path,
3 data_transform = None,
4 data_dir = '',
5 data_suffix = '.ply',
6 **kwargs)
PcdClsDataset
PcdClsDataset defines a classification dataset contains point clouds and the corresponding class labels.
1PcdClsDataset.__init__(self,
2 data_path,
3 data_transform = None,
4 label_transform = None,
5 mode = 'train',
6 data_suffix = '.ply',
7 pre_shuffle = True,
8 ## a dict that maps class names to class labels.
9 ## we pre-define the dicts of shapenet and medshapes,
10 ## therefore, cls_label can also be 'shapenet' or 'medshapes'.
11 cls_label = {},
12 **kwargs)
PcdSegDataset
PcdSegDataset defines a dataset contains images and segmentation labels.
1PcdSegDataset.__init__(self,
2 data_path,
3 data_transform = None,
4 label_transform = None,
5 mode = 'train',
6 data_dir = 'pcd',
7 label_dir = 'label',
8 data_suffix = '.ply',
9 label_suffix='.seg',
10 **kwargs):
PcdReconWithClassLabelDataset
PcdReconWithClassLabelDataset defines a completion dataset contains partial or noisy point clouds, reconstruction targets, and corresponding category labels.
1PcdReconWithClassLabelDataset.__init__(self, data_path,
2 data_transform = None,
3 ## transform for class labels
4 label_transform = None,
5 target_transform = None,
6 mode = 'train',
7 data_dir = 'partial',
8 target_dir = 'target',
9 data_suffix = '.ply',
10 target_suffix='.ply',
11 cls_label = {},
12 pre_shuffle = True,
13 **kwargs):
The following block defines a loader for MedPointS completion dataset:
1loader:
2 dataset:
3 name: PcdReconWithClassLabelDataset
4 data_dir: partial
5 target_dir: target
6 data_suffix: .ply
7 target_suffix: .ply
8 cls_label: MedPointS
9 data_path_list:
10 - /media/wlsdzyzl/DATA/datasets/pcd/MedPointS/completion/fold1
11 - /media/wlsdzyzl/DATA/datasets/pcd/MedPointS/completion/fold2
12 - /media/wlsdzyzl/DATA/datasets/pcd/MedPointS/completion/fold3
13 batch_size: 64
14 num_workers: 8
15 shuffle: true
16 data_transforms:
17 - name: Normalize
18 - name: FixedPoints
19 num: 2048
20 - name: ToTensor
21 dtype: float
22 target_transforms:
23 - name: Normalize
24 - name: FixedPoints
25 num: 2048
26 - name: ToTensor
27 dtype: float
28 label_transforms:
29 - name: ToOneHot
30 num_classes: 47
31 ignore_background: true
32 - name: ToTensor
33 dtype: float
Data Augmentations
2D Image Data Augmentation
For 2D image, we adopt the following common transforms from torchvision.transforms:
ToTensor
RandomHorizontalFlip
RandomVerticalFlip
Normalize
RandomRotation
GaussianBlur
CenterCrop
RandomCrop
The required parameters can refer to torchvision.transforms.
Beside of these, we also implement the following transforms (some of them are wrapped to have the same parameters as their 3D counterparts):
1# resize image to certain shape
2## size should be a list
3## mode should be one of ['nearest', 'bilinear', 'bicubic']
4Resize.__init__(self, size, mode = 'nearest')
5# to one hot label
6## number of classes
7## if ignore_background is true, the one-hot encoding of background (zero values) will be zero vectors
8ToOneHot.__init__(self, num_classes = None, ignore_background = False, **kwargs)
9# to binary mask
10## all values larger than threshold will be set as 1, others will be set as 0.
11ToBinaryMask.__init__(self, threshold=0)
12# to gray image
13## some version of torch vision doesn't contain this transforms
14## out_channel is the number of channel of output image
15GrayScale.__init__(self, out_channel = 1)
16# inverse color: white to black and black to white
17InverseColor.__init__(self)
18# Relabellabels into a consecutive numbers: [10, 10, 0, 6, 6] -> [2, 2, 0, 1, 1].
19## a optional map can be provided. Map should be a list with a shape of (n, 2).
20## A map like [[4, 3], [8, 4]] will relabel [4, 8] to [3, 4].
21Relabel.__init__(self, map = [], **kwargs)
22# perform elastic deformation on image
23## parameters can refer to elastic deformation. Default values are good choices.
24## for label transform, set spline_order = 1
25ElasticDeform.__init__(self, spline_order = 3,
26 alpha=2000,
27 sigma=50,
28 execution_probability=0.1)
For 2D image, most of the transforms need to be called after ToTensor because they should be performed on tensor.
3D Image Data Augmentation
We adopt and implement some common augmentations for 3D images.
1# Randomly flips the image.
2## axis_prob define the flip probability for each axis
3RandomFlip.__init__(self, axis_prob=0.5)
4# Rotate an array by 90 degrees around z-axis
5RandomRotate90.__init__(self)
6# Rotate an array by a random degrees from taken from (-angle_spectrum, angle_spectrum) interval
7## Rotation axis is picked at random from the list of provided axes.
8## mode should be one of ['reflect', 'grid-mirror', 'constant', 'grid-constant', 'nearest', 'mirror', 'grid-wrap', 'wrap'] from scipy
9RandomRotate.__init__(self, angle_spectrum=30, axes=None, mode='reflect', order=0)
10# Adjust contrast by scaling each voxel to `mean + alpha * (v - mean)`.
11RandomContrast.__init__(self, alpha=(0.5, 1.5), mean=0.0, execution_probability=0.1)
12# elastic deformation, similar to the 2D counterpart
13## if apply_3d is false, elastic deformation will be performed on each 2D slices
14ElasticDeform._init__(self, spline_order = 3, alpha=2000, sigma=50,
15 execution_probability=0.1, apply_3d=True)
16# crop width and height (x, y) to fixed shape
17## centered: always crop center region
18CropToFixed.__init__(self, size=(256, 256), centered=False)
19# normalize image with mean and std.
20## if mean and std are not specified, they will be computed based on the image
21## if channelwise is true and image has multiple channels,
22## image will be normalized in a channelwise manner.
23Normalize.__init__(self, mean=None, std=None, channelwise=False)
24# apply simple min-max normalization
25## if min_value and max_value are not specified, they will be computed based on the image
26## if norm01 is true, image will be normalized to [0, 1].
27## Otherwise, it will be normalized to [-1, 1]
28MinMaxNormalize.__init__(self, min_value=None, max_value=None,
29 norm01=True, channelwise=False)
30# resize the volume, similar to the 2D counterpart
31Resize.__init__(self, size, mode = 'nearest')
32# Converts a given input numpy.ndarray into torch.Tensor.
33## expand_dims (bool): if True, adds a channel dimension to the input data
34## dtype (np.dtype): the desired output data type
35ToTensor.__init__(self, expand_dims=True, dtype=np.float32)
36# relabel, similar to the 2D counterpart
37Relabel.__init__(self, map = [], **kwargs)
38# to one hot label, similar to the 2D counterpart
39ToOneHot.__init__(self, num_classes = None, ignore_background = False)
40# perform Gaussian blur with a certain probability.
41GaussianBlur.__init__(self, sigma=[.1, 2.], execution_probability=0.5)
42# perform binary closing for several iterations
43RemoveSmallGap.__init__(self, iterations)
44# perform binary opening for several iterations
45RemoveThinConnection.__init__(self, iterations)
46# to binary mask, similar to the 2D counterpart
47ToBinaryMask.__init__(self, threshold=0)
Point Cloud Augmentation
We adopt and implement some common augmentations for point clouds.
1# Numpy to tensor
2ToTensor.__init__(self, dtype = 'float')
3# Centers and normalizes node positions to the interval :math:`(-1, 1)`.
4# method should be one of ['minmax', 'mean']
5Normalize.__init__(self, method = 'minmax')
6# Samples a fixed number of points and features from a point cloud.
7FixedPoints.__init__(self, num, replace=True)
8# Transforms node positions with a square transformation matrix computed offline.
9LinearTransformation.__init__(self, matrix)
10# Rotates node positions around a specific axis by a randomly sampled factor within a given interval.
11# Args:
12# degrees (tuple or float): rotation degree
13# axis (int, optional): The rotation axis. (default: :obj:`0`)
14Rotate.__init__(self, degree, axis=0)
15# Rotates node positions around a specific axis by a randomly sampled factor within a given interval.
16# Args:
17# degrees (tuple or float): Rotation interval from which the rotation
18# angle is sampled. If :obj:`degrees` is a number instead of a
19# tuple, the interval is given by :math:`[-\mathrm{degrees},
20# \mathrm{degrees}]`.
21# axis (int, optional): The rotation axis. (default: :obj:`0`)
22RandomRotate.__init__(self, degrees, axis=0)
23# Add gaussian noise to points
24AddNoise.__init__(self, std=0.01)
25# Add gaussian noise with random std to points
26AddRandomNoise.__init__(self, std_range=[0, 0.10])
27# Scales node positions by a randomly sampled factor s within given interval, *e.g.*, resulting in the transformation matrix
28RandomScale.__init__(self, scales)
29# Translates node positions by randomly sampled translation values within a given interval.
30# In contrast to other random transformations, translation is applied separately at each position.
31# Args:
32# translate (sequence or float or int): Maximum translation in each
33# dimension, defining the range
34# :math:`(-\mathrm{translate}, +\mathrm{translate})` to sample from.
35# If :obj:`translate` is a number instead of a sequence, the same
36# range is used for each dimension.
37RandomTranslate.__init__(self, translate)
38# Shuffle order of points in point cloud
39ShufflePoints.__init__(self)
40# reorder points by a specified axis
41ReorderByAxis.__init__(self, axis=0)
42# reorder points by Hilbert curve
43ReorderByHilbert.__init__(self, bins = 16, radius = 1.0, origin = (0,0,0))
44# To one hot label, background value should be 0
45ToOneHot.__init__(self, num_classes = None, ignore_background = False, **kwargs)