DCT-UNet/data_loader.py at main · sunyoungIT/DCT-UNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
"""
Data load with pre-processing
"""

import monai.transforms as mt
from copy import deepcopy


def pre_transforms(crop_size, resample_spacing):

    train_transform = mt.Compose(
        [
            mt.LoadImaged(keys=["image", "label"]), # load image
            mt.EnsureChannelFirstd(keys=["image", "label"]),
            mt.Orientationd(keys=["image", "label"], axcodes="RAS"), # adjust orientation (RAS)
            mt.Spacingd(keys=["image", "label"], pixdim=(resample_spacing[0], resample_spacing[1],resample_spacing[2]), mode=("bilinear", "nearest")),
            mt.ScaleIntensityRangePercentilesd(keys="image",lower=10, upper=90, b_min=0, b_max=1, relative=True, channel_wise=True),
            mt.RandCropByPosNegLabeld(keys=["image", "label"], label_key = "label", spatial_size=crop_size, pos=1.0, neg=0.0, num_samples=1),
            mt.ToTensord(keys=["image", "label"]),
        ]
    )

    val_transform = mt.Compose(
        [
            mt.LoadImaged(keys=["image", "label"]),
            mt.EnsureChannelFirstd(keys=["image", "label"]),
            mt.Orientationd(keys=["image", "label"], axcodes="RAS"),
            mt.Spacingd(keys=["image", "label"], pixdim=(resample_spacing[0], resample_spacing[1],resample_spacing[2]), mode=("bilinear", "nearest")),
            mt.ScaleIntensityRangePercentilesd(keys="image",lower=10, upper=90, b_min=0, b_max=1, relative=True, channel_wise=True),
            mt.RandCropByPosNegLabeld(keys=["image", "label"], label_key = "label", spatial_size=crop_size, pos=1.0, neg=0.0, num_samples=1),
            mt.ToTensord(keys=["image", "label"]),
        ]
    )

    test_transform = mt.Compose(
        [
            mt.LoadImaged(keys=["image", "label"]),
            mt.EnsureChannelFirstd(keys=["image", "label"]),
            mt.Orientationd(keys=["image", "label"], axcodes="RAS"),
            mt.Spacingd(keys=["image", "label"], pixdim=(resample_spacing[0], resample_spacing[1],resample_spacing[2]), mode=("bilinear", "nearest")),
            mt.ScaleIntensityRangePercentilesd(keys="image",lower=10, upper=90, b_min=0, b_max=1, relative=True, channel_wise=True),
            mt.ToTensord(keys=["image", "label"]),
        ]
    )

    return train_transform, val_transform, test_transform