def build(self): SEED = 42 data = pd.read_csv(self.data) ab = data.label ############################################ transforms = [ RescaleIntensity((0, 1)), RandomAffine(), transformss.ToTensor(), ] transform = Compose(transforms) ############################################# dataset_dir = self.dataset_dir dataset_dir = Path(dataset_dir) images_dir = dataset_dir labels_dir = dataset_dir image_paths = sorted(images_dir.glob('**/*.nii')) label_paths = sorted(labels_dir.glob('**/*.nii')) assert len(image_paths) == len(label_paths) # These two names are arbitrary MRI = 'features' BRAIN = 'targets' #split dataset into training and validation from catalyst.utils import split_dataframe_train_test train_image_paths, valid_image_paths = split_dataframe_train_test( image_paths, test_size=0.2, random_state=SEED) #training data subjects = [] i = 0 for (image_path, label_path) in zip(train_image_paths, label_paths): subject_dict = { MRI: torchio.Image(image_path, torchio.INTENSITY), BRAIN: ab[i], } i = i + 1 subject = torchio.Subject(subject_dict) subjects.append(subject) train_data = torchio.ImagesDataset(subjects) #validation data subjects = [] for (image_path, label_path) in zip(valid_image_paths, label_paths): subject_dict = { MRI: torchio.Image(image_path, torchio.INTENSITY), BRAIN: ab[i], } i = i + 1 subject = torchio.Subject(subject_dict) subjects.append(subject) test_data = torchio.ImagesDataset(subjects) return train_data, test_data
def __init__(self, images_dir, labels_dir): self.subjects = [] if (hp.in_class == 1) and (hp.out_class == 1): images_dir = Path(images_dir) self.image_paths = sorted(images_dir.glob(hp.fold_arch)) labels_dir = Path(labels_dir) self.label_paths = sorted(labels_dir.glob(hp.fold_arch)) for (image_path, label_path) in zip(self.image_paths, self.label_paths): subject = tio.Subject( source=tio.ScalarImage(image_path), label=tio.LabelMap(label_path), ) self.subjects.append(subject) else: images_dir = Path(images_dir) self.image_paths = sorted(images_dir.glob(hp.fold_arch)) artery_labels_dir = Path(labels_dir + '/artery') self.artery_label_paths = sorted( artery_labels_dir.glob(hp.fold_arch)) lung_labels_dir = Path(labels_dir + '/lung') self.lung_label_paths = sorted(lung_labels_dir.glob(hp.fold_arch)) trachea_labels_dir = Path(labels_dir + '/trachea') self.trachea_label_paths = sorted( trachea_labels_dir.glob(hp.fold_arch)) vein_labels_dir = Path(labels_dir + '/vein') self.vein_label_paths = sorted(vein_labels_dir.glob(hp.fold_arch)) for (image_path, artery_label_path, lung_label_path, trachea_label_path, vein_label_path) in zip( self.image_paths, self.artery_label_paths, self.lung_label_paths, self.trachea_label_paths, self.vein_label_paths): subject = tio.Subject( source=tio.ScalarImage(image_path), atery=tio.LabelMap(artery_label_path), lung=tio.LabelMap(lung_label_path), trachea=tio.LabelMap(trachea_label_path), vein=tio.LabelMap(vein_label_path), ) self.subjects.append(subject) self.transforms = self.transform() self.training_set = tio.SubjectsDataset(self.subjects, transform=self.transforms)
def get_dhcp(max_subjects=500): subjects = [] data_path = home + '/Sync-Exp/Data/DHCP/' all_seg = glob.glob(data_path + '**/*fusion_space-T2w_dseg.nii.gz', recursive=True) all_t2s = glob.glob(data_path + '**/*desc-restore_T2w.nii.gz', recursive=True) all_t1s = glob.glob(data_path + '**/*desc-restore_space-T2w_T1w.nii.gz', recursive=True) all_seg = all_seg[:max_subjects] for seg_file in all_seg: id_subject = seg_file.split('/')[6].split('_')[0:2] id_subject = id_subject[0] + '_' + id_subject[1] t2_file = [s for s in all_t2s if id_subject in s][0] t1_file = [s for s in all_t1s if id_subject in s][0] subject = tio.Subject( t2=tio.ScalarImage(t2_file), t1=tio.ScalarImage(t1_file), label=tio.LabelMap(seg_file), ) subjects.append(subject) print('DHCP Dataset size:', len(subjects), 'subjects') return subjects
def load_subject_(self, index): sample = self.patients[index % len(self.patients)] # load mr and turs file if it hasn't already been loaded if sample not in self.subjects: # print(f'loading patient {sample}') if self.load_mask: subject = torchio.Subject(mr=torchio.ScalarImage(sample + "/mr.mhd"), trus=torchio.ScalarImage(sample + "/trus.mhd"), mr_tree=torchio.LabelMap(sample + "/mr_tree.mhd")) else: subject = torchio.Subject(mr=torchio.ScalarImage(sample + "/mr.mhd"), trus=torchio.Image(sample + "/trus.mhd")) self.subjects[sample] = subject subject = self.subjects[sample] return sample, subject
def cache(dataset, resection_params, augment=True, caches_dir='/tmp/val_set_cache', num_workers=12): caches_dir = Path(caches_dir) wm_lesion_p = resection_params['wm_lesion_p'] clot_p = resection_params['clot_p'] shape = resection_params['shape'] texture = resection_params['texture'] augment_string = '_no_augmentation' if not augment else '' dir_name = f'wm_{wm_lesion_p}_clot_{clot_p}_{shape}_{texture}{augment_string}' cache_dir = caches_dir / dir_name image_dir = cache_dir / 'image' label_dir = cache_dir / 'label' if not cache_dir.is_dir(): print('Caching validation set') image_dir.mkdir(parents=True) label_dir.mkdir(parents=True) loader = torch.utils.data.DataLoader( dataset, num_workers=num_workers, collate_fn=lambda x: x[0], ) for subject in tqdm(loader): image_path = image_dir / subject.image.path.name label_path = label_dir / subject.image.path.name # label has no path because it was created not loaded subject.image.save(image_path) subject.label.save(label_path) subjects = [] for im_path, label_path in zip(sglob(image_dir), sglob(label_dir)): subject = tio.Subject( image=tio.ScalarImage(im_path), label=tio.LabelMap(label_path), ) subjects.append(subject) return tio.SubjectsDataset(subjects)
def train_test_val_split(): ground_truths = Path("IXI-T1/Actual_Images") ground_paths = sorted(ground_truths.glob('*.nii.gz')) compressed_dirs = [ sorted(Path((os.path.join("IXI-T1", comp))).glob('*.nii.gz')) for comp in os.listdir("IXI-T1") if "Compressed" in comp ] # compressed_images = Path("IXI-T1/Compressed_3x3x1.2") # compressed_paths = sorted(compressed_images.glob('*.nii.gz')) training_subjects = [] test_subjects = [] validation_subjects = [] for compressed_paths in compressed_dirs: subjects = [] for gt, comp in zip(ground_paths, compressed_paths): subject = tio.Subject( ground_truth=tio.ScalarImage(gt), compressed=tio.ScalarImage(comp), ) subjects.append(subject) train_split, test_split = train_test_split(subjects, test_size=0.3) test_split, validation_split = train_test_split(test_split, test_size=0.2) training_subjects += train_split validation_subjects += validation_split test_subjects += test_split return training_subjects, test_subjects, validation_subjects
def add_evaluation_labels(subjects: Sequence[tio.Subject]): for subject in subjects: transform = subject.get_composed_history() label_transform_types = [LabelTransform, CopyProperty, RenameProperty, ConcatenateImages] label_transform = filter_transform(transform, include_types=label_transform_types) evaluation_transform = label_transform.inverse(warn=False) if 'y_pred' in subject: pred_subject = tio.Subject({'y': subject['y_pred']}) y_pred_eval = evaluation_transform(pred_subject).get_first_image() subject.add_image(y_pred_eval, 'y_pred_eval') if 'y' in subject: target_subject = tio.Subject({'y': subject['y']}) y_eval = evaluation_transform(target_subject).get_first_image() subject.add_image(y_eval, 'y_eval')
def test_different_spaces(self): t1 = self.sample_subject.t1 label = tio.Resample(2)(self.sample_subject.label) new_subject = tio.Subject(t1=t1, label=label) with self.assertRaises(RuntimeError): tio.RandomAffine()(new_subject) tio.RandomAffine(check_shape=False)(new_subject)
def __getitem__(self, index): # update the seed to avoid workers sample the same augmentation parameters np.random.seed(datetime.datetime.now().second + datetime.datetime.now().microsecond) # load the nifti images input, _ = load_nifti_img(self.image_filenames[index], dtype=np.int16) target, _ = load_nifti_img(self.target_filenames[index], dtype=np.uint8) #check_exceptions(input, target) if self.transform != None: sub = tio.Subject( input=tio.ScalarImage(tensor=input[None, :, :, :]), target=tio.LabelMap(tensor=target[None, :, :, :])) sub = self.transform(sub) input = np.array(sub['input'])[0, ...] target = np.array(sub['target'])[0, ...] if self.hot == 1: target = self._toEvaluationOneHot(target) input = torch.from_numpy(input[None, :, :, :]).float() target = torch.from_numpy(target).long() # print(target.shape if self.mode == 'test': pid = torch.from_numpy(np.array([self.data_splits[index]])) return pid, input, target return input, target
def visualize_training(self, visuals, affine, epoch, index): """ Save training image for visualization. :param affine: :param visuals: images. :type visuals: dict :param epoch: epoch :type epoch: int :param index: index :type index: float """ ct = visuals['ct'].cpu().transpose_(2, 4) visuals['segmentation_mask'] = torch.unsqueeze(visuals['segmentation_mask'].cpu(), 1).transpose_(2, 4) visuals['fake_segmentation_mask'] = visuals['fake_segmentation_mask'].cpu().transpose_(2, 4) segmentation_mask = visuals['segmentation_mask'] fake_segmentation_mask = visuals['fake_segmentation_mask'].argmax(dim=1, keepdim=True) for i in range(ct.shape[0]): subject = tio.Subject( ct=tio.ScalarImage(tensor=ct[i], affine=affine[i]), segmentation_mask=tio.LabelMap(tensor=segmentation_mask[i], affine=affine[i]), fake_segmentation_mask=tio.LabelMap(tensor=fake_segmentation_mask[i], affine=affine[i]) ) save_path = os.path.join(self.expr_dir, "training_visuals") save_path = os.path.join(save_path, 'cycle_%02d_%04d_%02d.png' % (epoch, index, i)) subject.plot(show=False, output_path=save_path) plt.close('all')
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 test_different_interpolation(self): def model_probs(subject): subject = copy.deepcopy(subject) subject.im.set_data(torch.rand_like(subject.im.data)) return subject def model_label(subject): subject = model_probs(subject) subject.im.set_data(torch.bernoulli(subject.im.data)) return subject transform = tio.RandomAffine(image_interpolation='bspline') subject = copy.deepcopy(self.sample_subject) tensor = (torch.rand(1, 20, 20, 20) > 0.5).float() # 0s and 1s subject = tio.Subject(im=tio.ScalarImage(tensor=tensor)) transformed = transform(subject) assert transformed.im.data.min() < 0 assert transformed.im.data.max() > 1 subject_probs = model_probs(transformed) transformed_back = subject_probs.apply_inverse_transform() assert transformed_back.im.data.min() < 0 assert transformed_back.im.data.max() > 1 transformed_back_linear = subject_probs.apply_inverse_transform( image_interpolation='linear', ) assert transformed_back_linear.im.data.min() >= 0 assert transformed_back_linear.im.data.max() <= 1 subject_label = model_label(transformed) transformed_back = subject_label.apply_inverse_transform() assert transformed_back.im.data.min() < 0 assert transformed_back.im.data.max() > 1 transformed_back_linear = subject_label.apply_inverse_transform( image_interpolation='nearest', ) assert transformed_back_linear.im.data.unique().tolist() == [0, 1]
def get_torchio_subject(item_data, item_label): # converting to inputs to torchio img t1 = torchio.ScalarImage(tensor=item_data) label = torchio.LabelMap(tensor=item_label) subject = torchio.Subject(t1=t1, label=label) return subject
def load_kidney_seg(data_shape, batch=3, workers=4, transform=None): #take input transform and apply it after clip, normalization, resize if transform == None: transform = tio.RandomFlip(p=0.) #preprocess all clippy = Lambda(lambda x: torch.clip(x, -80, 300), types_to_apply=[tio.INTENSITY]) normal = RescaleIntensity((0., 1.)) resize = Lambda(lambda x: torch.squeeze( interpolate(torch.unsqueeze(x, dim=0), data_shape), dim=0)) rounding = Lambda(lambda x: torch.round(x), types_to_apply=[tio.LABEL]) transform = tio.Compose([clippy, normal, resize, rounding, transform]) subject_list = [] for i in range(210): pt_image = ("data/case_{:05d}/imaging.nii.gz".format(i)) pt_label = ("data/case_{:05d}/segmentation.nii.gz".format(i)) subject_list.append( tio.Subject(img=tio.ScalarImage(pt_image), label=tio.LabelMap(pt_label))) dataset = tio.SubjectsDataset(subject_list, transform=transform) return DataLoader(dataset, num_workers=workers, batch_size=batch, pin_memory=True)
def test_inconsistent_spatial_shape(self): subject = tio.Subject( a=tio.ScalarImage(tensor=torch.rand(1, 3, 3, 4)), b=tio.ScalarImage(tensor=torch.rand(2, 2, 3, 4)), ) with self.assertRaises(RuntimeError): subject.spatial_shape
def get_subjects(path, structures, transform): """ Browse the path folder to build a dataset. Folder must contains the subjects with the CT and masks. :param path: root folder. :type path: str :param structures: list of structures. :type structures: list[str] :param transform: transforms to be applied. :type transform: :class:`tio.transforms.Transform` :return: Base TorchIO dataset. :rtype: :class:`tio.SubjectsDataset` """ subject_ids = os.listdir(path) subjects = [] for subject_id in subject_ids: ct_path = os.path.join(path, subject_id, 'ct.nii') structures_path_dict = {k: os.path.join(path, subject_id, k + '.nii') for k in structures} subject = tio.Subject( ct=tio.ScalarImage(ct_path), ) label_map = torch.zeros(subject["ct"].shape, dtype=torch.long) for i, (k, v) in enumerate(structures_path_dict.items()): label_map += tio.LabelMap(v).data * (i + 1) label_map[label_map > len(structures)] = 0 subject.add_image(tio.LabelMap(tensor=label_map, affine=subject["ct"].affine), 'label_map') subjects.append(subject) return tio.SubjectsDataset(subjects, transform=transform)
def get_torchio_dataset(inputs, targets, transform): """ Function creates a torchio.SubjectsDataset from inputs and targets lists and applies transform to that dataset Arguments: * inputs (list): list of paths to MR images * targets (list): list of paths to ground truth segmentation of MR images * transform (False/torchio.transforms): transformations which will be applied to MR images and ground truth segmentation of MR images (but not all of them) Output: * datasets (torchio.SubjectsDataset): it's kind of torchio list of torchio.data.subject.Subject entities """ subjects = [] for (image_path, label_path) in zip(inputs, targets ): subject_dict = { 'MRI' : torchio.Image(image_path, torchio.INTENSITY), 'LABEL': torchio.Image(label_path, torchio.LABEL), #intensity transformations won't be applied to torchio.LABEL } subject = torchio.Subject(subject_dict) subjects.append(subject) if transform: dataset = torchio.SubjectsDataset(subjects, transform = transform) elif not transform: dataset = torchio.SubjectsDataset(subjects) return dataset
def do_subject(image_paths, label_paths): for (image_path, label_path) in zip(image_paths, label_paths): subject = tio.Subject( pred=tio.ScalarImage(image_path), gt=tio.LabelMap(label_path), ) subjects.append(subject)
def test_repr(self): subject = tio.Subject( t1=tio.ScalarImage(self.get_image_path('repr_test')), ) assert 'memory' not in repr(subject['t1']) subject.load() assert 'memory' in repr(subject['t1'])
def infer_with_patches(self, model_inference_function, features): # This function infers using multiple patches, fusing corresponding outputs # model_inference_function is a list to suport recursive calls to similar function subject_dict = {} for i in range(0, features.shape[1]): # 0 is batch subject_dict[str(i)] = torchio.Image(tensor=features[:, i, :, :, :], type=torchio.INTENSITY) grid_sampler = torchio.inference.GridSampler( torchio.Subject(subject_dict), self.psize) patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=1) aggregator = torchio.inference.GridAggregator(grid_sampler) for patches_batch in patch_loader: # concatenate the different modalities into a tensor image = torch.cat([ patches_batch[str(i)][torchio.DATA] for i in range(0, features.shape[1]) ], dim=1) locations = patches_batch[ torchio.LOCATION] # get location of patch pred_mask = model_inference_function[0]( model_inference_function=model_inference_function[1:], features=image) aggregator.add_batch(pred_mask, locations) output = aggregator.get_output_tensor() # this is the final mask output = output.unsqueeze( 0) # increasing the number of dimension of the mask return output
def test_no_sample(self): with tempfile.NamedTemporaryFile(delete=False) as f: input_dict = {'image': tio.ScalarImage(f.name)} subject = tio.Subject(input_dict) with self.assertRaises(RuntimeError): with self.assertWarns(UserWarning): tio.RandomFlip()(subject)
def get_subject_with_labels(self, labels): return tio.Subject( label=tio.LabelMap( self.get_image_path( 'label_multi', labels=labels ) ) )
def get_subject_with_partial_volume_label_map(self, components=1): """Return a subject with a partial-volume label map.""" return tio.Subject( t1=tio.ScalarImage(self.get_image_path('t1_d'), ), label=tio.LabelMap( self.get_image_path('label_d2', binary=False, components=components)), )
def load_pretrain_datasets(data_shape, batch=3, workers=4, transform=None): data_path = '/home/mitch/Data/MSD/' directories = sorted(glob.glob(data_path + '*/')) loaders = [] #var to store dataloader for each task datasets = [] #store dataset objects before turning into loaders if transform == None: transform = tio.RandomFlip(p=0.) #preprocess all clippy = Lambda(lambda x: torch.clip(x, -80, 300), types_to_apply=[tio.INTENSITY]) normal = RescaleIntensity((0., 1.)) resize = Lambda(lambda x: torch.squeeze( interpolate(torch.unsqueeze(x, dim=0), data_shape), dim=0)) rounding = Lambda(lambda x: torch.round(x), types_to_apply=[tio.LABEL]) transform = tio.Compose([clippy, normal, resize, rounding, transform]) #deal with weird shapes braintransform = Lambda(lambda x: torch.unsqueeze(x[:, :, :, 2], dim=0), types_to_apply=[tio.INTENSITY]) braintransform = tio.Compose([braintransform, transform]) prostatetransform = Lambda(lambda x: torch.unsqueeze(x[:, :, :, 1], dim=0), types_to_apply=[tio.INTENSITY]) prostatetransform = tio.Compose([prostatetransform, transform]) for i, directory in enumerate(directories): images = sorted(glob.glob(directory + 'imagesTr/*')) segs = sorted(glob.glob(directory + 'labelsTr/*')) subject_list = [] for image, seg in zip(images, segs): subject_list.append( tio.Subject(img=tio.ScalarImage(image), label=tio.LabelMap(seg))) #handle special cases if i == 0: datasets.append( tio.SubjectsDataset(subject_list, transform=braintransform)) elif i == 4: datasets.append( tio.SubjectsDataset(subject_list, transform=prostatetransform)) else: datasets.append( tio.SubjectsDataset(subject_list, transform=transform)) loaders.append( DataLoader(datasets[-1], num_workers=workers, batch_size=batch, pin_memory=True)) return loaders
def test_incosistent_shape(self): # https://github.com/fepegar/torchio/issues/234#issuecomment-675029767 subject = torchio.Subject( im1=torchio.ScalarImage(tensor=torch.rand(1, 4, 5, 6)), im2=torchio.ScalarImage(tensor=torch.rand(2, 4, 5, 6)), ) patch_size = 2 sampler = torchio.data.WeightedSampler(patch_size, 'im1') next(sampler(subject))
def test_inconsistent_shape(self): # https://github.com/fepegar/torchio/issues/234#issuecomment-675029767 sample = torchio.Subject( im1=torchio.ScalarImage(tensor=torch.rand(2, 4, 5, 6)), im2=torchio.LabelMap(tensor=torch.rand(1, 4, 5, 6)), ) patch_size = 2 sampler = LabelSampler(patch_size, 'im2') next(sampler(sample))
def setUp(self): """Set up test fixtures, if any.""" self.dir = Path(tempfile.gettempdir()) / '.torchio_tests' self.dir.mkdir(exist_ok=True) random.seed(42) np.random.seed(42) registration_matrix = np.array([ [1, 0, 0, 10], [0, 1, 0, 0], [0, 0, 1.2, 0], [0, 0, 0, 1] ]) subject_a = tio.Subject( t1=tio.ScalarImage(self.get_image_path('t1_a')), ) subject_b = tio.Subject( t1=tio.ScalarImage(self.get_image_path('t1_b')), label=tio.LabelMap(self.get_image_path('label_b', binary=True)), ) subject_c = tio.Subject( label=tio.LabelMap(self.get_image_path('label_c', binary=True)), ) subject_d = tio.Subject( t1=tio.ScalarImage( self.get_image_path('t1_d'), pre_affine=registration_matrix, ), t2=tio.ScalarImage(self.get_image_path('t2_d')), label=tio.LabelMap(self.get_image_path('label_d', binary=True)), ) subject_a4 = tio.Subject( t1=tio.ScalarImage(self.get_image_path('t1_a'), components=2), ) self.subjects_list = [ subject_a, subject_a4, subject_b, subject_c, subject_d, ] self.dataset = tio.SubjectsDataset(self.subjects_list) self.sample_subject = self.dataset[-1] # subject_d
def test_empty_map(self): # https://github.com/fepegar/torchio/issues/392 im = tio.ScalarImage(tensor=torch.rand(1, 6, 6, 6)) label = torch.zeros(1, 6, 6, 6) label[..., 0] = 1 # voxels far from center label_im = tio.LabelMap(tensor=label) subject = tio.Subject(image=im, label=label_im) sampler = tio.LabelSampler(4) with self.assertRaises(RuntimeError): next(sampler(subject))
def get_sample(self, image_shape): t1 = torch.rand(*image_shape) prob = torch.zeros_like(t1) prob[0, 3, 3, 3] = 1 subject = torchio.Subject( t1=torchio.ScalarImage(tensor=t1), prob=torchio.ScalarImage(tensor=prob), ) subject = torchio.SubjectsDataset([subject])[0] return subject
def test_label_probabilities(self): labels = torch.Tensor((0, 0, 1, 1, 2, 1, 0)).reshape(1, 1, 1, -1) subject = torchio.Subject(label=torchio.Image(tensor=labels, type=torchio.LABEL), ) sample = torchio.SubjectsDataset([subject])[0] probs_dict = {0: 0, 1: 50, 2: 25, 3: 25} sampler = LabelSampler(5, 'label', label_probabilities=probs_dict) probabilities = sampler.get_probability_map(sample) fixture = torch.Tensor((0, 0, 2 / 12, 2 / 12, 3 / 12, 2 / 12, 0)) assert torch.all(probabilities.squeeze().eq(fixture))