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)