def test_no_labelmap(self):
im = tio.ScalarImage(tensor=torch.rand(1, 1, 1, 1))
subject = tio.Subject(image=im, no_label=im)
sampler = tio.LabelSampler(1)
with self.assertRaises(RuntimeError):
next(sampler(subject))
def test_pil_2(self):
with self.assertRaises(RuntimeError):
tio.ScalarImage(tensor=torch.rand(2, 2, 3, 1)).as_pil()
def assert_shape(shape_in, shape_out):
np.save(path, np.random.rand(*shape_in))
image = tio.ScalarImage(path, reader=numpy_reader)
assert image.shape == shape_out
def test_with_a_list_of_images_with_different_shapes(self):
path1 = self.get_image_path('path1', shape=(5, 5, 5))
path2 = self.get_image_path('path2', shape=(7, 5, 5))
image = tio.ScalarImage(path=[path1, path2])
with self.assertRaises(RuntimeError):
image.load()
def test_data_type_uint32_array(self):
tensor = np.random.rand(1, 3, 3, 3).astype(np.uint32)
image = tio.ScalarImage(tensor=tensor)
self.assertEqual(image.data.dtype, torch.int64)
def test():
parser = argparse.ArgumentParser(
description='PyTorch Medical Segmentation Testing')
parser = parse_training_args(parser)
args, _ = parser.parse_known_args()
args = parser.parse_args()
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.enabled = args.cudnn_enabled
torch.backends.cudnn.benchmark = args.cudnn_benchmark
from data_function import MedData_test
os.makedirs(output_dir_test, exist_ok=True)
if hp.mode == '2d':
from models.two_d.unet import Unet
model = Unet(in_channels=hp.in_class, classes=hp.out_class)
# from models.two_d.miniseg import MiniSeg
# model = MiniSeg(in_input=hp.in_class, classes=hp.out_class)
# from models.two_d.fcn import FCN32s as fcn
# model = fcn(in_class =hp.in_class,n_class=hp.out_class)
# from models.two_d.segnet import SegNet
# model = SegNet(input_nbr=hp.in_class,label_nbr=hp.out_class)
# from models.two_d.deeplab import DeepLabV3
# model = DeepLabV3(in_class=hp.in_class,class_num=hp.out_class)
# from models.two_d.unetpp import ResNet34UnetPlus
# model = ResNet34UnetPlus(num_channels=hp.in_class,num_class=hp.out_class)
# from models.two_d.pspnet import PSPNet
# model = PSPNet(in_class=hp.in_class,n_classes=hp.out_class)
elif hp.mode == '3d':
from models.three_d.unet3d import UNet
model = UNet(in_channels=hp.in_class,
n_classes=hp.out_class,
base_n_filter=2)
#from models.three_d.fcn3d import FCN_Net
#model = FCN_Net(in_channels =hp.in_class,n_class =hp.out_class)
#from models.three_d.highresnet import HighRes3DNet
#model = HighRes3DNet(in_channels=hp.in_class,out_channels=hp.out_class)
#from models.three_d.densenet3d import SkipDenseNet3D
#model = SkipDenseNet3D(in_channels=hp.in_class, classes=hp.out_class)
# from models.three_d.densevoxelnet3d import DenseVoxelNet
# model = DenseVoxelNet(in_channels=hp.in_class, classes=hp.out_class)
#from models.three_d.vnet3d import VNet
#model = VNet(in_channels=hp.in_class, classes=hp.out_class)
model = torch.nn.DataParallel(model,
device_ids=devicess,
output_device=[1])
print("load model:", args.ckpt)
print(os.path.join(args.output_dir, args.latest_checkpoint_file))
ckpt = torch.load(os.path.join(args.output_dir,
args.latest_checkpoint_file),
map_location=lambda storage, loc: storage)
model.load_state_dict(ckpt["model"])
model.cuda()
test_dataset = MedData_test(source_test_dir, label_test_dir)
znorm = ZNormalization()
if hp.mode == '3d':
patch_overlap = 4, 4, 4
patch_size = hp.patch_size, hp.patch_size, hp.patch_size
elif hp.mode == '2d':
patch_overlap = 4, 4, 0
patch_size = hp.patch_size, hp.patch_size, 1
for i, subj in enumerate(test_dataset.subjects):
subj = znorm(subj)
grid_sampler = torchio.inference.GridSampler(
subj,
patch_size,
patch_overlap,
)
patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=16)
aggregator = torchio.inference.GridAggregator(grid_sampler)
aggregator_1 = torchio.inference.GridAggregator(grid_sampler)
model.eval()
with torch.no_grad():
for patches_batch in tqdm(patch_loader):
input_tensor = patches_batch['source'][torchio.DATA].to(device)
locations = patches_batch[torchio.LOCATION]
if hp.mode == '2d':
input_tensor = input_tensor.squeeze(4)
outputs = model(input_tensor)
if hp.mode == '2d':
outputs = outputs.unsqueeze(4)
logits = torch.sigmoid(outputs)
labels = logits.clone()
labels[labels > 0.5] = 1
labels[labels <= 0.5] = 0
aggregator.add_batch(logits, locations)
aggregator_1.add_batch(labels, locations)
output_tensor = aggregator.get_output_tensor()
output_tensor_1 = aggregator_1.get_output_tensor()
affine = subj['source']['affine']
if (hp.in_class == 1) and (hp.out_class == 1):
label_image = torchio.ScalarImage(tensor=output_tensor.numpy(),
affine=affine)
label_image.save(
os.path.join(output_dir_test,
str(i) + "_result_float.mhd"))
output_image = torchio.ScalarImage(tensor=output_tensor_1.numpy(),
affine=affine)
output_image.save(
os.path.join(output_dir_test,
str(i) + "_result_int.mhd"))
else:
output_tensor = output_tensor.unsqueeze(1)
output_tensor_1 = output_tensor_1.unsqueeze(1)
output_image_artery_float = torchio.ScalarImage(
tensor=output_tensor[0].numpy(), affine=affine)
output_image_artery_float.save(
os.path.join(output_dir_test,
str(i) + "_result_float_artery.mhd"))
output_image_artery_int = torchio.ScalarImage(
tensor=output_tensor_1[0].numpy(), affine=affine)
output_image_artery_int.save(
os.path.join(output_dir_test,
str(i) + "_result_int_artery.mhd"))
output_image_lung_float = torchio.ScalarImage(
tensor=output_tensor[1].numpy(), affine=affine)
output_image_lung_float.save(
os.path.join(output_dir_test,
str(i) + "_result_float_lung.mhd"))
output_image_lung_int = torchio.ScalarImage(
tensor=output_tensor_1[1].numpy(), affine=affine)
output_image_lung_int.save(
os.path.join(output_dir_test,
str(i) + "_result_int_lung.mhd"))
output_image_trachea_float = torchio.ScalarImage(
tensor=output_tensor[2].numpy(), affine=affine)
output_image_trachea_float.save(
os.path.join(output_dir_test,
str(i) + "_result_float_trachea.mhd"))
output_image_trachea_int = torchio.ScalarImage(
tensor=output_tensor_1[2].numpy(), affine=affine)
output_image_trachea_int.save(
os.path.join(output_dir_test,
str(i) + "_result_int_trachea.mhd"))
output_image_vein_float = torchio.ScalarImage(
tensor=output_tensor[3].numpy(), affine=affine)
output_image_vein_float.save(
os.path.join(output_dir_test,
str(i) + "_result_float_veiny.mhd"))
output_image_vein_int = torchio.ScalarImage(
tensor=output_tensor_1[3].numpy(), affine=affine)
output_image_vein_int.save(
os.path.join(output_dir_test,
str(i) + "_result_int_vein.mhd"))
def __init__(self, images_dir, labels_dir):
if hp.mode == '3d':
patch_size = hp.patch_size
elif hp.mode == '2d':
patch_size = (hp.patch_size, hp.patch_size, 1)
else:
raise Exception('no such kind of mode!')
queue_length = 5
samples_per_volume = 5
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)
self.queue_dataset = Queue(
self.training_set,
queue_length,
samples_per_volume,
UniformSampler(patch_size),
)
def test_wrong_scalar_image_type(self):
data = torch.ones((1, 10, 10, 10))
with self.assertRaises(ValueError):
tio.ScalarImage(tensor=data, type=tio.LABEL)
def test_no_input(self):
with self.assertRaises(ValueError):
tio.ScalarImage()
def test_wrong_path_type(self):
with self.assertRaises(TypeError):
tio.ScalarImage(5)
def test_wrong_affine(self):
with self.assertRaises(TypeError):
tio.ScalarImage(5, affine=1)
def test_wrong_path_value(self):
with self.assertRaises(RuntimeError):
tio.ScalarImage('~&./@#"!?X7=+')
def createTIOSubDS(
root_gt: Union[str, Sequence[str]],
root_input: Optional[Union[str, Sequence[str]]] = None,
filename_filter: Optional[Union[str, Sequence[str]]] = None,
split_csv: str = "",
split: str = "",
data_mode: Literal['NIFTI', 'DICOM'] = "NIFTI",
isKSpace: bool = False,
isGTNonImg: bool = False,
init_transforms: Optional[Callable] = None,
aug_transforms: Optional[Callable] = None,
transforms: Optional[Callable] = None,
) -> tio.SubjectsDataset:
if type(root_gt) is not list:
root_gt = [root_gt]
if bool(root_input) and type(root_input) is not list:
root_input = [root_input]
files = []
for gt in root_gt:
if data_mode == "NIFTI":
files += glob(gt+"/**/*.nii", recursive=True) + glob(gt+"/**/*.nii.gz", recursive=True) +\
glob(gt+"/**/*.img", recursive=True) + \
glob(gt+"/**/*.img.gz", recursive=True)
else:
# TODO: DICOM read
sys.exit("DICOM read not implemented inside createTIOSubDS")
if bool(filename_filter):
if type(filename_filter) is str:
filename_filter = [filename_filter]
files = [
f for f in files if any(filt in f for filt in filename_filter)
]
if bool(split_csv):
df = pd.read_csv(split_csv)[split]
df.dropna(inplace=True)
df = list(df)
files = [f for f in files if any(d in f for d in df)]
subjects = []
filenames = []
print("Preparing dataset .....")
for file in tqdm(files):
filenames.append(os.path.basename(file))
if bool(root_input):
gt_id = [i for i, g in enumerate(root_gt) if g in file][0]
file_in = file.replace(root_gt[gt_id], root_input[gt_id])
if not os.path.isfile(file_in):
if data_mode == "NIFTI" and ".gz" not in file_in:
file_in += ".gz"
if not os.path.isfile(file_in):
continue
else:
continue
subjects.append(
tio.Subject(inp=tio.ScalarImage(file_in),
gt=tio.LabelMap(file)
if isGTNonImg else tio.ScalarImage(file),
filename=os.path.basename(file),
processed=False))
else:
subjects.append(
tio.Subject(gt=tio.ScalarImage(file),
filename=os.path.basename(file),
processed=False))
if isKSpace:
# TODO: Image to kSpace transform
sys.exit(
"Image to kSpace transform not implemented inside createTIOSubDS")
if init_transforms is not None:
aug_transforms = init_transforms if aug_transforms is None else (
init_transforms + aug_transforms)
if aug_transforms is not None:
transforms = aug_transforms if transforms is None else (
aug_transforms + transforms)
if transforms is not None:
transforms = tio.Compose(transforms)
subjects_dataset = tio.SubjectsDataset(subjects, transform=transforms)
return subjects_dataset, filenames
transformed = dataset[0]
print('Applied transforms:') # noqa: T001
pprint.pprint(transformed.history) # noqa: T003
print('\nComposed transform to reproduce history:') # noqa: T001
print(transformed.get_composed_history()) # noqa: T001
print('\nComposed transform to invert applied transforms when possible:'
) # noqa: T001, E501
print(transformed.get_inverse_transform(ignore_intensity=False)) # noqa: T001
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
collate_fn=tio.utils.history_collate,
)
batch = tio.utils.get_first_item(loader)
print('\nTransforms applied to subjects in batch:') # noqa: T001
pprint.pprint(batch[tio.HISTORY]) # noqa: T003
for i in range(batch_size):
tensor = batch['t1'][tio.DATA][i]
affine = batch['t1'][tio.AFFINE][i]
image = tio.ScalarImage(tensor=tensor, affine=affine)
image.plot(show=False)
history = batch[tio.HISTORY][i]
title = ', '.join(t.name for t in history)
plt.suptitle(title)
plt.tight_layout()
plt.show()
def swap_image(filename):
""" Image manipulation method - swap patches of image for context restoration """
image = tio.ScalarImage(filename)
swap = tio.RandomSwap()
swapped = swap(image)
return swapped
def test_bad_key(self):
with self.assertRaises(ValueError):
tio.ScalarImage(path='', data=5)
def train():
parser = argparse.ArgumentParser(
description='PyTorch Medical Segmentation Training')
parser = parse_training_args(parser)
args, _ = parser.parse_known_args()
args = parser.parse_args()
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.enabled = args.cudnn_enabled
torch.backends.cudnn.benchmark = args.cudnn_benchmark
from data_function import MedData_train
os.makedirs(args.output_dir, exist_ok=True)
if hp.mode == '2d':
from models.two_d.unet import Unet
model = Unet(in_channels=hp.in_class, classes=hp.out_class)
# from models.two_d.miniseg import MiniSeg
# model = MiniSeg(in_input=hp.in_class, classes=hp.out_class)
# from models.two_d.fcn import FCN32s as fcn
# model = fcn(in_class =hp.in_class,n_class=hp.out_class)
# from models.two_d.segnet import SegNet
# model = SegNet(input_nbr=hp.in_class,label_nbr=hp.out_class)
# from models.two_d.deeplab import DeepLabV3
# model = DeepLabV3(in_class=hp.in_class,class_num=hp.out_class)
# from models.two_d.unetpp import ResNet34UnetPlus
# model = ResNet34UnetPlus(num_channels=hp.in_class,num_class=hp.out_class)
# from models.two_d.pspnet import PSPNet
# model = PSPNet(in_class=hp.in_class,n_classes=hp.out_class)
elif hp.mode == '3d':
from models.three_d.unet3d import UNet3D
model = UNet3D(in_channels=hp.in_class,
out_channels=hp.out_class,
init_features=32)
# from models.three_d.residual_unet3d import UNet
# model = UNet(in_channels=hp.in_class, n_classes=hp.out_class, base_n_filter=2)
#from models.three_d.fcn3d import FCN_Net
#model = FCN_Net(in_channels =hp.in_class,n_class =hp.out_class)
#from models.three_d.highresnet import HighRes3DNet
#model = HighRes3DNet(in_channels=hp.in_class,out_channels=hp.out_class)
#from models.three_d.densenet3d import SkipDenseNet3D
#model = SkipDenseNet3D(in_channels=hp.in_class, classes=hp.out_class)
# from models.three_d.densevoxelnet3d import DenseVoxelNet
# model = DenseVoxelNet(in_channels=hp.in_class, classes=hp.out_class)
#from models.three_d.vnet3d import VNet
#model = VNet(in_channels=hp.in_class, classes=hp.out_class)
model = torch.nn.DataParallel(model,
device_ids=devicess,
output_device=[1])
optimizer = torch.optim.Adam(model.parameters(), lr=args.init_lr)
# scheduler = ReduceLROnPlateau(optimizer, 'min',factor=0.5, patience=20, verbose=True)
scheduler = StepLR(optimizer, step_size=30, gamma=0.8)
# scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=5e-6)
if args.ckpt is not None:
print("load model:", args.ckpt)
print(os.path.join(args.output_dir, args.latest_checkpoint_file))
ckpt = torch.load(os.path.join(args.output_dir,
args.latest_checkpoint_file),
map_location=lambda storage, loc: storage)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optim"])
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()
# scheduler.load_state_dict(ckpt["scheduler"])
elapsed_epochs = ckpt["epoch"]
else:
elapsed_epochs = 0
model.cuda()
from loss_function import Binary_Loss, DiceLoss
criterion = Binary_Loss().cuda()
writer = SummaryWriter(args.output_dir)
train_dataset = MedData_train(source_train_dir, label_train_dir)
train_loader = DataLoader(train_dataset.queue_dataset,
batch_size=args.batch,
shuffle=True,
pin_memory=True,
drop_last=True)
model.train()
epochs = args.epochs - elapsed_epochs
iteration = elapsed_epochs * len(train_loader)
for epoch in range(1, epochs + 1):
print("epoch:" + str(epoch))
epoch += elapsed_epochs
train_epoch_avg_loss = 0.0
num_iters = 0
for i, batch in enumerate(train_loader):
if hp.debug:
if i >= 1:
break
print(f"Batch: {i}/{len(train_loader)} epoch {epoch}")
optimizer.zero_grad()
if (hp.in_class == 1) and (hp.out_class == 1):
x = batch['source']['data']
y = batch['label']['data']
x = x.type(torch.FloatTensor).cuda()
y = y.type(torch.FloatTensor).cuda()
else:
x = batch['source']['data']
y_atery = batch['atery']['data']
y_lung = batch['lung']['data']
y_trachea = batch['trachea']['data']
y_vein = batch['atery']['data']
x = x.type(torch.FloatTensor).cuda()
y = torch.cat((y_atery, y_lung, y_trachea, y_vein), 1)
y = y.type(torch.FloatTensor).cuda()
if hp.mode == '2d':
x = x.squeeze(4)
y = y.squeeze(4)
y = y / 255.
# print(y.max())
outputs = model(x)
# for metrics
logits = torch.sigmoid(outputs)
labels = logits.clone()
labels[labels > 0.5] = 1
labels[labels <= 0.5] = 0
loss = criterion(outputs, y)
num_iters += 1
loss.backward()
optimizer.step()
iteration += 1
false_positive_rate, false_negtive_rate, dice = metric(
y.cpu(), labels.cpu())
## log
writer.add_scalar('Training/Loss', loss.item(), iteration)
writer.add_scalar('Training/false_positive_rate',
false_positive_rate, iteration)
writer.add_scalar('Training/false_negtive_rate',
false_negtive_rate, iteration)
writer.add_scalar('Training/dice', dice, iteration)
print("loss:" + str(loss.item()))
print('lr:' + str(scheduler._last_lr[0]))
scheduler.step()
# Store latest checkpoint in each epoch
torch.save(
{
"model": model.state_dict(),
"optim": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"epoch": epoch,
},
os.path.join(args.output_dir, args.latest_checkpoint_file),
)
# Save checkpoint
if epoch % args.epochs_per_checkpoint == 0:
torch.save(
{
"model": model.state_dict(),
"optim": optimizer.state_dict(),
"epoch": epoch,
},
os.path.join(args.output_dir, f"checkpoint_{epoch:04d}.pt"),
)
with torch.no_grad():
if hp.mode == '2d':
x = x.unsqueeze(4)
y = y.unsqueeze(4)
outputs = outputs.unsqueeze(4)
x = x[0].cpu().detach().numpy()
y = y[0].cpu().detach().numpy()
outputs = outputs[0].cpu().detach().numpy()
affine = batch['source']['affine'][0].numpy()
if (hp.in_class == 1) and (hp.out_class == 1):
source_image = torchio.ScalarImage(tensor=x, affine=affine)
source_image.save(
os.path.join(args.output_dir,
("step-{}-source.mhd").format(epoch)))
label_image = torchio.ScalarImage(tensor=y, affine=affine)
label_image.save(
os.path.join(args.output_dir,
("step-{}-gt.mhd").format(epoch)))
output_image = torchio.ScalarImage(tensor=outputs,
affine=affine)
output_image.save(
os.path.join(args.output_dir,
("step-{}-predict.mhd").format(epoch)))
else:
y = np.expand_dims(y, axis=1)
outputs = np.expand_dims(outputs, axis=1)
source_image = torchio.ScalarImage(tensor=x, affine=affine)
source_image.save(
os.path.join(args.output_dir,
("step-{}-source.mhd").format(epoch)))
label_image_artery = torchio.ScalarImage(tensor=y[0],
affine=affine)
label_image_artery.save(
os.path.join(args.output_dir,
("step-{}-gt_artery.mhd").format(epoch)))
output_image_artery = torchio.ScalarImage(
tensor=outputs[0], affine=affine)
output_image_artery.save(
os.path.join(
args.output_dir,
("step-{}-predict_artery.mhd").format(epoch)))
label_image_lung = torchio.ScalarImage(tensor=y[1],
affine=affine)
label_image_lung.save(
os.path.join(args.output_dir,
("step-{}-gt_lung.mhd").format(epoch)))
output_image_lung = torchio.ScalarImage(tensor=outputs[1],
affine=affine)
output_image_lung.save(
os.path.join(
args.output_dir,
("step-{}-predict_lung.mhd").format(epoch)))
label_image_trachea = torchio.ScalarImage(tensor=y[2],
affine=affine)
label_image_trachea.save(
os.path.join(args.output_dir,
("step-{}-gt_trachea.mhd").format(epoch)))
output_image_trachea = torchio.ScalarImage(
tensor=outputs[2], affine=affine)
output_image_trachea.save(
os.path.join(
args.output_dir,
("step-{}-predict_trachea.mhd").format(epoch)))
label_image_vein = torchio.ScalarImage(tensor=y[3],
affine=affine)
label_image_vein.save(
os.path.join(args.output_dir,
("step-{}-gt_vein.mhd").format(epoch)))
output_image_vein = torchio.ScalarImage(tensor=outputs[3],
affine=affine)
output_image_vein.save(
os.path.join(
args.output_dir,
("step-{}-predict_vein.mhd").format(epoch)))
writer.close()
def test_repr(self):
subject = tio.Subject(t1=tio.ScalarImage(
self.get_image_path('repr_test')), )
assert 'shape' not in repr(subject['t1'])
subject.load()
assert 'shape' in repr(subject['t1'])
if not args.output_probabilities:
if args.remove_holes:
out, hole_voxels_removed = remove_holes(out, hole_size=64)
pbar.write(
f"\tFilled {hole_voxels_removed} voxels from detected holes."
)
out = torch.from_numpy(out).unsqueeze(0)
out = out.int()
image = tio.LabelMap(tensor=out)
else:
raise NotImplementedError
image = tio.ScalarImage(tensor=probs)
# inverse_transforms = subject.get_composed_history().inverse(warn=False)
# image = inverse_transforms(image)
orig_subject = dataset.subjects_map[
subject["name"]] # access subject without applying transformations
# compare shapes ignoring the channel dimension
if image.shape[1:] != orig_subject.shape[1:]:
resample_transform = tio.Resample(orig_subject.get_images()[0])
image = resample_transform(image)
assert orig_subject.shape[1:] == image.shape[
1:], "Segmentation shape and original image shape do not match"
def test_bad_affine(self):
with self.assertRaises(ValueError):
tio.ScalarImage(tensor=torch.rand(1, 2, 3, 4), affine=np.eye(3))
def main(
input_path,
checkpoint_path,
output_dir,
landmarks_path,
num_iterations,
csv_path,
batch_size,
num_workers,
gpu,
threshold,
augmentation,
save_volumes,
interpolation,
std_noise,
):
import torch
import pandas as pd
import numpy as np
import torchio as tio
from tqdm import tqdm
import models
device = torch.device(
'cuda' if torch.cuda.is_available() and gpu else 'cpu')
checkpoint = torch.load(checkpoint_path, map_location=device)
model = models.get_unet().to(device)
model.load_state_dict(checkpoint['model'])
output_dir = Path(output_dir)
model.eval()
torch.set_grad_enabled(False)
fps = get_paths(input_path)
mean_dir = output_dir / 'mean'
std_dir = output_dir / 'std'
# entropy_dir = output_dir / 'entropy'
mean_dir.mkdir(parents=True, exist_ok=True)
std_dir.mkdir(parents=True, exist_ok=True)
# entropy_dir.mkdir(parents=True, exist_ok=True)
records = []
progress = tqdm(fps, unit='subject')
for fp in progress:
subject_id = fp.name[:4]
progress.set_description(subject_id)
image = tio.ScalarImage(fp)
subject = tio.Subject(
image=image) # key must be 'image' as in get_test_transform
transform = get_transform(augmentation, landmarks_path)
dataset = tio.SubjectsDataset(num_iterations * [subject],
transform=transform)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
num_workers=num_workers,
collate_fn=lambda x: x,
)
all_results = []
for subjects_list_batch in tqdm(loader, leave=False, unit='batch'):
inputs = torch.stack([
subject.image.data for subject in subjects_list_batch
]).float().to(device)
with torch.cuda.amp.autocast():
segs = model(inputs).softmax(dim=1)[:, 1:].cpu()
iterable = list(zip(subjects_list_batch, segs))
for subject, seg in tqdm(iterable, leave=False, unit='subject'):
subject.image.set_data(seg)
inverse_transform = subject.get_inverse_transform(warn=False)
inverse_transforms = inverse_transform.transforms
first = inverse_transforms[0]
if hasattr(first, 'image_interpolation'
) and first.image_interpolation != 'linear':
first.image_interpolation = 'linear' # force interp to be lin so probs stay in [0,1]
subject_back = inverse_transform(subject)
result = subject_back.image.data
assert np.count_nonzero(
result.numpy() < 0) == 0, 'neg values found in result'
if threshold:
result = (result >= 0.5).float()
all_results.append(result)
result = torch.stack(all_results)
volumes = result.sum(dim=(-3, -2, -1)).numpy()
mean_volumes = volumes.mean()
std_volumes = volumes.std()
volume_variation_coefficient = std_volumes / mean_volumes
q1, q3 = np.percentile(volumes, (25, 75))
quartile_coefficient_of_dispersion = (q3 - q1) / (q3 + q1)
record = dict(
Subject=subject_id,
VolumeMean=mean_volumes,
VolumeSTD=std_volumes,
VVC=volume_variation_coefficient,
Q1=q1,
Q3=q3,
QCD=quartile_coefficient_of_dispersion,
)
if save_volumes:
for i, volume in enumerate(volumes):
record[f'Volume_{i}'] = volume
records.append(record)
mean = result.mean(dim=0)
std = result.std(dim=0)
# entropy = utils.get_entropy(result)
mean_image = tio.ScalarImage(tensor=mean, affine=image.affine)
std_image = tio.ScalarImage(tensor=std, affine=image.affine)
# entropy_image = tio.ScalarImage(tensor=entropy, affine=image.affine)
mean_path = mean_dir / fp.name.replace('.nii', '_mean.nii')
std_path = std_dir / fp.name.replace('.nii', '_std.nii')
# entropy_path = entropy_dir / fp.name.replace('.nii', '_entropy.nii')
mean_image.save(mean_path)
std_image.save(std_path)
# entropy_image.save(entropy_path)
# So it's updated during execution
df = pd.DataFrame.from_records(records)
df.to_csv(csv_path)
return 0
def test_nans_tensor(self):
tensor = np.random.rand(1, 2, 3, 4)
tensor[0, 0, 0, 0] = np.nan
with self.assertWarns(RuntimeWarning):
image = tio.ScalarImage(tensor=tensor, check_nans=True)
image.set_check_nans(False)
def test_with_a_list_of_images_with_different_affines(self):
path1 = self.get_image_path('path1', spacing=(1, 1, 1))
path2 = self.get_image_path('path2', spacing=(1, 2, 1))
image = tio.ScalarImage(path=[path1, path2])
with self.assertWarns(RuntimeWarning):
image.load()
def test_with_list_of_missing_files(self):
with self.assertRaises(FileNotFoundError):
tio.ScalarImage(path=['nopath', 'error'])
def test_save_image_with_data_type_boolean(self):
tensor = np.random.rand(1, 3, 3, 3).astype(bool)
image = tio.ScalarImage(tensor=tensor)
image.save(self.dir / 'image.nii')
def create_train_and_test_data_loaders(df, count_train):
images = []
regression_targets = []
sizes = {}
artifact_column_indices = [
df.columns.get_loc(c) + 1 for c in artifacts if c in df
]
for row in df.itertuples():
try:
exists = row.exists
except AttributeError:
exists = True # assume that it exists by default
if exists:
images.append(row.file_path)
row_targets = [row.overall_qa_assessment]
for i in range(len(artifacts)):
artifact_value = row[artifact_column_indices[i]]
converted_result = convert_bool_to_int(artifact_value)
row_targets.append(converted_result)
regression_targets.append(row_targets)
try:
size = row.dimensions
if size not in sizes:
sizes[size] = 1
else:
sizes[size] += 1
except AttributeError:
pass
ground_truth = np.asarray(regression_targets)
count_val = df.shape[0] - count_train
train_files = [
torchio.Subject({
'img': torchio.ScalarImage(img),
'info': info
})
for img, info in zip(images[:count_train], ground_truth[:count_train])
]
val_files = [
torchio.Subject({
'img': torchio.ScalarImage(img),
'info': info
}) for img, info in zip(images[-count_val:], ground_truth[-count_val:])
]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# calculate class weights
class_count = len(artifacts)
count0 = [0] * class_count
count1 = [0] * class_count
for s in range(count_train):
for i in range(class_count):
if regression_targets[s][i + regression_count] == 0:
count0[i] += 1
elif regression_targets[s][i + regression_count] == 1:
count1[i] += 1
# else ignore the missing data
weights_array = np.zeros(class_count)
for i in range(class_count):
weights_array[i] = count0[i] / (count0[i] + count1[i])
logger.info(f'weights_array: {weights_array}')
class_weights = torch.tensor(weights_array, dtype=torch.float).to(device)
rescale = torchio.RescaleIntensity(out_min_max=(0, 1))
ghosting = CustomGhosting(p=0.2, intensity=(0.2, 0.8))
motion = CustomMotion(p=0.15,
degrees=5.0,
translation=5.0,
num_transforms=1)
inhomogeneity = CustomBiasField(p=0.05)
spike = CustomSpike(p=0.03, num_spikes=(1, 1))
# gamma = CustomGamma(p=0.1) # after quick experimentation: gamma does not appear to help
noise = CustomNoise(p=0.05)
transforms = torchio.Compose(
[rescale, ghosting, motion, inhomogeneity, spike, noise])
# create a training data loader
train_ds = torchio.SubjectsDataset(train_files, transform=transforms)
train_loader = DataLoader(train_ds,
batch_size=1,
shuffle=True,
num_workers=4,
pin_memory=torch.cuda.is_available())
# create a validation data loader
val_ds = torchio.SubjectsDataset(val_files, transform=rescale)
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=4,
pin_memory=torch.cuda.is_available())
return train_loader, val_loader, class_weights, sizes
def test_pil_3(self):
tio.ScalarImage(tensor=torch.rand(3, 2, 3, 1)).as_pil()
def add_noise_to_image(filename):
""" Image manipulation method - add noise """
image = tio.ScalarImage(filename)
noise = tio.RandomNoise()
noised = noise(image)
return noised
def test_image_not_found(self):
with self.assertRaises(FileNotFoundError):
tio.ScalarImage('nopath')
recursive=True)
all_t1s = glob.glob(data_path + '**/*desc-restore_space-T2w_T1w.nii.gz',
recursive=True)
all_t1s = all_t1s[:max_subjects]
subjects = []
for t1_file in all_t1s:
id_subject = t1_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]
seg_file = [s for s in all_seg if id_subject in s][0]
subject = tio.Subject(
t1=tio.ScalarImage(t1_file),
t2=tio.ScalarImage(t2_file),
label=tio.LabelMap(seg_file),
)
subjects.append(subject)
dataset = tio.SubjectsDataset(subjects)
print('Dataset size:', len(dataset), 'subjects')
#%%
normalization = tio.ZNormalization(masking_method='label')
onehot = tio.OneHot()
spatial = tio.RandomAffine(scales=0.1, degrees=10, translation=0, p=0.75)
bias = tio.RandomBiasField(coefficients=0.5, p=0.3)
