def byol_aug(filename):
"""
BYOL minimizes the distance between representations of each sample and a transformation of that sample.
Examples of transformations include: translation, rotation, blurring, color inversion, color jitter, gaussian noise.
Return an augmented dataset that consisted the above mentioned transformation. Will be used in the training.
"""
image = tio.ScalarImage(filename)
get_foreground = tio.ZNormalization.mean
training_transform = tio.Compose([
tio.CropOrPad((180, 220, 170)), # zero mean, unit variance of foreground
tio.ZNormalization(
masking_method=get_foreground),
tio.RandomBlur(p=0.25), # blur 25% of times
tio.RandomNoise(p=0.25), # Gaussian noise 25% of times
tio.OneOf({ # either
tio.RandomAffine(): 0.8, # random affine
tio.RandomElasticDeformation(): 0.2, # or random elastic deformation
}, p=0.8), # applied to 80% of images
tio.RandomBiasField(p=0.3), # magnetic field inhomogeneity 30% of times
tio.OneOf({ # either
tio.RandomMotion(): 1, # random motion artifact
tio.RandomSpike(): 2, # or spikes
tio.RandomGhosting(): 2, # or ghosts
}, p=0.5), # applied to 50% of images
])
tfs_image = training_transform(image)
return tfs_image
def __getitem__(self, idx):
# Generate one batch of data
# ScalarImage expect 4DTensor, so add a singleton dimension
image = self.CT_partition[idx].unsqueeze(0)
mask = self.mask_partition[idx].unsqueeze(0)
if self.augment:
aug = tio.Compose([tio.OneOf\
({tio.RandomAffine(scales= (0.9, 1.1, 0.9, 1.1, 1, 1),
degrees= (5.0, 5.0, 0)): 0.35,
tio.RandomElasticDeformation(num_control_points=9,
max_displacement= (0.1, 0.1, 0.1),
locked_borders= 2,
image_interpolation= 'linear'): 0.35,
tio.RandomFlip(axes=(2,)):.3}),
])
subject = tio.Subject(ct=tio.ScalarImage(tensor=image),
mask=tio.ScalarImage(tensor=mask))
output = aug(subject)
augmented_image = output['ct']
augmented_mask = output['mask']
image = augmented_image.data
mask = augmented_mask.data
# note that mask is integer
mask = mask.type(torch.IntTensor)
image = image.type(torch.FloatTensor)
#The tensor we pass into ScalarImage is C x W x H x D, so permute axes to
# C x D x H x W. At the end we have N x 1 x D x H x W.
image = image.permute(0, 3, 2, 1)
mask = mask.permute(0, 3, 2, 1)
# Return image and mask pair tensors
return image, mask
def get_train_transform(landmarks_path, resection_params=None):
spatial_transform = tio.Compose((
tio.OneOf({
tio.RandomAffine(): 0.9,
tio.RandomElasticDeformation(): 0.1,
}),
tio.RandomFlip(),
))
resolution_transform = tio.OneOf(
(
tio.RandomAnisotropy(),
tio.RandomBlur(),
),
p=0.75,
)
transforms = []
if resection_params is not None:
transforms.append(get_simulation_transform(resection_params))
if landmarks_path is not None:
transforms.append(
tio.HistogramStandardization({'image': landmarks_path}))
transforms.extend([
# tio.RandomGamma(p=0.2),
resolution_transform,
tio.RandomGhosting(p=0.2),
tio.RandomSpike(p=0.2),
tio.RandomMotion(p=0.2),
tio.RandomBiasField(p=0.5),
tio.ZNormalization(masking_method=tio.ZNormalization.mean),
tio.RandomNoise(p=0.75), # always after ZNorm and after blur!
spatial_transform,
get_tight_crop(),
])
return tio.Compose(transforms)
def test_transforms(self):
landmarks_dict = dict(
t1=np.linspace(0, 100, 13),
t2=np.linspace(0, 100, 13),
)
elastic = torchio.RandomElasticDeformation(max_displacement=1)
transforms = (
torchio.CropOrPad((9, 21, 30)),
torchio.ToCanonical(),
torchio.Resample((1, 1.1, 1.25)),
torchio.RandomFlip(axes=(0, 1, 2), flip_probability=1),
torchio.RandomMotion(),
torchio.RandomGhosting(axes=(0, 1, 2)),
torchio.RandomSpike(),
torchio.RandomNoise(),
torchio.RandomBlur(),
torchio.RandomSwap(patch_size=2, num_iterations=5),
torchio.Lambda(lambda x: 2 * x, types_to_apply=torchio.INTENSITY),
torchio.RandomBiasField(),
torchio.RescaleIntensity((0, 1)),
torchio.ZNormalization(masking_method='label'),
torchio.HistogramStandardization(landmarks_dict=landmarks_dict),
elastic,
torchio.RandomAffine(),
torchio.OneOf({
torchio.RandomAffine(): 3,
elastic: 1
}),
torchio.Pad((1, 2, 3, 0, 5, 6), padding_mode='constant', fill=3),
torchio.Crop((3, 2, 8, 0, 1, 4)),
)
transform = torchio.Compose(transforms)
transform(self.sample)
def get_heterogeneous_dataset(self):
# Keys missing in one of the samples will not be present in the batch
# This is relevant for the case in which a transform is applied to some
# samples only, according to its probability (p argument)
transform_no = tio.RandomElasticDeformation(p=0, max_displacement=1)
transform_yes = tio.RandomElasticDeformation(p=1, max_displacement=1)
sample_no = transform_no(self.sample_subject)
sample_yes = transform_yes(self.sample_subject)
data = sample_no, sample_yes
class Dataset:
def __init__(self, data):
self.data = data
def __len__(self):
return len(self.data)
def __getitem__(self, index):
return self.data[index]
return Dataset(data)
def get_transform(augmentation, landmarks_path):
import datasets
import torchio as tio
if augmentation:
return datasets.get_train_transform(landmarks_path)
else:
preprocess = datasets.get_test_transform(landmarks_path)
augment = tio.Compose((tio.RandomFlip(),
tio.OneOf({
tio.RandomAffine(): 0.8,
tio.RandomElasticDeformation(): 0.2,
})))
return tio.Compose((preprocess, augment))
def get_transform(self, channels, is_3d=True, labels=True):
landmarks_dict = {
channel: np.linspace(0, 100, 13)
for channel in channels
}
disp = 1 if is_3d else (1, 1, 0.01)
elastic = tio.RandomElasticDeformation(max_displacement=disp)
cp_args = (9, 21, 30) if is_3d else (21, 30, 1)
resize_args = (10, 20, 30) if is_3d else (10, 20, 1)
flip_axes = axes_downsample = (0, 1, 2) if is_3d else (0, 1)
swap_patch = (2, 3, 4) if is_3d else (3, 4, 1)
pad_args = (1, 2, 3, 0, 5, 6) if is_3d else (0, 0, 3, 0, 5, 6)
crop_args = (3, 2, 8, 0, 1, 4) if is_3d else (0, 0, 8, 0, 1, 4)
remapping = {1: 2, 2: 1, 3: 20, 4: 25}
transforms = [
tio.CropOrPad(cp_args),
tio.EnsureShapeMultiple(2, method='crop'),
tio.Resize(resize_args),
tio.ToCanonical(),
tio.RandomAnisotropy(downsampling=(1.75, 2), axes=axes_downsample),
tio.CopyAffine(channels[0]),
tio.Resample((1, 1.1, 1.25)),
tio.RandomFlip(axes=flip_axes, flip_probability=1),
tio.RandomMotion(),
tio.RandomGhosting(axes=(0, 1, 2)),
tio.RandomSpike(),
tio.RandomNoise(),
tio.RandomBlur(),
tio.RandomSwap(patch_size=swap_patch, num_iterations=5),
tio.Lambda(lambda x: 2 * x, types_to_apply=tio.INTENSITY),
tio.RandomBiasField(),
tio.RescaleIntensity(out_min_max=(0, 1)),
tio.ZNormalization(),
tio.HistogramStandardization(landmarks_dict),
elastic,
tio.RandomAffine(),
tio.OneOf({
tio.RandomAffine(): 3,
elastic: 1,
}),
tio.RemapLabels(remapping=remapping, masking_method='Left'),
tio.RemoveLabels([1, 3]),
tio.SequentialLabels(),
tio.Pad(pad_args, padding_mode=3),
tio.Crop(crop_args),
]
if labels:
transforms.append(tio.RandomLabelsToImage(label_key='label'))
return tio.Compose(transforms)
def get_context(device, variables, augmentation_mode, **kwargs):
context = base_config.get_context(device, variables, **kwargs)
context.file_paths.append(os.path.abspath(__file__))
context.config.update({'augmentation_mode': augmentation_mode})
# training_transform is a tio.Compose where the second transform is the augmentation
dataset_defn = context.get_component_definition("dataset")
training_transform = dataset_defn['params']['transforms']['training']
dwi_augmentation = ReconstructMeanDWI(num_dwis=(1, 7),
num_directions=(1, 3),
directionality=(4, 10))
noise = tio.RandomNoise(std=0.035, p=0.3)
blur = tio.RandomBlur((0, 1), p=0.2)
standard_augmentations = tio.Compose([
tio.RandomFlip(axes=(0, 1, 2)),
tio.RandomElasticDeformation(p=0.5,
num_control_points=(7, 7, 4),
locked_borders=1,
image_interpolation='bspline',
exclude="full_dwi"),
tio.RandomBiasField(p=0.5),
tio.RescaleIntensity((0, 1), (0.01, 99.9)),
tio.RandomGamma(p=0.8),
tio.RescaleIntensity((-1, 1)),
tio.OneOf([
tio.Compose([blur, noise]),
tio.Compose([noise, blur]),
])
],
exclude="full_dwi")
if augmentation_mode == 'no_augmentation':
training_transform.transforms.pop(1)
elif augmentation_mode == 'standard':
training_transform.transforms[1] = standard_augmentations
elif augmentation_mode == 'dwi_reconstruction':
training_transform.transforms[1] = dwi_augmentation
elif augmentation_mode == 'combined':
training_transform.transforms[1] = tio.Compose(
[dwi_augmentation, standard_augmentations])
else:
raise ValueError(f"Invalid augmentation mode {augmentation_mode}")
return context
def get_transform(self, channels, is_3d=True, labels=True):
landmarks_dict = {
channel: np.linspace(0, 100, 13)
for channel in channels
}
disp = 1 if is_3d else (1, 1, 0.01)
elastic = torchio.RandomElasticDeformation(max_displacement=disp)
cp_args = (9, 21, 30) if is_3d else (21, 30, 1)
flip_axes = axes_downsample = (0, 1, 2) if is_3d else (0, 1)
swap_patch = (2, 3, 4) if is_3d else (3, 4, 1)
pad_args = (1, 2, 3, 0, 5, 6) if is_3d else (0, 0, 3, 0, 5, 6)
crop_args = (3, 2, 8, 0, 1, 4) if is_3d else (0, 0, 8, 0, 1, 4)
transforms = [
torchio.CropOrPad(cp_args),
torchio.ToCanonical(),
torchio.RandomDownsample(downsampling=(1.75, 2),
axes=axes_downsample),
torchio.Resample((1, 1.1, 1.25)),
torchio.RandomFlip(axes=flip_axes, flip_probability=1),
torchio.RandomMotion(),
torchio.RandomGhosting(axes=(0, 1, 2)),
torchio.RandomSpike(),
torchio.RandomNoise(),
torchio.RandomBlur(),
torchio.RandomSwap(patch_size=swap_patch, num_iterations=5),
torchio.Lambda(lambda x: 2 * x, types_to_apply=torchio.INTENSITY),
torchio.RandomBiasField(),
torchio.RescaleIntensity((0, 1)),
torchio.ZNormalization(),
torchio.HistogramStandardization(landmarks_dict),
elastic,
torchio.RandomAffine(),
torchio.OneOf({
torchio.RandomAffine(): 3,
elastic: 1,
}),
torchio.Pad(pad_args, padding_mode=3),
torchio.Crop(crop_args),
]
if labels:
transforms.append(torchio.RandomLabelsToImage(label_key='label'))
return torchio.Compose(transforms)
def __init__(self):
# ID and Name
self.id = "506b95"
self.experiment_name = "ma_crosstr_v{}".format(self.id)
self.debug = False
# System
self.checkpointsBasePath = "./checkpoints/"
self.checkpointsBasePathMod = self.checkpointsBasePath + 'models/'
self.labelpath = '/local/DEEPLEARNING/MULTI_ATLAS/MULTI_ATLAS/nnUNet_preprocessed/Task017_BCV/nnUNetData_plans_v2.1_stage1/'
self.datapath = self.labelpath
self.input_shape = [512,512,256]
# self.filters = [16, 32, 64, 128]
# self.filters = [64, 192, 448, 704]
# self.filters = [16, 32, 64, 128, 256]
self.filters = [32, 64, 128, 256, 512]
d_model = self.filters[-1]
# skip_idx = [1,3,5,6]
# self.patch_size=(128,128,128)
self.patch_size=(192,192,48)
# n_layers=6
self.clip = False
self.patched = True
# GPU
self.gpu = '1'
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu
# torch.backends.cudnn.benchmark = False
# Model
number_of_cross_heads = 8
number_of_self_heads = 8
number_of_self_layer = 6
self.n_classes = 14
self.net = CrossPatch3DTr(filters=self.filters,patch_size=[1,1,1],
d_model=d_model,n_classes=self.n_classes,
n_cheads=1,n_sheads=number_of_self_heads,
bn=True,up_mode='deconv',
n_strans=number_of_self_layer, do_cross=True,
enc_grad=False)
self.net.inference_apply_nonlin = softmax_helper
self.n_parameters = count_parameters(self.net)
print("N PARAMS : {}".format(self.n_parameters))
# self.model_path = './checkpoints/models/deep_crosstr.pth'
self.model_path = './checkpoints/models/506/modlast.pt'
max_displacement = 5,5,5
deg = (0,5,10)
scales = 0
self.transform = tio.Compose([
tio.RandomElasticDeformation(max_displacement=max_displacement),
tio.RandomAffine(scales=scales, degrees=deg)
])
# Training
self.start_epoch = 1000
self.epoch = 2000
# self.loss = torch.nn.CrossEntropyLoss()
self.loss = DC_and_CE_loss({'batch_dice': True, 'smooth': 1e-5, 'do_bg': False}, {})
self.ds_scales = ((1, 1, 1), (0.5, 0.5, 0.5), (0.25, 0.25, 0.25), (0.125,0.125,0.125))
################# Here we wrap the loss for deep supervision ############
# we need to know the number of outputs of the network
net_numpool = 4
# we give each output a weight which decreases exponentially (division by 2) as the resolution decreases
# this gives higher resolution outputs more weight in the loss
weights = np.array([1 / (2 ** i) for i in range(net_numpool)])
# we don't use the lowest 2 outputs. Normalize weights so that they sum to 1
mask = np.array([True] + [True if i < net_numpool - 1 else False for i in range(1, net_numpool)])
weights[~mask] = 0
weights = weights / weights.sum()
self.ds_loss_weights = weights
# now wrap the loss
self.loss = MultipleOutputLoss2(self.loss, self.ds_loss_weights)
################# END ###################
self.batchsize = 2
self.lr_rate = 1e-3
self.load_lr = False
self.load_model()
self.net.reinit_decoder()
self.net.reinit_crostrans(dim=d_model, depth=1, heads=number_of_cross_heads, dim_head=1024, mlp_dim=1024, dropout = 0.1)
self.optimizer = optim.SGD(self.net.parameters(), lr = self.lr_rate, weight_decay=3e-5, momentum=0.99, nesterov=True)
self.optimizer.zero_grad()
self.validate_every_k_epochs = 1
# self.decay = (self.lr_rate/self.final_lr_rate - 1)/self.epoch
self.lr_scheduler = get_scheduler(self.optimizer, "poly", self.lr_rate, max_epochs=self.epoch)
# Other
self.classes_name = ['background','spleen','right kidney','left kidney','gallbladder','esophagus','liver','stomach','aorta','inferior vena cava','portal vein and splenic vein','pancreas','right adrenal gland','left adrenal gland']
def __init__(self):
# ID and Name
self.id = 601
self.experiment_name = "ma_cotr_v{}".format(self.id)
self.debug = False
# System
self.checkpointsBasePath = "./checkpoints/"
self.checkpointsBasePathMod = self.checkpointsBasePath + 'models/'
# self.labelpath = "/local/DEEPLEARNING/MULTI_ATLAS/multi_atlas//512_512_256/"
self.labelpath = '/local/DEEPLEARNING/VP_multiorgan_v2/'
self.datapath = self.labelpath
self.input_shape = [512,512,256]
# filters = [4, 8, 16, 32]
# skip_idx = [1,3,5,6]
# self.patch_size=(128,128,128)
self.patch_size=(192,192,48)
# n_layers=6
self.clip = True
self.patched = True
# GPU
self.gpu = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu
# Model
self.n_classes = 8
self.net = ResTranUnet(norm_cfg='IN', activation_cfg='LeakyReLU', img_size=self.patch_size, num_classes=self.n_classes, weight_std=False, deep_supervision=True)
self.net.inference_apply_nonlin = softmax_helper
self.n_parameters = count_parameters(self.net)
print("N PARAMS : {}".format(self.n_parameters))
self.model_path = './checkpoints/models/vp_cotr.pth'
# self.model_path = './checkpoints/models/403/mod.pt'
max_displacement = 5,5,5
deg = (0,5,10)
scales = 0
self.transform = tio.Compose([
tio.RandomElasticDeformation(max_displacement=max_displacement),
tio.RandomAffine(scales=scales, degrees=deg)
])
# Training
self.start_epoch = 0
self.epoch = 1000
# self.loss = torch.nn.CrossEntropyLoss()
self.loss = DC_and_CE_loss({'batch_dice': True, 'smooth': 1e-5, 'do_bg': False}, {})
self.ds_scales = ((1, 1, 1), (0.5, 0.5, 1), (0.25, 0.25, 0.5))
################# Here we wrap the loss for deep supervision ############
# we need to know the number of outputs of the network
net_numpool = 4
# we give each output a weight which decreases exponentially (division by 2) as the resolution decreases
# this gives higher resolution outputs more weight in the loss
weights = np.array([1 / (2 ** i) for i in range(net_numpool)])
# we don't use the lowest 2 outputs. Normalize weights so that they sum to 1
mask = np.array([True] + [True if i < net_numpool - 1 else False for i in range(1, net_numpool)])
weights[~mask] = 0
weights = weights / weights.sum()
self.ds_loss_weights = weights
# now wrap the loss
self.loss = MultipleOutputLoss2(self.loss, self.ds_loss_weights)
################# END ###################
self.batchsize = 2
self.lr_rate = 2e-2
# self.final_lr_rate = 1e-5
# self.optimizer = optim.Adam(self.net.parameters(), lr = self.lr_rate)
self.optimizer = optim.SGD(self.net.parameters(), lr = self.lr_rate, weight_decay=3e-5, momentum=0.99)
self.optimizer.zero_grad()
self.validate_every_k_epochs = 10
# self.decay = (self.lr_rate/self.final_lr_rate - 1)/self.epoch
self.lr_scheduler = get_scheduler(self.optimizer, "poly", self.lr_rate, max_epochs=self.epoch)
self.load_model()
# Other
self.classes_name = ["Background", "Liver","Gallbladder","Spleen","Left_Kidney","Right_Kidney","Pancreas","Stomach"]
def __init__(self):
# ID and Name
self.id = 100
self.experiment_name = "multi_atlas_unet_016_e1000_CE_adam_wd6_da_id{}".format(self.id)
self.debug = False
# System
self.checkpointsBasePath = "./checkpoints/"
self.checkpointsBasePathMod = self.checkpointsBasePath + 'models/'
# self.labelpath = "/local/SSD_DEEPLEARNING/MULTI_ATLAS/multi_atlas/data_3D_size_512_512_198_res_1.0_1.0_1.0.hdf5"
# self.labelpath = "/local/SSD_DEEPLEARNING/MULTI_ATLAS/multi_atlas/data_3D_size_256_256_99_res_0.5_0.5.hdf5"
self.labelpath = "/local/SSD_DEEPLEARNING/MULTI_ATLAS/multi_atlas/data_3D_size_80_80_32_res_0.16.hdf5"
self.datapath = "/local/SSD_DEEPLEARNING/MULTI_ATLAS/multi_atlas/data_3D_size_80_80_32_res_0.16.hdf5"
# GPU
self.gpu = '1'
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu
# Model
self.channels = [64, 128, 256, 512, 1024]
self.channels = [int(x) for x in self.channels]
self.net = unet_3D(self.channels, n_classes=14, is_batchnorm=False, in_channels=1, interpolation = None)#1, self.channels, 12, interpolation = (512,512,198))
# self.net = RevUnet3D(1, self.channels, 12, interpolation = (256,256,99))
self.n_parameters = count_parameters(self.net)
print("N PARAMS : {}".format(self.n_parameters))
self.n_classes = 14
max_displacement = 5,5,5
deg = (0,5,10)
scales = 0
self.transform = tio.Compose([
tio.RandomElasticDeformation(max_displacement=max_displacement),
tio.RandomAffine(scales=scales, degrees=deg)
])
# Training
self.train_original_classes = False
self.epoch = 1000
# def loss(outputs, labels):
# return atlasUtils.atlasDiceLoss(outputs, labels, n_classe = self.n_classes)
# self.loss = loss
# self.loss = SoftDiceLoss(self.n_classes)
self.loss = torch.nn.CrossEntropyLoss()
self.hot = 0
self.batchsize = 1
# self.optimizer = optim.Ada(self.net.parameters(),
# lr= 0.01, #to do
# momentum=0.9,
# nesterov=True,
# weight_decay=1e-5) #todo
self.lr_rate = 5e-4
self.optimizer = optim.Adam(self.net.parameters(), lr = self.lr_rate, weight_decay=1e-6)
# self.optimizer = optim.SGD(self.net.parameters(),
# lr=self.lr_rate)
self.optimizer.zero_grad()
self.validate_every_k_epochs = 1
# Scheduler list : [lambdarule_1]
# self.lr_scheduler = get_scheduler(self.optimizer, "multistep")
self.lr_scheduler = get_scheduler(self.optimizer, "multistep", self.lr_rate)
# self.lr_scheduler = get_scheduler(self.optimizer, "lambdarule_1", self.lr_rate)
# Other
self.classes_name = ['background','spleen','right kidney','left kidney','gallbladder','esophagus','liver','stomach','aorta','inferior vena cava','portal vein and splenic vein','pancreas','right adrenal gland','left adrenal gland']
self.look_small = False
def __init__(self):
# ID and Name
self.id = -1
self.experiment_name = "ma_cotr_pred_v{}".format(self.id)
self.debug = False
# System
self.checkpointsBasePath = "./checkpoints/"
self.checkpointsBasePathMod = self.checkpointsBasePath + 'models/'
self.labelpath = "/local/DEEPLEARNING/MULTI_ATLAS/multi_atlas//512_512_256/"
self.datapath = self.labelpath
self.input_shape = [512,512,256]
# filters = [4, 8, 16, 32]
# skip_idx = [1,3,5,6]
self.patch_size=(128,128,128)
# self.patch_size=(192,192,48)
# n_layers=6
self.clip = True
self.patched = True
# GPU
self.gpu = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu
# Model
self.n_classes = 14
self.net = ResTranUnet(norm_cfg='IN', activation_cfg='LeakyReLU', img_size=self.patch_size, num_classes=self.n_classes, weight_std=False, deep_supervision=False)
self.net.inference_apply_nonlin = softmax_helper
self.n_parameters = count_parameters(self.net)
print("N PARAMS : {}".format(self.n_parameters))
# self.model_path = './checkpoints/models/cotr.pth'
self.model_path = './checkpoints/models/400/mod.pt'
max_displacement = 5,5,5
deg = (0,5,10)
scales = 0
self.transform = tio.Compose([
tio.RandomElasticDeformation(max_displacement=max_displacement),
tio.RandomAffine(scales=scales, degrees=deg)
])
# Training
self.start_epoch = 1000
self.epoch = 1000
self.loss = torch.nn.CrossEntropyLoss()
self.loss = DC_and_CE_loss({'batch_dice': True, 'smooth': 1e-5, 'do_bg': False}, {})
self.batchsize = 2
self.lr_rate = 2e-2
# self.final_lr_rate = 1e-5
# self.optimizer = optim.Adam(self.net.parameters(), lr = self.lr_rate)
self.optimizer = optim.SGD(self.net.parameters(), lr = self.lr_rate, weight_decay=3e-5, momentum=0.99)
self.optimizer.zero_grad()
self.validate_every_k_epochs = 10
# self.decay = (self.lr_rate/self.final_lr_rate - 1)/self.epoch
self.lr_scheduler = get_scheduler(self.optimizer, "poly", self.lr_rate, max_epochs=self.epoch)
self.load_model()
# Other
self.classes_name = ['background','spleen','right kidney','left kidney','gallbladder','esophagus','liver','stomach','aorta','inferior vena cava','portal vein and splenic vein','pancreas','right adrenal gland','left adrenal gland']
def __init__(self):
# ID and Name
self.id = 301
self.experiment_name = "pancreas_unetr_v{}".format(self.id)
self.debug = False
# System
self.checkpointsBasePath = "./checkpoints/"
self.checkpointsBasePathMod = self.checkpointsBasePath + 'models/'
self.labelpath = "/local/DEEPLEARNING/PANCREAS_MULTI_RES/512_512_256/"
self.datapath = self.labelpath
self.input_shape = [512,512,256]
filters = [64, 128, 256, 512]
skip_idx = [3,6,9,12]
patch_size=(16,16,16)
self.data_patch = [128,128,128]
n_layers=12
self.patched = True
# GPU
self.gpu = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu
# Model
self.n_classes = 2
self.net = UNETR(input_shape=self.data_patch,filters=filters,patch_size=patch_size, n_layers=n_layers, skip_idx=skip_idx)
self.n_parameters = count_parameters(self.net)
print("N PARAMS : {}".format(self.n_parameters))
max_displacement = 5,5,5
deg = (0,5,10)
scales = 0
self.transform = tio.Compose([
tio.RandomElasticDeformation(max_displacement=max_displacement),
tio.RandomAffine(scales=scales, degrees=deg)
])
# Training
self.start_epoch = 0
self.train_original_classes = False
self.epoch = 25
# self.model_path = './checkpoints/models/unetr.pth'
self.model_path = './checkpoints/models/unetr.pth'
self.load_model()
self.split = 1
self.loss = torch.nn.CrossEntropyLoss()
self.hot = 0
self.batchsize = 1
self.lr_rate = 1e-4
self.final_lr_rate = 1e-5
self.optimizer = optim.Adam(self.net.parameters(), lr = self.lr_rate)
self.optimizer.zero_grad()
self.validate_every_k_epochs = 1
self.decay = (self.lr_rate/self.final_lr_rate - 1)/self.epoch
self.lr_scheduler = get_scheduler(self.optimizer, "po", self.lr_rate, self.decay)
# Other
self.classes_name = ['background','pancreas']
self.look_small = False
#corresponds to the lateral axis. In TorchIO, you can use anatomical
#labels instead, so that you don't need to figure out image orientation
#to know which axis you would like to flip. To make sure the transform modifies
#the image, we will use the inferior-superior (longitudinal) axis and a flip
#probability of 1. If the flipping happened along any other axis, we might not
#notice it using this visualization.
random_flip = tio.RandomFlip(axes=['inferior-superior'], flip_probability=1)
fpg_flipped = random_flip(fpg_ras)
show_fpg(fpg_flipped)
#Random elastic deformation To simulate anatomical variations in our images,
#we can apply a non-linear deformation using RandomElasticDeformation.
max_displacement = 10, 10, 0 # in x, y and z directions
random_elastic = tio.RandomElasticDeformation(max_displacement=max_displacement, seed=0)
slice_elastic = random_elastic(slice_grid)
to_pil(slice_elastic)
#As explained in the documentation, one can change the number of grid control
#points to set the deformation smoothness.
random_elastic = tio.RandomElasticDeformation(
max_displacement=max_displacement,
num_control_points=7,
)
slice_large_displacement_more_cp = random_elastic(slice_grid)
to_pil(slice_large_displacement_more_cp)
#Let's look at the effect of using few control points
#but very large displacements.
sub = tio.datasets.Colin27()
sub.pop('brain')
sub.pop('head')
t = tio.Compose([
tio.Pad(padding=10, padding_mode="reflect"),
tio.Crop(bounds_parameters=10)
])
new_sub = t(sub)
new_sub.t1.affine
sub.t1.affine
#test inverse elastic
import torchio as tio
import SimpleITK as sitk
colin = tio.datasets.Colin27()
transform = tio.RandomElasticDeformation()
transformed = transform(colin)
trsfm_hist, seeds_hist = tio.compose_from_history(history=transformed.history)
trsfm_hist[0].get_inverse = True
colin_back = trsfm_hist[0](transformed, seed=seeds_hist[0])
transformed.t1.save('/tmp/tcolin.nii')
colin_back.t1.save('/tmp/t_back_colin.nii')
colin.t1.save('/tmp/colin.nii')
grid = -transformed.history[0][1]['coarse_grid']
def get_context(device, variables, fold=0, **kwargs):
context = TorchContext(device, name="msseg2", variables=variables)
context.file_paths.append(os.path.abspath(__file__))
context.config = config = {'fold': fold, 'patch_size': 96}
input_images = ["flair_time01", "flair_time02"]
subject_loader = ComposeLoaders([
ImageLoader(glob_pattern="flair_time01*",
image_name='flair_time01',
image_constructor=tio.ScalarImage),
ImageLoader(glob_pattern="flair_time02*",
image_name='flair_time02',
image_constructor=tio.ScalarImage),
ImageLoader(glob_pattern="brain_mask.*",
image_name='brain_mask',
image_constructor=tio.LabelMap,
label_values={"brain": 1}),
ImageLoader(glob_pattern="ground_truth.*",
image_name="ground_truth",
image_constructor=tio.LabelMap,
label_values={"lesion": 1}),
])
cohorts = {}
cohorts['all'] = RequireAttributes(input_images)
cohorts['validation'] = RandomFoldFilter(num_folds=5,
selection=fold,
seed=0xDEADBEEF)
cohorts['training'] = NegateFilter(cohorts['validation'])
common_transforms_1 = tio.Compose([
SetDataType(torch.float),
EnforceConsistentAffine(source_image_name='flair_time01'),
TargetResample(target_spacing=1, tolerance=0.11),
CropToMask('brain_mask'),
MinSizePad(config['patch_size'])
])
augmentations = tio.Compose([
RandomPermuteDimensions(),
tio.RandomFlip(axes=(0, 1, 2)),
tio.OneOf(
{
tio.RandomElasticDeformation():
0.2,
tio.RandomAffine(scales=0.2,
degrees=45,
default_pad_value='otsu'):
0.8,
},
p=0.75),
tio.RandomBiasField(p=0.5),
tio.RescaleIntensity((0, 1), (0.01, 99.9)),
tio.RandomGamma(p=0.8),
tio.RescaleIntensity((-1, 1)),
tio.RandomBlur((0, 1), p=0.2),
tio.RandomNoise(std=0.1, p=0.35)
])
common_transforms_2 = tio.Compose([
tio.RescaleIntensity((-1, 1.), (0.05, 99.5)),
ConcatenateImages(image_names=["flair_time01", "flair_time02"],
image_channels=[1, 1],
new_image_name="X"),
RenameProperty(old_name='ground_truth', new_name='y'),
CustomOneHot(include="y"),
])
transforms = {
'default':
tio.Compose([common_transforms_1, common_transforms_2]),
'training':
tio.Compose([
common_transforms_1, augmentations, common_transforms_2,
ImageFromLabels(new_image_name="patch_probability",
label_weights=[('brain_mask', 'brain', 1),
('y', 'lesion', 100)])
]),
}
context.add_component("dataset",
SubjectFolder,
root='$DATASET_PATH',
subject_path="",
subject_loader=subject_loader,
cohorts=cohorts,
transforms=transforms)
context.add_component("model",
ModularUNet,
in_channels=2,
out_channels=2,
filters=[40, 40, 80, 80, 120, 120],
depth=6,
block_params={'residual': True},
downsample_class=BlurConv3d,
downsample_params={
'kernel_size': 3,
'stride': 2,
'padding': 1
},
upsample_class=BlurConvTranspose3d,
upsample_params={
'kernel_size': 3,
'stride': 2,
'padding': 1,
'output_padding': 0
})
context.add_component("optimizer",
SGD,
params="self.model.parameters()",
lr=0.001,
momentum=0.95)
context.add_component("criterion",
HybridLogisticDiceLoss,
logistic_class_weights=[1, 100])
training_evaluators = [
ScheduledEvaluation(evaluator=SegmentationEvaluator(
'y_pred_eval', 'y_eval'),
log_name='training_segmentation_eval',
interval=15),
ScheduledEvaluation(evaluator=LabelMapEvaluator('y_pred_eval'),
log_name='training_label_eval',
interval=15),
ScheduledEvaluation(evaluator=ContourImageEvaluator(
"random",
'flair_time02',
'y_pred_eval',
'y_eval',
slice_id=0,
legend=True,
ncol=2,
interesting_slice=True,
split_subjects=False),
log_name=f"contour_image",
interval=15),
]
validation_evaluators = [
ScheduledEvaluation(evaluator=SegmentationEvaluator(
"y_pred_eval", "y_eval"),
log_name="segmentation_eval",
cohorts=["validation"],
interval=50),
ScheduledEvaluation(evaluator=ContourImageEvaluator(
"interesting",
'flair_time02',
'y_pred_eval',
'y_eval',
slice_id=0,
legend=True,
ncol=1,
interesting_slice=True,
split_subjects=True),
log_name=f"contour_image",
cohorts=["validation"],
interval=50),
]
def scoring_function(evaluation_dict):
# Grab the output of the SegmentationEvaluator
seg_eval = evaluation_dict['segmentation_eval']['validation']
# Take mean dice, while accounting for subjects which have no lesions.
# Dice is 0/0 = nan when the model correctly outputs no lesions. This is counted as a score of 1.0.
# Dice is (>0)/0 = posinf when the model incorrectly predicts lesions when there are none.
# This is counted as a score of 0.0.
dice = torch.tensor(seg_eval["subject_stats"]['dice.lesion'])
dice = dice.nan_to_num(nan=1.0, posinf=0.0)
score = dice.mean()
return score
train_predictor = StandardPredict(image_names=['X', 'y'])
validation_predictor = PatchPredict(patch_batch_size=32,
patch_size=config['patch_size'],
patch_overlap=(config['patch_size'] //
8),
padding_mode=None,
overlap_mode='average',
image_names=['X'])
patch_sampler = tio.WeightedSampler(patch_size=config['patch_size'],
probability_map='patch_probability')
train_dataloader_factory = PatchDataLoader(max_length=100,
samples_per_volume=1,
sampler=patch_sampler)
validation_dataloader_factory = StandardDataLoader(
sampler=SequentialSampler)
context.add_component(
"trainer",
SegmentationTrainer,
training_batch_size=4,
save_rate=100,
scoring_interval=50,
scoring_function=scoring_function,
one_time_evaluators=[],
training_evaluators=training_evaluators,
validation_evaluators=validation_evaluators,
max_iterations_with_no_improvement=2000,
train_predictor=train_predictor,
validation_predictor=validation_predictor,
train_dataloader_factory=train_dataloader_factory,
validation_dataloader_factory=validation_dataloader_factory)
return context
def dataloader(handles, mode = 'train'):
# If pickle exists, load it
try:
with open('../inputs/flpickles/' + mode + '.pickle', 'rb') as f:
images = pickle.load(f)
except:
images = {}
images['Image'] = []
images['Label'] = []
images['Gap'] = []
images['ID'] = []
# Data augmentations
random_flip = tio.RandomFlip(axes=1)
random_flip2 = tio.RandomFlip(axes=2)
random_affine = tio.RandomAffine(seed=0, scales=(3, 3))
random_elastic = tio.RandomElasticDeformation(
max_displacement=(0, 20, 40),
num_control_points=20,
seed=0,
)
rescale = tio.RescaleIntensity((-1, 1), percentiles=(1, 99))
standardize_foreground = tio.ZNormalization(masking_method=lambda x: x > x.mean())
blur = tio.RandomBlur(seed=0)
standardize = tio.ZNormalization()
add_noise = tio.RandomNoise(std=0.5, seed=42)
add_spike = tio.RandomSpike(seed=42)
add_ghosts = tio.RandomGhosting(intensity=1.5, seed=42)
add_motion = tio.RandomMotion(num_transforms=6, image_interpolation='nearest', seed=42)
swap = tio.RandomSwap(patch_size = 7)
# For each image
for idx, row in handles.iterrows():
im_aug = []
lb_aug = []
gap_aug = []
imgs = np.zeros(shape=(1, 1,7,1024, 1024), dtype=np.float32) # change patch shape if necessary
lbs = np.zeros(shape=(1, 1,7,1024, 1024), dtype=np.float32)
gaps = np.zeros(shape=(1, 1,7,1024, 1024), dtype=np.float32)
im = io.imread(row['Image'])
im = im / 255 # Normalization
im = np.expand_dims(im, axis=0)
imgs[0] = im
im_aug.append(imgs)
images['ID'].append(row['ID'])
if mode == 'train':
im_flip1 = random_flip(im)
imgs[0] = im_flip1
im_aug.append(imgs)
im_flip2 = random_flip2(im)
imgs[0] = im_flip2
im_aug.append(imgs)
im_affine = random_affine(im)
imgs[0] = im_affine
im_aug.append(imgs)
im_elastic = random_elastic(im)
imgs[0] = im_elastic
im_aug.append(imgs)
im_rescale = rescale(im)
imgs[0] = im_rescale
im_aug.append(imgs)
im_standard = standardize_foreground(im)
imgs[0] = im_standard
im_aug.append(imgs)
im_blur = blur(im)
imgs[0] = im_blur
im_aug.append(imgs)
im_noisy = add_noise(standardize(im))
imgs[0] = im_noisy
im_aug.append(imgs)
im_spike = add_spike(im)
imgs[0] = im_spike
im_aug.append(imgs)
im_ghost = add_ghosts(im)
imgs[0] = im_ghost
im_aug.append(imgs)
im_motion = add_motion(im)
imgs[0] = im_motion
im_aug.append(imgs)
im_swap = swap(im)
imgs[0] = im_swap
im_aug.append(imgs)
images['Image'].append(np.array(im_aug))
if mode != 'test':
lb = io.imread(row['Label'])
lb = label_converter(lb)
lb = np.expand_dims(lb, axis=0)
lbs[0] = lb
lb_aug.append(lbs)
gap = io.imread(row['Gap'])
gap = np.expand_dims(gap, axis = 0)
gaps[0] = gap
gap_aug.append(gaps)
if mode == 'train':
lb_flip1 = random_flip(lb)
lbs[0] = lb_flip1
lb_aug.append(lbs)
lb_flip2 = random_flip2(lb)
lbs[0] = lb_flip2
lb_aug.append(lbs)
lb_affine = random_affine(lb)
lbs[0] = lb_affine
lb_aug.append(lbs)
lb_elastic = random_elastic(lb)
lbs[0] = lb_elastic
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
lbs[0] = lb
lb_aug.append(lbs)
gap_flip1 = random_flip(gap)
gaps[0] = gap_flip1
gap_aug.append(gaps)
gap_flip2 = random_flip2(gap)
gaps[0] = gap_flip2
gap_aug.append(gaps)
gap_affine = random_affine(gap)
gaps[0] = gap_affine
gap_aug.append(gaps)
gap_elastic = random_elastic(gap)
gaps[0] = gap_elastic
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
gaps[0] = gap
gap_aug.append(gaps)
images['Label'].append(np.array(lb_aug))
images['Gap'].append(np.array(gap_aug))
# Save images
with open("../inputs/flpickles/" + mode + '.pickle', 'wb') as f:
pickle.dump(images, f)
with open('../inputs/flpickles/' + mode + '.pickle', 'rb') as f:
images = pickle.load(f)
return images
def __init__(self):
# ID and Name
self.id = 305
self.experiment_name = "ma_unetr_v{}".format(self.id)
self.debug = False
# System
self.checkpointsBasePath = "./checkpoints/"
self.checkpointsBasePathMod = self.checkpointsBasePath + 'models/'
self.labelpath = "/local/DEEPLEARNING/MULTI_ATLAS/multi_atlas//512_512_256/"
self.datapath = self.labelpath
self.input_shape = [512, 512, 256]
filters = [64, 128, 256, 512]
skip_idx = [3, 6, 9, 12]
self.patch_size = (192, 192, 48)
n_layers = 12
# GPU
self.gpu = '1'
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu
# Model
self.n_classes = 14
self.net = UNETR(input_shape=self.patch_size,
n_classes=self.n_classes,
filters=filters,
patch_size=(16, 16, 16),
n_layers=n_layers,
skip_idx=skip_idx)
self.n_parameters = count_parameters(self.net)
print("N PARAMS : {}".format(self.n_parameters))
self.model_path = './checkpoints/models/ma_unetr.pth'
# self.model_path = './checkpoints/models/300/mod.pth'
max_displacement = 5, 5, 5
deg = (0, 5, 10)
scales = 0
self.transform = tio.Compose([
tio.RandomElasticDeformation(max_displacement=max_displacement),
tio.RandomAffine(scales=scales, degrees=deg)
])
# Training
self.start_epoch = 0
self.epoch = 1000
self.loss = torch.nn.CrossEntropyLoss()
self.batchsize = 2
self.lr_rate = 2e-2
# self.final_lr_rate = 1e-5
# self.optimizer = optim.Adam(self.net.parameters(), lr = self.lr_rate)
self.optimizer = optim.SGD(self.net.parameters(),
lr=self.lr_rate,
weight_decay=3e-5,
momentum=0.99)
self.optimizer.zero_grad()
self.validate_every_k_epochs = 10
# self.decay = (self.lr_rate/self.final_lr_rate - 1)/self.epoch
self.lr_scheduler = get_scheduler(self.optimizer,
"poly",
self.lr_rate,
max_epochs=self.epoch)
self.load_model()
# Other
self.classes_name = [
'background', 'spleen', 'right kidney', 'left kidney',
'gallbladder', 'esophagus', 'liver', 'stomach', 'aorta',
'inferior vena cava', 'portal vein and splenic vein', 'pancreas',
'right adrenal gland', 'left adrenal gland'
]
def __init__(self):
# ID and Name
self.id = 206
self.experiment_name = "tcia_revunet_03_d3_e1000_CE_adam_wd0_da_f1_lr5_gr1_id{}".format(self.id)
self.debug = False
# System
self.checkpointsBasePath = "./checkpoints/"
self.checkpointsBasePathMod = self.checkpointsBasePath + 'models/'
self.labelpath = "/local/SSD_DEEPLEARNING/PANCREAS_MULTI_RES/160_160_64/"
self.datapath = self.labelpath
self.im_dim = (160,160,64)
# GPU
self.gpu = '2'
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu
# Model
self.channels = [64, 128, 256, 512, 1024]
self.channels = [int(x) for x in self.channels]
self.n_classes = 2
self.n_groups = 1
self.net = RevUnet3D(1, self.channels, self.n_classes , depth = 3 ,interpolation = None, groups = self.n_groups)#(512,512,198))
# self.net = RevUnet3D(1, self.channels, 12, interpolation = (256,256,99))
self.n_parameters = count_parameters(self.net)
print("N PARAMS : {}".format(self.n_parameters))
self.model_path = './checkpoints/models/revunet_tcia_160_160_64_d3_gr1.pth'
self.load_model()
self.split = 1
max_displacement = 5,5,5
deg = (0,5,10)
scales = 0
self.transform = tio.Compose([
tio.RandomElasticDeformation(max_displacement=max_displacement),
tio.RandomAffine(scales=scales, degrees=deg)
])
# Training
self.train_original_classes = False
self.start_epoch = 0
self.epoch = 1000
self.loss = torch.nn.CrossEntropyLoss()
# self.loss = SoftDiceLoss(self.n_classes)
self.hot = 0
self.batchsize = 2
self.lr_rate = 5e-5 #5e-4 # 1e-2 #5e-5
self.optimizer = optim.Adam(self.net.parameters(), lr = self.lr_rate, weight_decay=0)
self.optimizer.zero_grad()
self.validate_every_k_epochs = 1
# Scheduler list : [lambdarule_1]
# self.lr_scheduler = get_scheduler(self.optimizer, "multistep")
self.lr_scheduler = get_scheduler(self.optimizer, "constant", self.lr_rate)
# self.lr_scheduler = get_scheduler(self.optimizer, "lambdarule_1", self.lr_rate)
# Other
self.classes_name = ['background','pancreas']
self.look_small = False
# print(0)
#
# # plt.imshow(one_subject.mri.data[0, 40, :, :])
#
# print(1)
training_transform = tio.Compose(
[
tio.ToCanonical(),
tio.RandomBlur(std=(0, 1), seed=seed, p=0.1), # blur 50% of times
tio.RandomNoise(std=5, seed=1, p=0.5), # Gaussian noise 50% of times
tio.OneOf(
{ # either
tio.RandomAffine(scales=(0.95, 1.05), degrees=5, seed=seed):
0.75, # random affine
tio.RandomElasticDeformation(max_displacement=(5, 5, 5),
seed=seed):
0.25, # or random elastic deformation
},
p=0.8), # applied to 80% of images
])
for one_subject in dataset:
image0 = one_subject.mri
plt.imshow(image0.data[0, int(image0.shape[1] / 2), :, :])
plt.show()
break
dataset_augmented = SubjectsDataset(subjects, transform=training_transform)
for one_subject in dataset_augmented:
image = one_subject.mri
plt.imshow(image.data[0, int(image.shape[1] / 2), :, :])
plt.show()
import os
from typing import Tuple
import numpy as np
import pandas as pd
import torch
import torchio as tio
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
from torchio import Image
from torchvision.transforms import Compose
transforms_dict = {
tio.RandomAffine(): 0.55,
tio.RandomElasticDeformation(): 0.25
}
transforms_dict2 = {tio.RandomBlur(): 0.25, tio.RandomMotion(): 0.25}
# for aumentation
transform_flip = tio.OneOf(transforms_dict)
class ADNIDataloaderAllData(Dataset):
def __init__(self, df, root_dir, transform):
self.df = df
self.root_dir = root_dir
self.transform = transform
def __len__(self):
return len(self.df)
def get_context(
device,
variables,
fold=0,
predict_hbt=False,
training_batch_size=4,
):
context = TorchContext(device, name="dmri-hippo", variables=variables)
context.file_paths.append(os.path.abspath(__file__))
context.config.update({'fold': fold})
input_images = ["mean_dwi", "md", "fa"]
output_labels = ["whole_roi", "hbt_roi"]
subject_loader = ComposeLoaders([
ImageLoader(glob_pattern="mean_dwi.*",
image_name='mean_dwi',
image_constructor=tio.ScalarImage),
ImageLoader(glob_pattern="md.*",
image_name='md',
image_constructor=tio.ScalarImage),
ImageLoader(glob_pattern="fa.*",
image_name='fa',
image_constructor=tio.ScalarImage),
# ImageLoader(glob_pattern="full_dwi.*", image_name='full_dwi', image_constructor=tio.ScalarImage),
# TensorLoader(glob_pattern="full_dwi_grad.b", tensor_name="grad", belongs_to="full_dwi"),
ImageLoader(glob_pattern="whole_roi.*",
image_name="whole_roi",
image_constructor=tio.LabelMap,
label_values={
"left_whole": 1,
"right_whole": 2
}),
ImageLoader(glob_pattern="whole_roi_alt.*",
image_name="whole_roi_alt",
image_constructor=tio.LabelMap,
label_values={
"left_whole": 1,
"right_whole": 2
}),
ImageLoader(glob_pattern="hbt_roi.*",
image_name="hbt_roi",
image_constructor=tio.LabelMap,
label_values={
"left_head": 1,
"left_body": 2,
"left_tail": 3,
"right_head": 4,
"right_body": 5,
"right_tail": 6
}),
ImageLoader(glob_pattern="../../atlas/whole_roi_union.*",
image_name="whole_roi_union",
image_constructor=tio.LabelMap,
uniform=True),
AttributeLoader(glob_pattern='attributes.*'),
AttributeLoader(
glob_pattern='../../attributes/cross_validation_split.json',
multi_subject=True,
uniform=True),
AttributeLoader(
glob_pattern='../../attributes/ab300_validation_subjects.json',
multi_subject=True,
uniform=True),
AttributeLoader(
glob_pattern='../../attributes/cbbrain_test_subjects.json',
multi_subject=True,
uniform=True),
])
cohorts = {}
cohorts['all'] = RequireAttributes(input_images)
cohorts['cross_validation'] = RequireAttributes(['fold'])
cohorts['training'] = ComposeFilters(
[cohorts['cross_validation'],
ForbidAttributes({"fold": fold})])
cohorts['cbbrain_validation'] = ComposeFilters(
[cohorts['cross_validation'],
RequireAttributes({"fold": fold})])
cohorts['cbbrain_test'] = RequireAttributes({'cbbrain_test': True})
cohorts['ab300_validation'] = RequireAttributes({'ab300_validation': True})
cohorts['ab300_validation_plot'] = ComposeFilters(
[cohorts['ab300_validation'],
RandomSelectFilter(num_subjects=20)])
cohorts['cbbrain'] = RequireAttributes({"protocol": "cbbrain"})
cohorts['ab300'] = RequireAttributes({"protocol": "ab300"})
cohorts['rescans'] = ForbidAttributes({"rescan_id": "None"})
cohorts['fasd'] = RequireAttributes({"pathologies": "FASD"})
cohorts['inter_rater'] = RequireAttributes(["whole_roi_alt"])
common_transforms_1 = tio.Compose([
tio.CropOrPad((96, 88, 24),
padding_mode='minimum',
mask_name='whole_roi_union'),
CustomRemapLabels(remapping=[("right_whole", 2, 1)],
masking_method="Right",
include=["whole_roi"]),
CustomRemapLabels(remapping=[("right_head", 4, 1),
("right_body", 5, 2),
("right_tail", 6, 3)],
masking_method="Right",
include=["hbt_roi"]),
])
noise = tio.RandomNoise(std=0.035, p=0.3)
blur = tio.RandomBlur((0, 1), p=0.2)
standard_augmentations = tio.Compose([
tio.RandomFlip(axes=(0, 1, 2)),
tio.RandomElasticDeformation(p=0.5,
num_control_points=(7, 7, 4),
locked_borders=1,
image_interpolation='bspline',
exclude=["full_dwi"]),
tio.RandomBiasField(p=0.5),
tio.RescaleIntensity((0, 1), (0.01, 99.9)),
tio.RandomGamma(p=0.8),
tio.RescaleIntensity((-1, 1)),
tio.OneOf([
tio.Compose([blur, noise]),
tio.Compose([noise, blur]),
])
],
exclude="full_dwi")
common_transforms_2 = tio.Compose([
tio.RescaleIntensity((-1., 1.), (0.5, 99.5)),
ConcatenateImages(image_names=["mean_dwi", "md", "fa"],
image_channels=[1, 1, 1],
new_image_name="X"),
RenameProperty(old_name="hbt_roi" if predict_hbt else "whole_roi",
new_name="y"),
CustomOneHot(include=["y"])
])
transforms = {
'default':
tio.Compose([common_transforms_1, common_transforms_2]),
'training':
tio.Compose(
[common_transforms_1, standard_augmentations,
common_transforms_2]),
}
context.add_component("dataset",
SubjectFolder,
root='$DATASET_PATH',
subject_path="subjects",
subject_loader=subject_loader,
cohorts=cohorts,
transforms=transforms)
context.add_component("model",
NestedResUNet,
input_channels=3,
output_channels=4 if predict_hbt else 2,
filters=40,
dropout_p=0.2)
context.add_component("optimizer",
Adam,
params="self.model.parameters()",
lr=0.0002)
context.add_component("criterion", HybridLogisticDiceLoss)
training_evaluators = [
ScheduledEvaluation(evaluator=SegmentationEvaluator(
'y_pred_eval', 'y_eval'),
log_name='training_segmentation_eval',
interval=10),
ScheduledEvaluation(evaluator=ContourImageEvaluator(
"Axial",
'mean_dwi',
'y_pred_eval',
'y_eval',
slice_id=12,
legend=True,
ncol=2,
split_subjects=False),
log_name=f"contour_image_training",
interval=50),
]
curve_params = {
"left_whole":
np.array([-1.96312119e-01, 9.46668029e+00, 2.33635173e+03]),
"right_whole":
np.array([-2.68467331e-01, 1.67925603e+01, 2.07224236e+03])
}
validation_evaluators = [
ScheduledEvaluation(evaluator=LabelMapEvaluator(
'y_pred_eval',
curve_params=curve_params,
curve_attribute='age',
stats_to_output=('volume', 'error', 'absolute_error',
'squared_error', 'percent_diff')),
log_name="predicted_label_eval",
cohorts=['cbbrain_validation', 'ab300_validation'],
interval=50),
ScheduledEvaluation(evaluator=SegmentationEvaluator(
"y_pred_eval", "y_eval"),
log_name="segmentation_eval",
cohorts=['cbbrain_validation'],
interval=50),
ScheduledEvaluation(
evaluator=ContourImageEvaluator("Axial",
"mean_dwi",
"y_pred_eval",
"y_eval",
slice_id=10,
legend=True,
ncol=5,
split_subjects=False),
log_name="contour_image_axial",
cohorts=['cbbrain_validation', 'ab300_validation_plot'],
interval=250),
ScheduledEvaluation(
evaluator=ContourImageEvaluator("Coronal",
"mean_dwi",
"y_pred_eval",
"y_eval",
slice_id=44,
legend=True,
ncol=2,
split_subjects=False),
log_name="contour_image_coronal",
cohorts=['cbbrain_validation', 'ab300_validation_plot'],
interval=250),
]
def scoring_function(evaluation_dict):
# Grab the output of the SegmentationEvaluator
seg_eval_cbbrain = evaluation_dict['segmentation_eval'][
'cbbrain_validation']["summary_stats"]
# Get the mean dice for each label (the mean is across subjects)
cbbrain_dice = seg_eval_cbbrain['mean', :, 'dice']
# Now take the mean across all labels
cbbrain_dice = cbbrain_dice.mean()
score = cbbrain_dice
return score
train_predictor = StandardPredict(sagittal_split=True,
image_names=['X', 'y'])
validation_predictor = StandardPredict(sagittal_split=True,
image_names=['X'])
train_dataloader_factory = StandardDataLoader(sampler=RandomSampler)
validation_dataloader_factory = StandardDataLoader(
sampler=SequentialSampler)
context.add_component(
"trainer",
SegmentationTrainer,
training_batch_size=training_batch_size,
save_rate=100,
scoring_interval=50,
scoring_function=scoring_function,
one_time_evaluators=[],
training_evaluators=training_evaluators,
validation_evaluators=validation_evaluators,
max_iterations_with_no_improvement=2000,
train_predictor=train_predictor,
validation_predictor=validation_predictor,
train_dataloader_factory=train_dataloader_factory,
validation_dataloader_factory=validation_dataloader_factory)
return context
