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))
