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 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 test_incosistent_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_persistent_bounds_params(self):
# https://github.com/fepegar/torchio/issues/757
shape = (1, 5, 5, 5)
mask_a = np.zeros(shape)
mask_a[0, 2, 2, 2] = 1
mask_b = mask_a.copy()
mask_b[0, 1:4, 1:4, 1:4] = 1
tensor = np.ones(shape)
image_a = tio.ScalarImage(tensor=tensor)
mask_a = tio.LabelMap(tensor=mask_a)
subject_a = tio.Subject(image=image_a, mask=mask_a)
image_b = tio.ScalarImage(tensor=tensor)
mask_b = tio.LabelMap(tensor=mask_b)
subject_b = tio.Subject(image=image_b, mask=mask_b)
crop = tio.CropOrPad(mask_name='mask')
for _ in range(2):
shape_a = crop(subject_a).image.shape
shape_b = crop(subject_b).image.shape
assert shape_a != shape_b
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_inconsistent_shape_subject(self):
"""Return a subject containing images of different shape."""
subject = tio.Subject(
t1=tio.ScalarImage(self.get_image_path('t1_inc')),
t2=tio.ScalarImage(
self.get_image_path('t2_inc', shape=(10, 20, 31))),
label=tio.LabelMap(
self.get_image_path(
'label_inc',
shape=(8, 17, 25),
binary=True,
), ),
label2=tio.LabelMap(
self.get_image_path(
'label2_inc',
shape=(18, 17, 25),
binary=True,
), ),
)
return subject
def test_same_affine(self):
image = tio.ScalarImage(tensor=torch.rand(2, 2, 2, 2))
mask = tio.LabelMap(tensor=torch.rand(2, 2, 2, 2))
mask.affine *= 1.1
subject = tio.Subject(t1=image, mask=mask)
transform = tio.CopyAffine('t1')
transformed = transform(subject)
self.assertTensorEqual(
transformed['t1'].affine,
transformed['mask'].affine,
)
def get_torchio_registration_subject(item_data, item_label):
log_msg = f'get_torchio_registration_subject_0: {item_data.shape}, {item_label.shape}'
# print(log_msg)
logging.debug(log_msg)
# converting to inputs to torchio img
if not torch.is_tensor(item_data):
item_data = torch.tensor(item_data)
item_label = torch.tensor(item_label)
t1_input = torch.unsqueeze(item_data[0], 0)
t2_input = torch.unsqueeze(item_data[1], 0)
# print(f"{t1_input.shape}", {t2_input.shape})
t1 = torchio.ScalarImage(tensor=t1_input)
t2 = torchio.LabelMap(tensor=t2_input)
label = torchio.LabelMap(tensor=item_label)
subject = torchio.Subject(t1=t1, t2=t2, label=label)
return subject
def test_2d(self):
# https://github.com/fepegar/torchio/issues/434
image = np.random.rand(1, 16, 16, 1)
mask = np.zeros_like(image, dtype=bool)
mask[0, 7, 0] = True
subject = tio.Subject(
image=tio.ScalarImage(tensor=image),
mask=tio.LabelMap(tensor=mask),
)
transform = tio.CropOrPad((12, 12, 1), mask_name='mask')
transformed = transform(subject)
assert transformed.shape == (1, 12, 12, 1)
def __call__(self, inputVolumeNode, outputVolumeNode):
image = su.PullVolumeFromSlicer(inputVolumeNode)
data, affine = torchio.io.sitk_to_nib(image)
tensor = torch.from_numpy(data.astype(
np.float32)) # why do I need this? Open a TorchIO issue?
if inputVolumeNode.IsA('vtkMRMLScalarVolumeNode'):
image = torchio.ScalarImage(tensor=tensor, affine=affine)
elif inputVolumeNode.IsA('vtkMRMLLabelMapVolumeNode'):
image = torchio.LabelMap(tensor=tensor, affine=affine)
transformed = self.getTransform()(image)
image = torchio.io.nib_to_sitk(transformed.data, transformed.affine)
su.PushVolumeToSlicer(image, targetNode=outputVolumeNode)
return outputVolumeNode
def spatial(self, x, y):
trans = tio.OneOf({
tio.RandomAffine(scales=0.1, degrees=(20, 0, 0), translation=0):
0.5,
tio.RandomElasticDeformation(max_displacement=(0, 7.5, 7.5)):
0.5
})
image = torchio.ScalarImage(tensor=x[None, None, ...])
mask = torchio.LabelMap(tensor=y[None, None, ...])
sub = torchio.Subject({'image': image, 'mask': mask})
sub = trans(sub)
return sub.image.data[0, 0, ...], sub.mask.data[0, 0, ...]
def get_subjects_list_from_dir(dataset_dir):
dataset_dir = Path(dataset_dir)
mni_dir = dataset_dir / 'mni'
resection_dir = dataset_dir / 'resection'
noise_paths = sglob(resection_dir, '*noise*')
subjects_list = []
for noise_path in noise_paths:
stem = noise_path.stem.split('_noise')[0]
image_path = mni_dir / f'{stem}_on_mni.nii.gz'
gml_path = resection_dir / f'{stem}_gray_matter_left_seg.nii.gz'
gmr_path = resection_dir / f'{stem}_gray_matter_right_seg.nii.gz'
rl_path = resection_dir / f'{stem}_resectable_left_seg.nii.gz'
rr_path = resection_dir / f'{stem}_resectable_right_seg.nii.gz'
subject = tio.Subject(
image=tio.ScalarImage(image_path),
resection_noise=tio.ScalarImage(noise_path),
resection_gray_matter_left=tio.LabelMap(gml_path),
resection_gray_matter_right=tio.LabelMap(gmr_path),
resection_resectable_left=tio.LabelMap(rl_path),
resection_resectable_right=tio.LabelMap(rr_path),
)
subjects_list.append(subject)
return subjects_list
def test_mask_specified_label_small(self):
def to_image(*numbers):
return torch.as_tensor(numbers).reshape(1, 1, 1, len(numbers))
image_tensor = to_image(1, 6, 7, 3, 0)
label_tensor = to_image(0, 1, 2, 3, 4)
mask_labels = [1, 2]
subject = tio.Subject(
image=tio.ScalarImage(tensor=image_tensor),
label=tio.LabelMap(tensor=label_tensor),
)
transform = tio.Mask(masking_method='label', labels=mask_labels)
transformed = transform(subject)
masked_list = transformed.image.data.flatten().tolist()
assert masked_list == [0, 6, 7, 0, 0]
def get_real_resection_subjects(dataset_dir):
dataset_dir = Path(dataset_dir)
image_dir = dataset_dir / 'image'
label_dir = dataset_dir / 'label'
image_paths = sglob(image_dir)
label_paths = sglob(label_dir)
assert len(image_paths) == len(label_paths)
subjects = []
for image_path, label_path in zip(image_paths, label_paths):
subject = tio.Subject(
image=tio.ScalarImage(image_path),
label=tio.LabelMap(label_path),
)
subjects.append(subject)
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 test_multichannel_label_sampler(self):
sample = torchio.Subject(label=torchio.LabelMap(
tensor=torch.tensor([[[[1, 1]]], [[[0, 1]]]])))
patch_size = 1
sampler = LabelSampler(patch_size,
'label',
label_probabilities={
0: 1,
1: 1
})
# There are 2 voxels in the image, channels have same probabilities,
# 1st voxel has probability 0.5 * 0.5 + 0 * 0.5 of being chosen while
# 2nd voxel has probability 0.5 * 0.5 + 1 * 0.5 of being chosen.
probabilities = sampler.get_probability_map(sample)
fixture = torch.Tensor((1 / 4, 3 / 4))
assert torch.all(probabilities.squeeze().eq(fixture))
def __call__(self, inputVolumeNode, outputVolumeNode):
image = su.PullVolumeFromSlicer(inputVolumeNode)
tensor, affine = torchio.io.sitk_to_nib(image)
if inputVolumeNode.IsA('vtkMRMLScalarVolumeNode'):
image = torchio.ScalarImage(tensor=tensor, affine=affine)
elif inputVolumeNode.IsA('vtkMRMLLabelMapVolumeNode'):
image = torchio.LabelMap(tensor=tensor, affine=affine)
subject = torchio.Subject(image=image) # to get transform history
transformed = self.getTransform()(subject)
deterministicApplied = transformed.get_applied_transforms()[0]
logging.info(f'Applied transform: {deterministicApplied}')
transformedImage = transformed.image
image = torchio.io.nib_to_sitk(transformedImage.data,
transformedImage.affine)
su.PushVolumeToSlicer(image, targetNode=outputVolumeNode)
return outputVolumeNode
def test_get_subjects(self):
ct = tio.ScalarImage(tensor=torch.rand(1, 3, 3, 2))
structure = tio.LabelMap(
tensor=torch.ones((1, 3, 3, 2), dtype=torch.uint8))
subject_1_dir = "tests/test_data/subjects/subject_1"
os.makedirs(subject_1_dir, exist_ok=True)
ct.save(os.path.join(subject_1_dir, "ct.nii"))
structure.save(os.path.join(subject_1_dir, "structure.nii"))
transform = tio.Compose(
[tio.ToCanonical(),
tio.RescaleIntensity(1, (1, 99.0))])
subject_dataset = get_subjects(os.path.dirname(subject_1_dir),
structures=["structure"],
transform=transform)
self.assertEqual(len(subject_dataset), 1)
shutil.rmtree(os.path.dirname(subject_1_dir), ignore_errors=True)
def __getitem__(self, item_index):
#lazily open file
self.openFileIfNotOpen()
index = self.used_split[item_index]
#load from hdf5 file
image = self.file["images_" + 'train'][index, ...]
if self.hasMasks:
labels = self.labelFile["masks_" + 'train'][index, ...]
#Prepare data depeinding on soft/hard augmentation scheme
n_classes = self.n_classes
if self.transform != None and self.mode=="train":
sub = tio.Subject(image = tio.ScalarImage(tensor = image[None, :,:,:]),
labels = tio.LabelMap(tensor = labels[None, :,:,:]))
sub = self.transform(sub)
image = np.array(sub['image'])[0,...]
labels = np.array(sub['labels'])[0,...]
if 1 == self.hot:
labels = self._toEvaluationOneHot(labels)
if self.hasMasks:
labels = np.transpose(labels, (3, 0, 1, 2)) # bring into NCWH format
image = torch.from_numpy(image)
image = image.expand(1,-1,-1,-1)
if 1 == 1:
if self.hasMasks:
#labels = labels[:, 0:32, 0:32, 0:32]
labels = torch.from_numpy(labels).long()
if self.hasMasks:
return image, labels
else:
return image
def SubjectsDataset():
images_dir = dataset_dir / 'image'
labels_dir = dataset_dir / 'label'
image_paths = sorted(images_dir.glob('*.nii.gz'))
label_paths = sorted(labels_dir.glob('*.nii.gz'))
assert len(image_paths) == len(label_paths)
subjects = []
for (image_path, label_path) in zip(image_paths, label_paths):
subject = tio.Subject(
mri=tio.ScalarImage(image_path),
brain=tio.LabelMap(label_path),
)
subjects.append(subject)
# subjects = np.array(subjects)
dataset = tio.SubjectsDataset(subjects)
print('Dataset size:', len(dataset),
'subjects') ## => Dataset size : 566 subjects
return dataset, subjects
def __init__(self, mask, ct_path=None, ds_cts=None, structures=None):
if ct_path is None and ds_cts is None:
raise ValueError('At least ct_path should be provided')
self.masks = mask if not isinstance(mask, str) else tio.LabelMap(mask)
self.structures = structures
self.masks_itk = self.masks.as_sitk()
self.transform = tio.OneHot()
self.one_hot_masks = self.transform(self.masks)
self.n_masks = self.one_hot_masks.shape[0] - 1
self.ct_files = natsorted([
os.path.join(ct_path, ct) for ct in os.listdir(ct_path)
if ct.endswith("dcm")
])
self.ds_ct = ds_cts or [
dcmread(ct_file, force=True) for ct_file in self.ct_files
]
self.ds_ct.reverse()
def __init__(self, path, images=None, labels=None, transforms=None):
self.transforms = transforms
self.subjects = []
self.images = images
self.labels = labels
self.subject_folder_names = os.listdir(path)
self.subject_folders = [f"{path}/{folder}/" for folder in self.subject_folder_names]
for subject_folder in self.subject_folders:
subject_files = os.listdir(subject_folder)
subject_data = {}
attributes_file = "attributes.json"
if attributes_file in subject_files:
with open(f"{subject_folder}/{attributes_file}") as f:
subject_data = json.load(f)
subject_files.remove(attributes_file)
file_map = {file[:file.find(".")]: file for file in subject_files}
missing_name = False
all_names = []
if images is not None:
all_names += images
if labels is not None:
all_names += labels
for name in all_names:
if name not in file_map:
missing_name = True
if missing_name:
continue
if images is not None:
for name in images:
subject_data[name] = tio.ScalarImage(subject_folder + file_map[name])
if labels is not None:
for name in labels:
subject_data[name] = tio.LabelMap(subject_folder + file_map[name])
self.subjects.append(tio.Subject(**subject_data))
self.subject_dataset = tio.SubjectsDataset(self.subjects, transform=transforms)
def get_pseudo_loader(
threshold,
percentile,
metric,
summary_path,
dataset_name,
num_workers,
batch_size=2,
remove_zero_volume=False,
):
subjects = []
subject_ids = get_certain_subjects(
threshold,
percentile,
metric,
summary_path,
remove_zero_volume=remove_zero_volume,
)
dataset_dir = Path('/home/fernando/datasets/real/') / dataset_name
assert dataset_dir.is_dir()
image_dir = dataset_dir / 'image'
label_dir = dataset_dir / 'label'
for subject_id in subject_ids:
image_path = list(image_dir.glob(f'{subject_id}_*'))[0]
label_path = list(label_dir.glob(f'{subject_id}_*'))[0]
subject = tio.Subject(
image=tio.ScalarImage(image_path),
label=tio.LabelMap(label_path),
)
subjects.append(subject)
transform = get_train_transform(get_landmarks_path())
dataset = tio.SubjectsDataset(subjects, transform=transform)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
pin_memory=True,
shuffle=True,
num_workers=num_workers,
)
return loader
def main(image_dir, label_dir, checkpoint_path, output_dir, landmarks_path,
df_path, batch_size, num_workers, multi_gpu):
import torch
import torchio as tio
import models
import datasets
import engine
import utils
fps = get_paths(image_dir)
lfps = get_paths(label_dir)
assert len(fps) == len(lfps)
# key must be 'image' as in get_test_transform
subjects = [
tio.Subject(image=tio.ScalarImage(fp), label=tio.LabelMap(lfp))
for (fp, lfp) in zip(fps, lfps)
]
transform = datasets.get_test_transform(landmarks_path)
dataset = tio.SubjectsDataset(subjects, transform)
checkpoint = torch.load(checkpoint_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = models.get_unet().to(device)
if multi_gpu:
model = torch.nn.DataParallel(model)
model.module.load_state_dict(checkpoint['model'])
else:
model.load_state_dict(checkpoint['model'])
output_dir = Path(output_dir)
model.eval()
torch.set_grad_enabled(False)
loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers)
output_dir.mkdir(parents=True)
evaluator = engine.Evaluator()
df = evaluator.infer(model, loader, output_dir)
df.to_csv(df_path)
med, iqr = 100 * utils.get_median_iqr(df.Dice)
print(f'{med:.1f} ({iqr:.1f})')
return 0
def test(self, dataset, export_path=None, checkpoint=None, save=False):
"""
Model prediction for a SingleDataset.
:param dataset: training dataset.
:type dataset: :class:`DataLoader`
:param export_path: export path.
:type export_path: str
:return: MAE of the dataset.
:rtype: float
"""
checkpoint = checkpoint or os.path.join(self.expr_dir, "latest")
self.load(checkpoint)
self.eval()
prediction_path = export_path or self.expr_dir
if not os.path.exists(prediction_path):
os.makedirs(prediction_path)
start = time.time()
with torch.no_grad():
for i, data in enumerate(dataset.loader):
ct = data['ct'][tio.DATA].to(self.device)
ct = ct.transpose_(2, 4)
locations = data[tio.LOCATION]
fake_segmentation = self.netG.forward(ct)
fake_segmentation = fake_segmentation.transpose_(2, 4)
dataset.aggregator.add_batch(fake_segmentation, locations)
print(f"patch {i + 1}/{len(dataset.loader)}")
affine = dataset.transform(dataset.subject['ct']).affine
foreground = dataset.aggregator.get_output_tensor()
fake_segmentation_mask = foreground.argmax(dim=0, keepdim=True).type(torch.int8)
prediction = tio.LabelMap(tensor=fake_segmentation_mask, affine=affine)
print(f"{time.time() - start} sec. for evaluation")
if save:
prediction.save(os.path.join(prediction_path, 'fake_segmentation.nii'))
return prediction
def infer(self, model, loader, out_dir):
out_dir = Path(out_dir)
out_dir.mkdir(exist_ok=True, parents=True)
records = []
model.eval()
device = utils.get_device()
with torch.no_grad():
for batch in tqdm(loader):
inputs = batch['image'][tio.DATA].float().to(device)
segs = model(inputs).softmax(dim=1)[:, 1:].cpu() > 0.5
affines = batch['image'][tio.AFFINE]
paths = batch['image'][tio.PATH]
targets = batch['label'][tio.DATA]
for seg, target, affine, path in zip(segs, targets, affines,
paths):
image = tio.LabelMap(tensor=seg, affine=affine.numpy())
path = Path(path)
out_path = out_dir / path.name.replace(
'.nii', '_seg_cnn.nii')
image.save(out_path)
confusion = loss.get_confusion(seg[0].float(),
target[0].float())
precision = loss.get_precision(confusion)
recall = loss.get_recall(confusion)
dice = loss.get_dice_from_precision_and_recall(
precision, recall)
sid = path.name.split(
'_t1_post'
)[0] if '_t1_post' in path.name else path.name.split(
'.')[0]
record = dict(
Subject=sid,
Precision=precision.item(),
Recall=recall.item(),
Dice=dice.item(),
)
records.append(record)
df = pd.DataFrame.from_records(records)
return df
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 main(input_path, checkpoint_path, output_dir, landmarks_path, batch_size, num_workers, resample):
import torch
from tqdm import tqdm
import torchio as tio
import models
import datasets
fps = get_paths(input_path)
subjects = [tio.Subject(image=tio.ScalarImage(fp)) for fp in fps] # key must be 'image' as in get_test_transform
transform = tio.Compose((
tio.ToCanonical(),
datasets.get_test_transform(landmarks_path),
))
if resample:
transform = tio.Compose((
tio.Resample(),
transform,
# tio.CropOrPad((264, 268, 144)), # ################################# for BITE?
))
dataset = tio.SubjectsDataset(subjects, transform)
checkpoint = torch.load(checkpoint_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = models.get_unet().to(device)
model.load_state_dict(checkpoint['model'])
output_dir = Path(output_dir)
model.eval()
torch.set_grad_enabled(False)
loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
output_dir.mkdir(exist_ok=True, parents=True)
for batch in tqdm(loader):
inputs = batch['image'][tio.DATA].float().to(device)
seg = model(inputs).softmax(dim=1)[:, 1:].cpu() > 0.5
for tensor, affine, path in zip(seg, batch['image'][tio.AFFINE], batch['image'][tio.PATH]):
image = tio.LabelMap(tensor=tensor, affine=affine.numpy())
path = Path(path)
out_path = output_dir / path.name.replace('.nii', '_seg_cnn.nii')
image.save(out_path)
return 0
def apply_transforms(self, image, labels):
#inputs = np.asarray(image, dtype=np.float32)
inputs = image
inputs = torch.tensor(inputs, dtype=torch.float, requires_grad=False)
labels = torch.tensor(labels, dtype=torch.long, requires_grad=False)
""" Expected input is: (C x W x H x D) """
inputs = inputs.unsqueeze(0)
inputs = torch.moveaxis(inputs, 1, -1)
labels = labels.unsqueeze(0)
labels = torch.moveaxis(labels, 1, -1)
subject_a = tio.Subject(
one_image=tio.ScalarImage(tensor=inputs), # *** must be tensors!!!
a_segmentation=tio.LabelMap(tensor=labels))
subjects_list = [subject_a]
subjects_dataset = tio.SubjectsDataset(subjects_list,
transform=self.transforms)
subject_sample = subjects_dataset[0]
X = subject_sample['one_image']['data'].numpy()
Y = subject_sample['a_segmentation']['data'].numpy()
""" Re-arrange channels for Pytorch into (D, H, W) """
X = X[0]
X = np.moveaxis(X, -1, 0)
Y = Y[0]
Y = np.moveaxis(Y, -1, 0)
""" DEBUG """
#plot_max(X)
#plot_max(Y)
return X, Y
