def train_pre_transforms(self, context: Context): # Dataset preparation t: List[Any] = [ LoadImaged(keys=("image", "label")), AddChanneld(keys=("image", "label")), SpatialCropForegroundd(keys=("image", "label"), source_key="label", spatial_size=self.roi_size), Resized(keys=("image", "label"), spatial_size=self.model_size, mode=("area", "nearest")), NormalizeIntensityd(keys="image", subtrahend=208.0, divisor=388.0), # type: ignore ] if self.dimension == 3: t.append(FindAllValidSlicesd(label="label", sids="sids")) t.extend([ AddInitialSeedPointd(label="label", guidance="guidance", sids="sids"), AddGuidanceSignald(image="image", guidance="guidance"), EnsureTyped(keys=("image", "label"), device=context.device), SelectItemsd(keys=("image", "label", "guidance")), ]) return t
def pre_transforms(self, data=None) -> Sequence[Callable]: t = [ LoadImaged(keys="image"), AsChannelFirstd(keys="image"), Spacingd(keys="image", pixdim=[1.0] * self.dimension, mode="bilinear"), AddGuidanceFromPointsd(ref_image="image", guidance="guidance", dimensions=self.dimension), ] if self.dimension == 2: t.append(Fetch2DSliced(keys="image", guidance="guidance")) t.extend([ AddChanneld(keys="image"), SpatialCropGuidanced(keys="image", guidance="guidance", spatial_size=self.spatial_size), Resized(keys="image", spatial_size=self.model_size, mode="area"), ResizeGuidanced(guidance="guidance", ref_image="image"), NormalizeIntensityd(keys="image", subtrahend=208, divisor=388), # type: ignore AddGuidanceSignald(image="image", guidance="guidance"), EnsureTyped(keys="image", device=data.get("device") if data else None), ]) return t
def __init__(self, tranforms): self.tranform_list = [] for tranform in tranforms: if 'LoadImaged' == tranform: self.tranform_list.append(LoadImaged(keys=["image", "label"])) elif 'AsChannelFirstd' == tranform: self.tranform_list.append(AsChannelFirstd(keys="image")) elif 'ConvertToMultiChannelBasedOnBratsClassesd' == tranform: self.tranform_list.append(ConvertToMultiChannelBasedOnBratsClassesd(keys="label")) elif 'Spacingd' == tranform: self.tranform_list.append(Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest"))) elif 'Orientationd' == tranform: self.tranform_list.append(Orientationd(keys=["image", "label"], axcodes="RAS")) elif 'CenterSpatialCropd' == tranform: self.tranform_list.append(CenterSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64])) elif 'NormalizeIntensityd' == tranform: self.tranform_list.append(NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True)) elif 'ToTensord' == tranform: self.tranform_list.append(ToTensord(keys=["image", "label"])) elif 'Activations' == tranform: self.tranform_list.append(Activations(sigmoid=True)) elif 'AsDiscrete' == tranform: self.tranform_list.append(AsDiscrete(threshold_values=True)) else: raise ValueError( f"Unsupported tranform: {tranform}. Please add it to support it." ) super().__init__(self.tranform_list)
def test_channel_wise(self): key = 'img' normalizer = NormalizeIntensityd(keys=key, nonzero=True, channel_wise=True) input_data = {key: np.array([[0., 3., 0., 4.], [0., 4., 0., 5.]])} expected = np.array([[0., -1., 0., 1.], [0., -1., 0., 1.]]) np.testing.assert_allclose(expected, normalizer(input_data)[key])
def test_nonzero(self, input_param, input_data, expected_data): key = "img" normalizer = NormalizeIntensityd(**input_param) normalized = normalizer(input_data)[key] self.assertEqual(type(input_data[key]), type(normalized)) if isinstance(normalized, torch.Tensor): self.assertEqual(input_data[key].device, normalized.device) assert_allclose(normalized, expected_data)
def test_channel_wise(self): key = "img" normalizer = NormalizeIntensityd(keys=key, nonzero=True, channel_wise=True) input_data = { key: np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0, 5.0]]) } expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]]) np.testing.assert_allclose(expected, normalizer(input_data)[key])
def get_transforms(self): self.logger.info("Getting transforms...") # Setup transforms of data sets train_transforms = Compose([ LoadNiftid(keys=["image", "label"]), AddChanneld(keys=["image", "label"]), Orientationd(keys=["image", "label"], axcodes="RAS"), NormalizeIntensityd(keys=["image"]), SpatialPadd(keys=["image", "label"], spatial_size=self.pad_crop_shape), RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0), RandSpatialCropd(keys=["image", "label"], roi_size=self.pad_crop_shape, random_center=True, random_size=False), ToTensord(keys=["image", "label"]), ]) val_transforms = Compose([ LoadNiftid(keys=["image", "label"]), AddChanneld(keys=["image", "label"]), Orientationd(keys=["image", "label"], axcodes="RAS"), NormalizeIntensityd(keys=["image"]), SpatialPadd(keys=["image", "label"], spatial_size=self.pad_crop_shape), RandSpatialCropd( keys=["image", "label"], roi_size=self.pad_crop_shape, random_center=True, random_size=False, ), ToTensord(keys=["image", "label"]), ]) test_transforms = Compose([ LoadNiftid(keys=["image", "label"]), AddChanneld(keys=["image", "label"]), Orientationd(keys=["image", "label"], axcodes="RAS"), NormalizeIntensityd(keys=["image"]), ToTensord(keys=["image", "label"]), ]) return train_transforms, val_transforms, test_transforms
def test_image_normalize_intensityd(self, im_type): key = "img" im = im_type(self.imt) normalizer = NormalizeIntensityd(keys=[key]) normalized = normalizer({key: im})[key] expected = (self.imt - np.mean(self.imt)) / np.std(self.imt) self.assertEqual(type(im), type(normalized)) if isinstance(normalized, torch.Tensor): self.assertEqual(im.device, normalized.device) assert_allclose(normalized, im_type(expected), rtol=1e-3)
def test_image_normalize_intensityd(self, im_type): key = "img" im = im_type(self.imt) normalizer = NormalizeIntensityd(keys=[key]) normalized = normalizer({key: im})[key] expected = (self.imt - np.mean(self.imt)) / np.std(self.imt) assert_allclose(normalized, im_type(expected), rtol=1e-3, type_test="tensor")
def test_channel_wise(self, im_type): key = "img" normalizer = NormalizeIntensityd(keys=key, nonzero=True, channel_wise=True) input_data = { key: im_type(np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0, 5.0]])) } normalized = normalizer(input_data)[key] expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]]) assert_allclose(normalized, im_type(expected), type_test="tensor")
def get_test_transforms(end_image_shape): test_transforms = Compose([ LoadNiftid(keys=["image"]), AddChanneld(keys=["image"]), Orientationd(keys=["image"], axcodes="RAS"), Winsorized(keys=["image"]), NormalizeIntensityd(keys=["image"]), ScaleIntensityd(keys=["image"]), SpatialPadd(keys=["image"], spatial_size=end_image_shape), ToTensord(keys=["image"]), ]) return test_transforms
def pre_transforms(self, data): return [ LoadImaged(keys="image"), AsChannelFirstd(keys="image"), Spacingd(keys="image", pixdim=[1.0, 1.0, 1.0], mode="bilinear"), AddGuidanceFromPointsd(ref_image="image", guidance="guidance", dimensions=3), AddChanneld(keys="image"), SpatialCropGuidanced(keys="image", guidance="guidance", spatial_size=self.spatial_size), Resized(keys="image", spatial_size=self.model_size, mode="area"), ResizeGuidanced(guidance="guidance", ref_image="image"), NormalizeIntensityd(keys="image", subtrahend=208, divisor=388), AddGuidanceSignald(image="image", guidance="guidance"), ]
def setup(self, stage): data_dir = 'data/' # Train imgs/masks train_imgs = [] with open(data_dir + 'train_imgs.txt', 'r') as f: train_imgs = [image.rstrip() for image in f.readlines()] train_masks = [] with open(data_dir + 'train_masks.txt', 'r') as f: train_masks = [mask.rstrip() for mask in f.readlines()] train_dicts = [{'image': image, 'mask': mask} for (image, mask) in zip(train_imgs, train_masks)] train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2) # Basic transforms data_keys = ["image", "mask"] data_transforms = Compose( [ LoadNiftid(keys=data_keys), AddChanneld(keys=data_keys), NormalizeIntensityd(keys="image"), RandCropByPosNegLabeld( keys=data_keys, label_key="mask", size=(256, 256, 16), num_samples=4, image_key="image" ), ] ) self.train_dataset = monai.data.CacheDataset( data=train_dicts, transform=Compose( [ data_transforms, self.augmentations, ToTensord(keys=data_keys) ] ), cache_rate=1.0 ) self.val_dataset = monai.data.CacheDataset( data=val_dicts, transform=Compose( [ data_transforms, ToTensord(keys=data_keys) ] ), cache_rate=1.0 )
def test_channel_wise(self, im_type): key = "img" normalizer = NormalizeIntensityd(keys=key, nonzero=True, channel_wise=True) input_data = { key: im_type(np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0, 5.0]])) } normalized = normalizer(input_data)[key] self.assertEqual(type(input_data[key]), type(normalized)) if isinstance(normalized, torch.Tensor): self.assertEqual(input_data[key].device, normalized.device) expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]]) assert_allclose(normalized, im_type(expected))
def setup(self, stage): data_dir = "data/" # Train imgs/masks train_imgs = [] with open(data_dir + "train_imgs.txt", "r") as f: train_imgs = [image.rstrip() for image in f.readlines()] train_masks = [] with open(data_dir + "train_masks.txt", "r") as f: train_masks = [mask.rstrip() for mask in f.readlines()] train_dicts = [{ "image": image, "mask": mask } for (image, mask) in zip(train_imgs, train_masks)] train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2) # Basic transforms data_keys = ["image", "mask"] data_transforms = Compose([ AddChanneld(keys=data_keys), NormalizeIntensityd(keys="image"), RandCropByPosNegLabeld( keys=data_keys, label_key="mask", spatial_size=self.hparams.patch_size, num_samples=4, image_key="image", ), ]) self.train_dataset = monai.data.CacheDataset( data=train_dicts, transform=Compose([data_transforms, ToTensord(keys=data_keys)]), cache_rate=1.0, ) self.val_dataset = monai.data.CacheDataset( data=val_dicts, transform=Compose([data_transforms, ToTensord(keys=data_keys)]), cache_rate=1.0, )
def pre_transforms(self, data=None): return [ LoadImaged(keys="image"), AsChannelFirstd(keys="image"), Spacingd(keys="image", pixdim=[1.0, 1.0], mode="bilinear"), AddGuidanceFromPointsd(ref_image="image", guidance="guidance", dimensions=2), Fetch2DSliced(keys="image", guidance="guidance"), AddChanneld(keys="image"), SpatialCropGuidanced(keys="image", guidance="guidance", spatial_size=[256, 256]), Resized(keys="image", spatial_size=[256, 256], mode="area"), ResizeGuidanced(guidance="guidance", ref_image="image"), NormalizeIntensityd(keys="image", subtrahend=208, divisor=388), # type: ignore AddGuidanceSignald(image="image", guidance="guidance"), ]
def get_pre_transforms(roi_size, model_size, dimensions): t = [ LoadImaged(keys=('image', 'label')), AddChanneld(keys=('image', 'label')), SpatialCropForegroundd(keys=('image', 'label'), source_key='label', spatial_size=roi_size), Resized(keys=('image', 'label'), spatial_size=model_size, mode=('area', 'nearest')), NormalizeIntensityd(keys='image', subtrahend=208.0, divisor=388.0) ] if dimensions == 3: t.append(FindAllValidSlicesd(label='label', sids='sids')) t.extend([ AddInitialSeedPointd(label='label', guidance='guidance', sids='sids'), AddGuidanceSignald(image='image', guidance='guidance'), ToTensord(keys=('image', 'label')) ]) return Compose(t)
def get_pre_transforms(roi_size, model_size, dimensions): t = [ LoadImaged(keys=("image", "label")), AddChanneld(keys=("image", "label")), SpatialCropForegroundd(keys=("image", "label"), source_key="label", spatial_size=roi_size), Resized(keys=("image", "label"), spatial_size=model_size, mode=("area", "nearest")), NormalizeIntensityd(keys="image", subtrahend=208.0, divisor=388.0), ] if dimensions == 3: t.append(FindAllValidSlicesd(label="label", sids="sids")) t.extend([ AddInitialSeedPointd(label="label", guidance="guidance", sids="sids"), AddGuidanceSignald(image="image", guidance="guidance"), EnsureTyped(keys=("image", "label")), ]) return Compose(t)
def setup(self, stage): data_df = pd.read_csv( '/data/shared/prostate/yale_prostate/input_lists/MR_yale.csv') train_imgs = data_df['IMAGE'][0:295].tolist() train_masks = data_df['SEGM'][0:295].tolist() train_dicts = [{ 'image': image, 'mask': mask } for (image, mask) in zip(train_imgs, train_masks)] train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.15) # Basic transforms data_keys = ["image", "mask"] data_transforms = Compose([ LoadNiftid(keys=data_keys), AddChanneld(keys=data_keys), NormalizeIntensityd(keys="image"), RandCropByPosNegLabeld(keys=data_keys, label_key="mask", spatial_size=self.hparams.patch_size, num_samples=4, image_key="image", pos=0.7, neg=0.3), ]) self.train_dataset = monai.data.CacheDataset( data=train_dicts, transform=Compose([data_transforms, ToTensord(keys=data_keys)]), cache_rate=1.0) self.val_dataset = monai.data.CacheDataset( data=val_dicts, transform=Compose([data_transforms, ToTensord(keys=data_keys)]), cache_rate=1.0)
def get_data_loaders(batch_size, patch_size): ### Data collection data_dir = "../../data/" print("Available directories: ", os.listdir(data_dir)) # Get paths for images and masks, organize into dictionaries images = sorted(glob.glob(data_dir + "**/*CTImg*", recursive=True)) masks = sorted(glob.glob(data_dir + "**/*Mask*", recursive=True)) # Dataset selection train_dicts = select_animals(images, masks, [12, 13, 14, 18, 20]) val_dicts = select_animals(images, masks, [25]) data_keys = ["image", "mask"] # Data transformation data_transforms = Compose([ LoadNiftid(keys=data_keys), AddChanneld(keys=data_keys), NormalizeIntensityd(keys="image"), RandCropByPosNegLabeld(keys=data_keys, label_key="mask", size=patch_size, num_samples=16, image_key="image"), ToTensord(keys=data_keys), ]) # Data loaders data_loaders = { "train": create_loader( train_dicts, batch_size=batch_size, transforms=data_transforms, shuffle=True, ), "val": create_loader(val_dicts, transforms=data_transforms), } for key in data_loaders: print(key, len(data_loaders[key])) return data_loaders
'label': label_name } for image_name, label_name in zip(train_images, train_segs)] monai_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), Orientationd(keys=["image", "label"], axcodes="RAS"), # ThresholdIntensityd(keys=['image'], threshold=-135, above=True, cval=-135), # ThresholdIntensityd(keys=['image'], threshold=215, above=False, cval=215), CropForegroundd( keys=['image', 'label'], source_key='image', start_coord_key='foreground_start_coord', end_coord_key='foreground_end_coord', ), # crop CropForeground NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), # Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')), SpatialPadd(keys=['image', 'label'], spatial_size=opt.patch_size, method='end'), RandSpatialCropd(keys=['image', 'label'], roi_size=opt.patch_size, random_size=False), ToTensord(keys=[ 'image', 'label', 'foreground_start_coord', 'foreground_end_coord' ], ) ] transform = Compose(monai_transforms) check_ds = monai.data.Dataset(data=data_dicts, transform=transform)
def test_nonzero(self, input_param, input_data, expected_data): normalizer = NormalizeIntensityd(**input_param) np.testing.assert_allclose(expected_data, normalizer(input_data)['img'])
def test_image_normalize_intensityd(self): key = 'img' normalizer = NormalizeIntensityd(keys=[key]) normalized = normalizer({key: self.imt}) expected = (self.imt - np.mean(self.imt)) / np.std(self.imt) np.testing.assert_allclose(normalized[key], expected, rtol=1e-6)
def main(): opt = Options().parse() # monai.config.print_config() logging.basicConfig(stream=sys.stdout, level=logging.INFO) set_determinism(seed=0) logging.basicConfig(stream=sys.stdout, level=logging.INFO) device = torch.device(opt.gpu_id) # ------- Data loader creation ---------- # images images = sorted(glob(os.path.join(opt.images_folder, 'image*.nii'))) segs = sorted(glob(os.path.join(opt.labels_folder, 'label*.nii'))) train_files = [] val_files = [] for i in range(opt.models_ensemble): train_files.append([{ "image": img, "label": seg } for img, seg in zip( images[:(opt.split_val * i)] + images[(opt.split_val * (i + 1)):(len(images) - opt.split_val)], segs[:(opt.split_val * i)] + segs[(opt.split_val * (i + 1)):(len(images) - opt.split_val)])]) val_files.append([{ "image": img, "label": seg } for img, seg in zip( images[(opt.split_val * i):(opt.split_val * (i + 1))], segs[(opt.split_val * i):(opt.split_val * (i + 1))])]) test_files = [{ "image": img, "label": seg } for img, seg in zip(images[(len(images) - opt.split_test):len(images)], segs[( len(images) - opt.split_test):len(images)])] # ----------- Transforms list -------------- if opt.resolution is not None: train_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36), padding_mode="zeros"), Rand3DElasticd(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, sigma_range=(5, 8), magnitude_range=(100, 200), scale_range=(0.15, 0.15, 0.15), padding_mode="zeros"), RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1), RandGaussianNoised(keys=['image'], prob=0.1, mean=np.random.uniform(0, 0.5), std=np.random.uniform(0, 1)), RandShiftIntensityd(keys=['image'], offsets=np.random.uniform(0, 0.3), prob=0.1), RandSpatialCropd(keys=['image', 'label'], roi_size=opt.patch_size, random_size=False), ToTensord(keys=['image', 'label']) ] val_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')), ToTensord(keys=['image', 'label']) ] else: train_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36), padding_mode="zeros"), Rand3DElasticd(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, sigma_range=(5, 8), magnitude_range=(100, 200), scale_range=(0.15, 0.15, 0.15), padding_mode="zeros"), RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1), RandGaussianNoised(keys=['image'], prob=0.1, mean=np.random.uniform(0, 0.5), std=np.random.uniform(0, 1)), RandShiftIntensityd(keys=['image'], offsets=np.random.uniform(0, 0.3), prob=0.1), RandSpatialCropd(keys=['image', 'label'], roi_size=opt.patch_size, random_size=False), ToTensord(keys=['image', 'label']) ] val_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), ToTensord(keys=['image', 'label']) ] train_transforms = Compose(train_transforms) val_transforms = Compose(val_transforms) # ---------- Creation of DataLoaders ------------- train_dss = [ CacheDataset(data=train_files[i], transform=train_transforms) for i in range(opt.models_ensemble) ] train_loaders = [ DataLoader(train_dss[i], batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers, pin_memory=torch.cuda.is_available()) for i in range(opt.models_ensemble) ] val_dss = [ CacheDataset(data=val_files[i], transform=val_transforms) for i in range(opt.models_ensemble) ] val_loaders = [ DataLoader(val_dss[i], batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available()) for i in range(opt.models_ensemble) ] test_ds = CacheDataset(data=test_files, transform=val_transforms) test_loader = DataLoader(test_ds, batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available()) def train(index): # ---------- Build the nn-Unet network ------------ if opt.resolution is None: sizes, spacings = opt.patch_size, opt.spacing else: sizes, spacings = opt.patch_size, opt.resolution strides, kernels = [], [] while True: spacing_ratio = [sp / min(spacings) for sp in spacings] stride = [ 2 if ratio <= 2 and size >= 8 else 1 for (ratio, size) in zip(spacing_ratio, sizes) ] kernel = [3 if ratio <= 2 else 1 for ratio in spacing_ratio] if all(s == 1 for s in stride): break sizes = [i / j for i, j in zip(sizes, stride)] spacings = [i * j for i, j in zip(spacings, stride)] kernels.append(kernel) strides.append(stride) strides.insert(0, len(spacings) * [1]) kernels.append(len(spacings) * [3]) net = monai.networks.nets.DynUNet( spatial_dims=3, in_channels=opt.in_channels, out_channels=opt.out_channels, kernel_size=kernels, strides=strides, upsample_kernel_size=strides[1:], res_block=True, # act=act_type, # norm=Norm.BATCH, ).to(device) from torch.autograd import Variable from torchsummaryX import summary data = Variable( torch.randn(int(opt.batch_size), int(opt.in_channels), int(opt.patch_size[0]), int(opt.patch_size[1]), int(opt.patch_size[2]))).cuda() out = net(data) summary(net, data) print("out size: {}".format(out.size())) # if opt.preload is not None: # net.load_state_dict(torch.load(opt.preload)) # ---------- ------------------------ ------------ optim = torch.optim.Adam(net.parameters(), lr=opt.lr) lr_scheduler = torch.optim.lr_scheduler.LambdaLR( optim, lr_lambda=lambda epoch: (1 - epoch / opt.epochs)**0.9) loss_function = monai.losses.DiceCELoss(sigmoid=True) val_post_transforms = Compose([ Activationsd(keys="pred", sigmoid=True), AsDiscreted(keys="pred", threshold_values=True), # KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]) ]) val_handlers = [ StatsHandler(output_transform=lambda x: None), CheckpointSaver(save_dir="./runs/", save_dict={"net": net}, save_key_metric=True), ] evaluator = SupervisedEvaluator( device=device, val_data_loader=val_loaders[index], network=net, inferer=SlidingWindowInferer(roi_size=opt.patch_size, sw_batch_size=opt.batch_size, overlap=0.5), post_transform=val_post_transforms, key_val_metric={ "val_mean_dice": MeanDice( include_background=True, output_transform=lambda x: (x["pred"], x["label"]), ) }, val_handlers=val_handlers) train_post_transforms = Compose([ Activationsd(keys="pred", sigmoid=True), AsDiscreted(keys="pred", threshold_values=True), # KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]), ]) train_handlers = [ ValidationHandler(validator=evaluator, interval=5, epoch_level=True), LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True), StatsHandler(tag_name="train_loss", output_transform=lambda x: x["loss"]), CheckpointSaver(save_dir="./runs/", save_dict={ "net": net, "opt": optim }, save_final=True, epoch_level=True), ] trainer = SupervisedTrainer( device=device, max_epochs=opt.epochs, train_data_loader=train_loaders[index], network=net, optimizer=optim, loss_function=loss_function, inferer=SimpleInferer(), post_transform=train_post_transforms, amp=False, train_handlers=train_handlers, ) trainer.run() return net models = [train(i) for i in range(opt.models_ensemble)] # -------- Test the models --------- def ensemble_evaluate(post_transforms, models): evaluator = EnsembleEvaluator( device=device, val_data_loader=test_loader, pred_keys=opt.pred_keys, networks=models, inferer=SlidingWindowInferer(roi_size=opt.patch_size, sw_batch_size=opt.batch_size, overlap=0.5), post_transform=post_transforms, key_val_metric={ "test_mean_dice": MeanDice( include_background=True, output_transform=lambda x: (x["pred"], x["label"]), ) }, ) evaluator.run() mean_post_transforms = Compose([ MeanEnsembled( keys=opt.pred_keys, output_key="pred", # in this particular example, we use validation metrics as weights weights=opt.weights_models, ), Activationsd(keys="pred", sigmoid=True), AsDiscreted(keys="pred", threshold_values=True), # KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]) ]) print('Results from MeanEnsembled:') ensemble_evaluate(mean_post_transforms, models) vote_post_transforms = Compose([ Activationsd(keys=opt.pred_keys, sigmoid=True), # transform data into discrete before voting AsDiscreted(keys=opt.pred_keys, threshold_values=True), # KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]), VoteEnsembled(keys=opt.pred_keys, output_key="pred"), ]) print('Results from VoteEnsembled:') ensemble_evaluate(vote_post_transforms, models)
def run_training(train_file_list, valid_file_list, config_info): """ Pipeline to train a dynUNet segmentation model in MONAI. It is composed of the following main blocks: * Data Preparation: Extract the filenames and prepare the training/validation processing transforms * Load Data: Load training and validation data to PyTorch DataLoader * Network Preparation: Define the network, loss function, optimiser and learning rate scheduler * MONAI Evaluator: Initialise the dynUNet evaluator, i.e. the class providing utilities to perform validation during training. Attach handlers to save the best model on the validation set. A 2D sliding window approach on the 3D volume is used at evaluation. The mean 3D Dice is used as validation metric. * MONAI Trainer: Initialise the dynUNet trainer, i.e. the class providing utilities to perform the training loop. * Run training: The MONAI trainer is run, performing training and validation during training. Args: train_file_list: .txt or .csv file (with no header) storing two-columns filenames for training: image filename in the first column and segmentation filename in the second column. The two columns should be separated by a comma. See monaifbs/config/mock_train_file_list_for_dynUnet_training.txt for an example of the expected format. valid_file_list: .txt or .csv file (with no header) storing two-columns filenames for validation: image filename in the first column and segmentation filename in the second column. The two columns should be separated by a comma. See monaifbs/config/mock_valid_file_list_for_dynUnet_training.txt for an example of the expected format. config_info: dict, contains configuration parameters for sampling, network and training. See monaifbs/config/monai_dynUnet_training_config.yml for an example of the expected fields. """ """ Read input and configuration parameters """ # print MONAI config information logging.basicConfig(stream=sys.stdout, level=logging.INFO) print_config() # print to log the parameter setups print(yaml.dump(config_info)) # extract network parameters, perform checks/set defaults if not present and print them to log if 'seg_labels' in config_info['training'].keys(): seg_labels = config_info['training']['seg_labels'] else: seg_labels = [1] nr_out_channels = len(seg_labels) print("Considering the following {} labels in the segmentation: {}".format(nr_out_channels, seg_labels)) patch_size = config_info["training"]["inplane_size"] + [1] print("Considering patch size = {}".format(patch_size)) spacing = config_info["training"]["spacing"] print("Bringing all images to spacing = {}".format(spacing)) if 'model_to_load' in config_info['training'].keys() and config_info['training']['model_to_load'] is not None: model_to_load = config_info['training']['model_to_load'] if not os.path.exists(model_to_load): raise FileNotFoundError("Cannot find model: {}".format(model_to_load)) else: print("Loading model from {}".format(model_to_load)) else: model_to_load = None # set up either GPU or CPU usage if torch.cuda.is_available(): print("\n#### GPU INFORMATION ###") print("Using device number: {}, name: {}\n".format(torch.cuda.current_device(), torch.cuda.get_device_name())) current_device = torch.device("cuda:0") else: current_device = torch.device("cpu") print("Using device: {}".format(current_device)) # set determinism if required if 'manual_seed' in config_info['training'].keys() and config_info['training']['manual_seed'] is not None: seed = config_info['training']['manual_seed'] else: seed = None if seed is not None: print("Using determinism with seed = {}\n".format(seed)) set_determinism(seed=seed) """ Setup data output directory """ out_model_dir = os.path.join(config_info['output']['out_dir'], datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + '_' + config_info['output']['out_postfix']) print("Saving to directory {}\n".format(out_model_dir)) # create cache directory to store results for Persistent Dataset if 'cache_dir' in config_info['output'].keys(): out_cache_dir = config_info['output']['cache_dir'] else: out_cache_dir = os.path.join(out_model_dir, 'persistent_cache') persistent_cache: Path = Path(out_cache_dir) persistent_cache.mkdir(parents=True, exist_ok=True) """ Data preparation """ # Read the input files for training and validation print("*** Loading input data for training...") train_files = create_data_list_of_dictionaries(train_file_list) print("Number of inputs for training = {}".format(len(train_files))) val_files = create_data_list_of_dictionaries(valid_file_list) print("Number of inputs for validation = {}".format(len(val_files))) # Define MONAI processing transforms for the training data. This includes: # - Load Nifti files and convert to format Batch x Channel x Dim1 x Dim2 x Dim3 # - CropForegroundd: Reduce the background from the MR image # - InPlaneSpacingd: Perform in-plane resampling to the desired spacing, but preserve the resolution along the # last direction (lowest resolution) to avoid introducing motion artefact resampling errors # - SpatialPadd: Pad the in-plane size to the defined network input patch size [N, M] if needed # - NormalizeIntensityd: Apply whitening # - RandSpatialCropd: Crop a random patch from the input with size [B, C, N, M, 1] # - SqueezeDimd: Convert the 3D patch to a 2D one as input to the network (i.e. bring it to size [B, C, N, M]) # - Apply data augmentation (RandZoomd, RandRotated, RandGaussianNoised, RandGaussianSmoothd, RandScaleIntensityd, # RandFlipd) # - ToTensor: convert to pytorch tensor train_transforms = Compose( [ LoadNiftid(keys=["image", "label"]), AddChanneld(keys=["image", "label"]), CropForegroundd(keys=["image", "label"], source_key="image"), InPlaneSpacingd( keys=["image", "label"], pixdim=spacing, mode=("bilinear", "nearest"), ), SpatialPadd(keys=["image", "label"], spatial_size=patch_size, mode=["constant", "edge"]), NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=True), RandSpatialCropd(keys=["image", "label"], roi_size=patch_size, random_size=False), SqueezeDimd(keys=["image", "label"], dim=-1), RandZoomd( keys=["image", "label"], min_zoom=0.9, max_zoom=1.2, mode=("bilinear", "nearest"), align_corners=(True, None), prob=0.16, ), RandRotated(keys=["image", "label"], range_x=90, range_y=90, prob=0.2, keep_size=True, mode=["bilinear", "nearest"], padding_mode=["zeros", "border"]), RandGaussianNoised(keys=["image"], std=0.01, prob=0.15), RandGaussianSmoothd( keys=["image"], sigma_x=(0.5, 1.15), sigma_y=(0.5, 1.15), sigma_z=(0.5, 1.15), prob=0.15, ), RandScaleIntensityd(keys=["image"], factors=0.3, prob=0.15), RandFlipd(["image", "label"], spatial_axis=[0, 1], prob=0.5), ToTensord(keys=["image", "label"]), ] ) # Define MONAI processing transforms for the validation data # - Load Nifti files and convert to format Batch x Channel x Dim1 x Dim2 x Dim3 # - CropForegroundd: Reduce the background from the MR image # - InPlaneSpacingd: Perform in-plane resampling to the desired spacing, but preserve the resolution along the # last direction (lowest resolution) to avoid introducing motion artefact resampling errors # - SpatialPadd: Pad the in-plane size to the defined network input patch size [N, M] if needed # - NormalizeIntensityd: Apply whitening # - ToTensor: convert to pytorch tensor # NOTE: The validation data is kept 3D as a 2D sliding window approach is used throughout the volume at inference val_transforms = Compose( [ LoadNiftid(keys=["image", "label"]), AddChanneld(keys=["image", "label"]), CropForegroundd(keys=["image", "label"], source_key="image"), InPlaneSpacingd( keys=["image", "label"], pixdim=spacing, mode=("bilinear", "nearest"), ), SpatialPadd(keys=["image", "label"], spatial_size=patch_size, mode=["constant", "edge"]), NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=True), ToTensord(keys=["image", "label"]), ] ) """ Load data """ # create training data loader train_ds = PersistentDataset(data=train_files, transform=train_transforms, cache_dir=persistent_cache) train_loader = DataLoader(train_ds, batch_size=config_info['training']['batch_size_train'], shuffle=True, num_workers=config_info['device']['num_workers']) check_train_data = misc.first(train_loader) print("Training data tensor shapes:") print("Image = {}; Label = {}".format(check_train_data["image"].shape, check_train_data["label"].shape)) # create validation data loader if config_info['training']['batch_size_valid'] != 1: raise Exception("Batch size different from 1 at validation ar currently not supported") val_ds = PersistentDataset(data=val_files, transform=val_transforms, cache_dir=persistent_cache) val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=config_info['device']['num_workers']) check_valid_data = misc.first(val_loader) print("Validation data tensor shapes (Example):") print("Image = {}; Label = {}\n".format(check_valid_data["image"].shape, check_valid_data["label"].shape)) """ Network preparation """ print("*** Preparing the network ...") # automatically extracts the strides and kernels based on nnU-Net empirical rules spacings = spacing[:2] sizes = patch_size[:2] strides, kernels = [], [] while True: spacing_ratio = [sp / min(spacings) for sp in spacings] stride = [2 if ratio <= 2 and size >= 8 else 1 for (ratio, size) in zip(spacing_ratio, sizes)] kernel = [3 if ratio <= 2 else 1 for ratio in spacing_ratio] if all(s == 1 for s in stride): break sizes = [i / j for i, j in zip(sizes, stride)] spacings = [i * j for i, j in zip(spacings, stride)] kernels.append(kernel) strides.append(stride) strides.insert(0, len(spacings) * [1]) kernels.append(len(spacings) * [3]) # initialise the network net = DynUNet( spatial_dims=2, in_channels=1, out_channels=nr_out_channels, kernel_size=kernels, strides=strides, upsample_kernel_size=strides[1:], norm_name="instance", deep_supervision=True, deep_supr_num=2, res_block=False, ).to(current_device) print(net) # define the loss function loss_function = choose_loss_function(nr_out_channels, config_info) # define the optimiser and the learning rate scheduler opt = torch.optim.SGD(net.parameters(), lr=float(config_info['training']['lr']), momentum=0.95) scheduler = torch.optim.lr_scheduler.LambdaLR( opt, lr_lambda=lambda epoch: (1 - epoch / config_info['training']['nr_train_epochs']) ** 0.9 ) """ MONAI evaluator """ print("*** Preparing the dynUNet evaluator engine...\n") # val_post_transforms = Compose( # [ # Activationsd(keys="pred", sigmoid=True), # ] # ) val_handlers = [ StatsHandler(output_transform=lambda x: None), TensorBoardStatsHandler(log_dir=os.path.join(out_model_dir, "valid"), output_transform=lambda x: None, global_epoch_transform=lambda x: trainer.state.iteration), CheckpointSaver(save_dir=out_model_dir, save_dict={"net": net, "opt": opt}, save_key_metric=True, file_prefix='best_valid'), ] if config_info['output']['val_image_to_tensorboad']: val_handlers.append(TensorBoardImageHandler(log_dir=os.path.join(out_model_dir, "valid"), batch_transform=lambda x: (x["image"], x["label"]), output_transform=lambda x: x["pred"], interval=2)) # Define customized evaluator class DynUNetEvaluator(SupervisedEvaluator): def _iteration(self, engine, batchdata): inputs, targets = self.prepare_batch(batchdata) inputs, targets = inputs.to(engine.state.device), targets.to(engine.state.device) flip_inputs_1 = torch.flip(inputs, dims=(2,)) flip_inputs_2 = torch.flip(inputs, dims=(3,)) flip_inputs_3 = torch.flip(inputs, dims=(2, 3)) def _compute_pred(): pred = self.inferer(inputs, self.network) # use random flipping as data augmentation at inference flip_pred_1 = torch.flip(self.inferer(flip_inputs_1, self.network), dims=(2,)) flip_pred_2 = torch.flip(self.inferer(flip_inputs_2, self.network), dims=(3,)) flip_pred_3 = torch.flip(self.inferer(flip_inputs_3, self.network), dims=(2, 3)) return (pred + flip_pred_1 + flip_pred_2 + flip_pred_3) / 4 # execute forward computation self.network.eval() with torch.no_grad(): if self.amp: with torch.cuda.amp.autocast(): predictions = _compute_pred() else: predictions = _compute_pred() return {"image": inputs, "label": targets, "pred": predictions} evaluator = DynUNetEvaluator( device=current_device, val_data_loader=val_loader, network=net, inferer=SlidingWindowInferer2D(roi_size=patch_size, sw_batch_size=4, overlap=0.0), post_transform=None, key_val_metric={ "Mean_dice": MeanDice( include_background=False, to_onehot_y=True, mutually_exclusive=True, output_transform=lambda x: (x["pred"], x["label"]), ) }, val_handlers=val_handlers, amp=False, ) """ MONAI trainer """ print("*** Preparing the dynUNet trainer engine...\n") # train_post_transforms = Compose( # [ # Activationsd(keys="pred", sigmoid=True), # ] # ) validation_every_n_epochs = config_info['training']['validation_every_n_epochs'] epoch_len = len(train_ds) // train_loader.batch_size validation_every_n_iters = validation_every_n_epochs * epoch_len # define event handlers for the trainer writer_train = SummaryWriter(log_dir=os.path.join(out_model_dir, "train")) train_handlers = [ LrScheduleHandler(lr_scheduler=scheduler, print_lr=True), ValidationHandler(validator=evaluator, interval=validation_every_n_iters, epoch_level=False), StatsHandler(tag_name="train_loss", output_transform=lambda x: x["loss"]), TensorBoardStatsHandler(summary_writer=writer_train, log_dir=os.path.join(out_model_dir, "train"), tag_name="Loss", output_transform=lambda x: x["loss"], global_epoch_transform=lambda x: trainer.state.iteration), CheckpointSaver(save_dir=out_model_dir, save_dict={"net": net, "opt": opt}, save_final=True, save_interval=2, epoch_level=True, n_saved=config_info['output']['max_nr_models_saved']), ] if model_to_load is not None: train_handlers.append(CheckpointLoader(load_path=model_to_load, load_dict={"net": net, "opt": opt})) # define customized trainer class DynUNetTrainer(SupervisedTrainer): def _iteration(self, engine, batchdata): inputs, targets = self.prepare_batch(batchdata) inputs, targets = inputs.to(engine.state.device), targets.to(engine.state.device) def _compute_loss(preds, label): labels = [label] + [interpolate(label, pred.shape[2:]) for pred in preds[1:]] return sum([0.5 ** i * self.loss_function(p, l) for i, (p, l) in enumerate(zip(preds, labels))]) self.network.train() self.optimizer.zero_grad() if self.amp and self.scaler is not None: with torch.cuda.amp.autocast(): predictions = self.inferer(inputs, self.network) loss = _compute_loss(predictions, targets) self.scaler.scale(loss).backward() self.scaler.step(self.optimizer) self.scaler.update() else: predictions = self.inferer(inputs, self.network) loss = _compute_loss(predictions, targets).mean() loss.backward() self.optimizer.step() return {"image": inputs, "label": targets, "pred": predictions, "loss": loss.item()} trainer = DynUNetTrainer( device=current_device, max_epochs=config_info['training']['nr_train_epochs'], train_data_loader=train_loader, network=net, optimizer=opt, loss_function=loss_function, inferer=SimpleInferer(), post_transform=None, key_train_metric=None, train_handlers=train_handlers, amp=False, ) """ Run training """ print("*** Run training...") trainer.run() print("Done!")
def main_worker(args): # disable logging for processes except 0 on every node if args.local_rank != 0: f = open(os.devnull, "w") sys.stdout = sys.stderr = f if not os.path.exists(args.dir): raise FileNotFoundError(f"Missing directory {args.dir}") # initialize the distributed training process, every GPU runs in a process dist.init_process_group(backend="nccl", init_method="env://") total_start = time.time() train_transforms = Compose([ # load 4 Nifti images and stack them together LoadNiftid(keys=["image", "label"]), AsChannelFirstd(keys="image"), ConvertToMultiChannelBasedOnBratsClassesd(keys="label"), Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")), Orientationd(keys=["image", "label"], axcodes="RAS"), RandSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64], random_size=False), NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True), RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0), RandScaleIntensityd(keys="image", factors=0.1, prob=0.5), RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5), ToTensord(keys=["image", "label"]), ]) # create a training data loader train_ds = BratsCacheDataset( root_dir=args.dir, transform=train_transforms, section="training", num_workers=4, cache_rate=args.cache_rate, shuffle=True, ) train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) # validation transforms and dataset val_transforms = Compose([ LoadNiftid(keys=["image", "label"]), AsChannelFirstd(keys="image"), ConvertToMultiChannelBasedOnBratsClassesd(keys="label"), Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")), Orientationd(keys=["image", "label"], axcodes="RAS"), CenterSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64]), NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True), ToTensord(keys=["image", "label"]), ]) val_ds = BratsCacheDataset( root_dir=args.dir, transform=val_transforms, section="validation", num_workers=4, cache_rate=args.cache_rate, shuffle=False, ) val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if dist.get_rank() == 0: # Logging for TensorBoard writer = SummaryWriter(log_dir=args.log_dir) # create UNet, DiceLoss and Adam optimizer device = torch.device(f"cuda:{args.local_rank}") if args.network == "UNet": model = UNet( dimensions=3, in_channels=4, out_channels=3, channels=(16, 32, 64, 128, 256), strides=(2, 2, 2, 2), num_res_units=2, ).to(device) else: model = SegResNet(in_channels=4, out_channels=3, init_filters=16, dropout_prob=0.2).to(device) loss_function = DiceLoss(to_onehot_y=False, sigmoid=True, squared_pred=True) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5, amsgrad=True) # wrap the model with DistributedDataParallel module model = DistributedDataParallel(model, device_ids=[args.local_rank]) # start a typical PyTorch training total_epoch = args.epochs best_metric = -1000000 best_metric_epoch = -1 epoch_time = AverageMeter("Time", ":6.3f") progress = ProgressMeter(total_epoch, [epoch_time], prefix="Epoch: ") end = time.time() print(f"Time elapsed before training: {end-total_start}") for epoch in range(total_epoch): train_loss = train(train_loader, model, loss_function, optimizer, epoch, args, device) epoch_time.update(time.time() - end) if epoch % args.print_freq == 0: progress.display(epoch) if dist.get_rank() == 0: writer.add_scalar("Loss/train", train_loss, epoch) if (epoch + 1) % args.val_interval == 0: metric, metric_tc, metric_wt, metric_et = evaluate( model, val_loader, device) if dist.get_rank() == 0: writer.add_scalar("Mean Dice/val", metric, epoch) writer.add_scalar("Mean Dice TC/val", metric_tc, epoch) writer.add_scalar("Mean Dice WT/val", metric_wt, epoch) writer.add_scalar("Mean Dice ET/val", metric_et, epoch) if metric > best_metric: best_metric = metric best_metric_epoch = epoch + 1 print( f"current epoch: {epoch + 1} current mean dice: {metric:.4f}" f" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}" f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}" ) end = time.time() print(f"Time elapsed after epoch {epoch + 1} is {end - total_start}") if dist.get_rank() == 0: print( f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}" ) # all processes should see same parameters as they all start from same # random parameters and gradients are synchronized in backward passes, # therefore, saving it in one process is sufficient torch.save(model.state_dict(), "final_model.pth") writer.flush() dist.destroy_process_group()
def main(): opt = Options().parse() # monai.config.print_config() logging.basicConfig(stream=sys.stdout, level=logging.INFO) if opt.gpu_ids != '-1': num_gpus = len(opt.gpu_ids.split(',')) else: num_gpus = 0 print('number of GPU:', num_gpus) # Data loader creation # train images train_images = sorted(glob(os.path.join(opt.images_folder, 'train', 'image*.nii'))) train_segs = sorted(glob(os.path.join(opt.labels_folder, 'train', 'label*.nii'))) train_images_for_dice = sorted(glob(os.path.join(opt.images_folder, 'train', 'image*.nii'))) train_segs_for_dice = sorted(glob(os.path.join(opt.labels_folder, 'train', 'label*.nii'))) # validation images val_images = sorted(glob(os.path.join(opt.images_folder, 'val', 'image*.nii'))) val_segs = sorted(glob(os.path.join(opt.labels_folder, 'val', 'label*.nii'))) # test images test_images = sorted(glob(os.path.join(opt.images_folder, 'test', 'image*.nii'))) test_segs = sorted(glob(os.path.join(opt.labels_folder, 'test', 'label*.nii'))) # augment the data list for training for i in range(int(opt.increase_factor_data)): train_images.extend(train_images) train_segs.extend(train_segs) print('Number of training patches per epoch:', len(train_images)) print('Number of training images per epoch:', len(train_images_for_dice)) print('Number of validation images per epoch:', len(val_images)) print('Number of test images per epoch:', len(test_images)) # Creation of data directories for data_loader train_dicts = [{'image': image_name, 'label': label_name} for image_name, label_name in zip(train_images, train_segs)] train_dice_dicts = [{'image': image_name, 'label': label_name} for image_name, label_name in zip(train_images_for_dice, train_segs_for_dice)] val_dicts = [{'image': image_name, 'label': label_name} for image_name, label_name in zip(val_images, val_segs)] test_dicts = [{'image': image_name, 'label': label_name} for image_name, label_name in zip(test_images, test_segs)] # Transforms list # Need to concatenate multiple channels here if you want multichannel segmentation # Check other examples on Monai webpage. if opt.resolution is not None: train_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36), padding_mode="zeros"), Rand3DElasticd(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, sigma_range=(5, 8), magnitude_range=(100, 200), scale_range=(0.15, 0.15, 0.15), padding_mode="zeros"), RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1), RandGaussianNoised(keys=['image'], prob=0.1, mean=np.random.uniform(0, 0.5), std=np.random.uniform(0, 1)), RandShiftIntensityd(keys=['image'], offsets=np.random.uniform(0,0.3), prob=0.1), RandSpatialCropd(keys=['image', 'label'], roi_size=opt.patch_size, random_size=False), ToTensord(keys=['image', 'label']) ] val_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')), ToTensord(keys=['image', 'label']) ] else: train_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0), RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36), padding_mode="zeros"), RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36), padding_mode="zeros"), Rand3DElasticd(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=0.1, sigma_range=(5, 8), magnitude_range=(100, 200), scale_range=(0.15, 0.15, 0.15), padding_mode="zeros"), RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1), RandGaussianNoised(keys=['image'], prob=0.1, mean=np.random.uniform(0, 0.5), std=np.random.uniform(0, 1)), RandShiftIntensityd(keys=['image'], offsets=np.random.uniform(0,0.3), prob=0.1), RandSpatialCropd(keys=['image', 'label'], roi_size=opt.patch_size, random_size=False), ToTensord(keys=['image', 'label']) ] val_transforms = [ LoadImaged(keys=['image', 'label']), AddChanneld(keys=['image', 'label']), NormalizeIntensityd(keys=['image']), ScaleIntensityd(keys=['image']), ToTensord(keys=['image', 'label']) ] train_transforms = Compose(train_transforms) val_transforms = Compose(val_transforms) # create a training data loader check_train = monai.data.Dataset(data=train_dicts, transform=train_transforms) train_loader = DataLoader(check_train, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers, pin_memory=torch.cuda.is_available()) # create a training_dice data loader check_val = monai.data.Dataset(data=train_dice_dicts, transform=val_transforms) train_dice_loader = DataLoader(check_val, batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available()) # create a validation data loader check_val = monai.data.Dataset(data=val_dicts, transform=val_transforms) val_loader = DataLoader(check_val, batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available()) # create a validation data loader check_val = monai.data.Dataset(data=test_dicts, transform=val_transforms) test_loader = DataLoader(check_val, batch_size=1, num_workers=opt.workers, pin_memory=torch.cuda.is_available()) # # try to use all the available GPUs # devices = get_devices_spec(None) # build the network net = build_net() net.cuda() if num_gpus > 1: net = torch.nn.DataParallel(net) if opt.preload is not None: net.load_state_dict(torch.load(opt.preload)) dice_metric = DiceMetric(include_background=True, reduction="mean") post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)]) # loss_function = monai.losses.DiceLoss(sigmoid=True) # loss_function = monai.losses.TverskyLoss(sigmoid=True, alpha=0.3, beta=0.7) loss_function = monai.losses.DiceCELoss(sigmoid=True) optim = torch.optim.Adam(net.parameters(), lr=opt.lr) net_scheduler = get_scheduler(optim, opt) # start a typical PyTorch training val_interval = 1 best_metric = -1 best_metric_epoch = -1 epoch_loss_values = list() metric_values = list() writer = SummaryWriter() for epoch in range(opt.epochs): print("-" * 10) print(f"epoch {epoch + 1}/{opt.epochs}") net.train() epoch_loss = 0 step = 0 for batch_data in train_loader: step += 1 inputs, labels = batch_data["image"].cuda(), batch_data["label"].cuda() optim.zero_grad() outputs = net(inputs) loss = loss_function(outputs, labels) loss.backward() optim.step() epoch_loss += loss.item() epoch_len = len(check_train) // train_loader.batch_size print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}") writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step) epoch_loss /= step epoch_loss_values.append(epoch_loss) print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}") update_learning_rate(net_scheduler, optim) if (epoch + 1) % val_interval == 0: net.eval() with torch.no_grad(): def plot_dice(images_loader): metric_sum = 0.0 metric_count = 0 val_images = None val_labels = None val_outputs = None for data in images_loader: val_images, val_labels = data["image"].cuda(), data["label"].cuda() roi_size = opt.patch_size sw_batch_size = 4 val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, net) val_outputs = post_trans(val_outputs) value, _ = dice_metric(y_pred=val_outputs, y=val_labels) metric_count += len(value) metric_sum += value.item() * len(value) metric = metric_sum / metric_count metric_values.append(metric) return metric, val_images, val_labels, val_outputs metric, val_images, val_labels, val_outputs = plot_dice(val_loader) # Save best model if metric > best_metric: best_metric = metric best_metric_epoch = epoch + 1 torch.save(net.state_dict(), "best_metric_model.pth") print("saved new best metric model") metric_train, train_images, train_labels, train_outputs = plot_dice(train_dice_loader) metric_test, test_images, test_labels, test_outputs = plot_dice(test_loader) # Logger bar print( "current epoch: {} Training dice: {:.4f} Validation dice: {:.4f} Testing dice: {:.4f} Best Validation dice: {:.4f} at epoch {}".format( epoch + 1, metric_train, metric, metric_test, best_metric, best_metric_epoch ) ) writer.add_scalar("Mean_epoch_loss", epoch_loss, epoch + 1) writer.add_scalar("Testing_dice", metric_test, epoch + 1) writer.add_scalar("Training_dice", metric_train, epoch + 1) writer.add_scalar("Validation_dice", metric, epoch + 1) # plot the last model output as GIF image in TensorBoard with the corresponding image and label val_outputs = (val_outputs.sigmoid() >= 0.5).float() plot_2d_or_3d_image(val_images, epoch + 1, writer, index=0, tag="validation image") plot_2d_or_3d_image(val_labels, epoch + 1, writer, index=0, tag="validation label") plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=0, tag="validation inference") print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}") writer.close()
def run_inference(input_data, config_info): """ Pipeline to run inference with MONAI dynUNet model. The pipeline reads the input filenames, applies the required preprocessing and creates the pytorch dataloader; it then performs evaluation on each input file using a trained dynUNet model (random flipping augmentation is applied at inference). It uses the dynUNet model implemented in the MONAI framework (https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/nets/dynunet.py) which is inspired by the nnU-Net framework (https://arxiv.org/abs/1809.10486) Inference is performed in 2D slice-by-slice, all slices are then recombined together into the 3D volume. Args: input_data: str or list of strings, filenames of images to be processed config_info: dict, contains the configuration parameters to reload the trained model """ """ Read input and configuration parameters """ val_files = create_data_list_of_dictionaries(input_data) # print MONAI config information logging.basicConfig(stream=sys.stdout, level=logging.INFO) print("*** MONAI config: ") print_config() # print to log the parameter setups print("*** Network inference config: ") print(yaml.dump(config_info)) # inference params nr_out_channels = config_info['inference']['nr_out_channels'] spacing = config_info["inference"]["spacing"] prob_thr = config_info['inference']['probability_threshold'] model_to_load = config_info['inference']['model_to_load'] if not os.path.exists(model_to_load): raise FileNotFoundError('Trained model not found') patch_size = config_info["inference"]["inplane_size"] + [1] print("Considering patch size = {}".format(patch_size)) # set up either GPU or CPU usage if torch.cuda.is_available(): print("\n#### GPU INFORMATION ###") print("Using device number: {}, name: {}".format( torch.cuda.current_device(), torch.cuda.get_device_name())) current_device = torch.device("cuda:0") else: current_device = torch.device("cpu") print("Using device: {}".format(current_device)) """ Data Preparation """ print("*** Preparing data ... ") # data preprocessing for inference: # - convert data to right format [batch, channel, dim, dim, dim] # - resample to the training resolution in-plane (not along z) # - apply whitening # - convert to tensor val_transforms = Compose([ LoadNiftid(keys=["image"]), AddChanneld(keys=["image"]), InPlaneSpacingd( keys=["image"], pixdim=spacing, mode="bilinear", ), NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=True), ToTensord(keys=["image"]), ]) # create a validation data loader val_ds = Dataset(data=val_files, transform=val_transforms) val_loader = DataLoader(val_ds, batch_size=1, num_workers=config_info['device']['num_workers']) def prepare_batch(batchdata): assert isinstance(batchdata, dict), "prepare_batch expects dictionary input data." return ((batchdata["image"], batchdata["label"]) if "label" in batchdata else (batchdata["image"], None)) """ Network preparation """ print("*** Preparing network ... ") # automatically extracts the strides and kernels based on nnU-Net empirical rules spacings = spacing[:2] sizes = patch_size[:2] strides, kernels = [], [] while True: spacing_ratio = [sp / min(spacings) for sp in spacings] stride = [ 2 if ratio <= 2 and size >= 8 else 1 for (ratio, size) in zip(spacing_ratio, sizes) ] kernel = [3 if ratio <= 2 else 1 for ratio in spacing_ratio] if all(s == 1 for s in stride): break sizes = [i / j for i, j in zip(sizes, stride)] spacings = [i * j for i, j in zip(spacings, stride)] kernels.append(kernel) strides.append(stride) strides.insert(0, len(spacings) * [1]) kernels.append(len(spacings) * [3]) net = DynUNet(spatial_dims=2, in_channels=1, out_channels=nr_out_channels, kernel_size=kernels, strides=strides, upsample_kernel_size=strides[1:], norm_name="instance", deep_supervision=True, deep_supr_num=2, res_block=False).to(current_device) """ Set ignite evaluator to perform inference """ print("*** Preparing evaluator ... ") if nr_out_channels == 1: do_sigmoid = True do_softmax = False elif nr_out_channels > 1: do_sigmoid = False do_softmax = True else: raise Exception("incompatible number of output channels") print("Using sigmoid={} and softmax={} as final activation".format( do_sigmoid, do_softmax)) val_post_transforms = Compose([ Activationsd(keys="pred", sigmoid=do_sigmoid, softmax=do_softmax), AsDiscreted(keys="pred", argmax=True, threshold_values=True, logit_thresh=prob_thr), KeepLargestConnectedComponentd(keys="pred", applied_labels=1) ]) val_handlers = [ StatsHandler(output_transform=lambda x: None), CheckpointLoader(load_path=model_to_load, load_dict={"net": net}, map_location=torch.device('cpu')), SegmentationSaver( output_dir=config_info['output']['out_dir'], output_ext='.nii.gz', output_postfix=config_info['output']['out_postfix'], batch_transform=lambda batch: batch["image_meta_dict"], output_transform=lambda output: output["pred"], ), ] # Define customized evaluator class DynUNetEvaluator(SupervisedEvaluator): def _iteration(self, engine, batchdata): inputs, targets = self.prepare_batch(batchdata) inputs = inputs.to(engine.state.device) if targets is not None: targets = targets.to(engine.state.device) flip_inputs_1 = torch.flip(inputs, dims=(2, )) flip_inputs_2 = torch.flip(inputs, dims=(3, )) flip_inputs_3 = torch.flip(inputs, dims=(2, 3)) def _compute_pred(): pred = self.inferer(inputs, self.network) # use random flipping as data augmentation at inference flip_pred_1 = torch.flip(self.inferer(flip_inputs_1, self.network), dims=(2, )) flip_pred_2 = torch.flip(self.inferer(flip_inputs_2, self.network), dims=(3, )) flip_pred_3 = torch.flip(self.inferer(flip_inputs_3, self.network), dims=(2, 3)) return (pred + flip_pred_1 + flip_pred_2 + flip_pred_3) / 4 # execute forward computation self.network.eval() with torch.no_grad(): if self.amp: with torch.cuda.amp.autocast(): predictions = _compute_pred() else: predictions = _compute_pred() return {"image": inputs, "label": targets, "pred": predictions} evaluator = DynUNetEvaluator( device=current_device, val_data_loader=val_loader, network=net, prepare_batch=prepare_batch, inferer=SlidingWindowInferer2D(roi_size=patch_size, sw_batch_size=4, overlap=0.0), post_transform=val_post_transforms, val_handlers=val_handlers, amp=False, ) """ Run inference """ print("*** Running evaluator ... ") evaluator.run() print("Done!") return
def main_worker(args): # disable logging for processes except 0 on every node if args.local_rank != 0: f = open(os.devnull, "w") sys.stdout = sys.stderr = f if not os.path.exists(args.dir): raise FileNotFoundError(f"missing directory {args.dir}") # initialize the distributed training process, every GPU runs in a process dist.init_process_group(backend="nccl", init_method="env://") device = torch.device(f"cuda:{args.local_rank}") torch.cuda.set_device(device) # use amp to accelerate training scaler = torch.cuda.amp.GradScaler() torch.backends.cudnn.benchmark = True total_start = time.time() train_transforms = Compose([ # load 4 Nifti images and stack them together LoadImaged(keys=["image", "label"]), EnsureChannelFirstd(keys="image"), ConvertToMultiChannelBasedOnBratsClassesd(keys="label"), Orientationd(keys=["image", "label"], axcodes="RAS"), Spacingd( keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest"), ), EnsureTyped(keys=["image", "label"]), ToDeviced(keys=["image", "label"], device=device), RandSpatialCropd(keys=["image", "label"], roi_size=[224, 224, 144], random_size=False), RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0), RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1), RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2), NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True), RandScaleIntensityd(keys="image", factors=0.1, prob=0.5), RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5), ]) # create a training data loader train_ds = BratsCacheDataset( root_dir=args.dir, transform=train_transforms, section="training", num_workers=4, cache_rate=args.cache_rate, shuffle=True, ) # ThreadDataLoader can be faster if no IO operations when caching all the data in memory train_loader = ThreadDataLoader(train_ds, num_workers=0, batch_size=args.batch_size, shuffle=True) # validation transforms and dataset val_transforms = Compose([ LoadImaged(keys=["image", "label"]), EnsureChannelFirstd(keys="image"), ConvertToMultiChannelBasedOnBratsClassesd(keys="label"), Orientationd(keys=["image", "label"], axcodes="RAS"), Spacingd( keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest"), ), NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True), EnsureTyped(keys=["image", "label"]), ToDeviced(keys=["image", "label"], device=device), ]) val_ds = BratsCacheDataset( root_dir=args.dir, transform=val_transforms, section="validation", num_workers=4, cache_rate=args.cache_rate, shuffle=False, ) # ThreadDataLoader can be faster if no IO operations when caching all the data in memory val_loader = ThreadDataLoader(val_ds, num_workers=0, batch_size=args.batch_size, shuffle=False) # create network, loss function and optimizer if args.network == "SegResNet": model = SegResNet( blocks_down=[1, 2, 2, 4], blocks_up=[1, 1, 1], init_filters=16, in_channels=4, out_channels=3, dropout_prob=0.0, ).to(device) else: model = UNet( spatial_dims=3, in_channels=4, out_channels=3, channels=(16, 32, 64, 128, 256), strides=(2, 2, 2, 2), num_res_units=2, ).to(device) loss_function = DiceFocalLoss( smooth_nr=1e-5, smooth_dr=1e-5, squared_pred=True, to_onehot_y=False, sigmoid=True, batch=True, ) optimizer = Novograd(model.parameters(), lr=args.lr) lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs) # wrap the model with DistributedDataParallel module model = DistributedDataParallel(model, device_ids=[device]) dice_metric = DiceMetric(include_background=True, reduction="mean") dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch") post_trans = Compose( [EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)]) # start a typical PyTorch training best_metric = -1 best_metric_epoch = -1 print(f"time elapsed before training: {time.time() - total_start}") train_start = time.time() for epoch in range(args.epochs): epoch_start = time.time() print("-" * 10) print(f"epoch {epoch + 1}/{args.epochs}") epoch_loss = train(train_loader, model, loss_function, optimizer, lr_scheduler, scaler) print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}") if (epoch + 1) % args.val_interval == 0: metric, metric_tc, metric_wt, metric_et = evaluate( model, val_loader, dice_metric, dice_metric_batch, post_trans) if metric > best_metric: best_metric = metric best_metric_epoch = epoch + 1 if dist.get_rank() == 0: torch.save(model.state_dict(), "best_metric_model.pth") print( f"current epoch: {epoch + 1} current mean dice: {metric:.4f}" f" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}" f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}" ) print( f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}" ) print( f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}," f" total train time: {(time.time() - train_start):.4f}") dist.destroy_process_group()
def test_image_normalize_intensityd(self): key = 'img' normalizer = NormalizeIntensityd(keys=[key]) normalised = normalizer({key: self.imt}) expected = (self.imt - np.mean(self.imt)) / np.std(self.imt) self.assertTrue(np.allclose(normalised[key], expected))