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