def test_shape(self, expected_shape):
test_image = nib.Nifti1Image(
np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
tempdir = tempfile.mkdtemp()
nib.save(test_image, os.path.join(tempdir, "test_image1.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_label1.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_extra1.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_image2.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_label2.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_extra2.nii.gz"))
test_data = [
{
"image": os.path.join(tempdir, "test_image1.nii.gz"),
"label": os.path.join(tempdir, "test_label1.nii.gz"),
"extra": os.path.join(tempdir, "test_extra1.nii.gz"),
},
{
"image": os.path.join(tempdir, "test_image2.nii.gz"),
"label": os.path.join(tempdir, "test_label2.nii.gz"),
"extra": os.path.join(tempdir, "test_extra2.nii.gz"),
},
]
test_transform = Compose([
LoadNiftid(keys=["image", "label", "extra"]),
SimulateDelayd(keys=["image", "label", "extra"],
delay_time=[1e-7, 1e-6, 1e-5]),
])
dataset = Dataset(data=test_data, transform=test_transform)
data1 = dataset[0]
data2 = dataset[1]
self.assertTupleEqual(data1["image"].shape, expected_shape)
self.assertTupleEqual(data1["label"].shape, expected_shape)
self.assertTupleEqual(data1["extra"].shape, expected_shape)
self.assertTupleEqual(data2["image"].shape, expected_shape)
self.assertTupleEqual(data2["label"].shape, expected_shape)
self.assertTupleEqual(data2["extra"].shape, expected_shape)
dataset = Dataset(
data=test_data,
transform=LoadNiftid(keys=["image", "label", "extra"]))
data1_simple = dataset[0]
data2_simple = dataset[1]
self.assertTupleEqual(data1_simple["image"].shape, expected_shape)
self.assertTupleEqual(data1_simple["label"].shape, expected_shape)
self.assertTupleEqual(data1_simple["extra"].shape, expected_shape)
self.assertTupleEqual(data2_simple["image"].shape, expected_shape)
self.assertTupleEqual(data2_simple["label"].shape, expected_shape)
self.assertTupleEqual(data2_simple["extra"].shape, expected_shape)
shutil.rmtree(tempdir)
def __init__(
self,
root_dir: str,
task: str,
section: str,
transform=LoadNiftid(["image", "label"]),
seed: int = 0,
val_frac: float = 0.2,
cache_num: int = sys.maxsize,
cache_rate: float = 1.0,
num_workers: int = 0,
shuffle=False,
) -> None:
if not os.path.isdir(root_dir):
raise ValueError("Root directory root_dir must be a directory.")
self.section = section
self.val_frac = val_frac
self.set_random_state(seed=seed)
dataset_dir = os.path.join(root_dir, task)
self.rann = None
if not os.path.exists(dataset_dir):
raise RuntimeError(
f"Cannot find dataset directory: {dataset_dir}, please use download=True to download it."
)
data = self._generate_data_list(dataset_dir)
data = partition_dataset(data, shuffle=shuffle)
super().__init__(data,
transform,
cache_num=cache_num,
cache_rate=cache_rate,
num_workers=num_workers)
def test_shape(self, expected_shape):
test_image = nib.Nifti1Image(
np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
tempdir = tempfile.mkdtemp()
nib.save(test_image, os.path.join(tempdir, "test_image1.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_label1.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_extra1.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_image2.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_label2.nii.gz"))
nib.save(test_image, os.path.join(tempdir, "test_extra2.nii.gz"))
test_data = [
{
"image": os.path.join(tempdir, "test_image1.nii.gz"),
"label": os.path.join(tempdir, "test_label1.nii.gz"),
"extra": os.path.join(tempdir, "test_extra1.nii.gz"),
},
{
"image": os.path.join(tempdir, "test_image2.nii.gz"),
"label": os.path.join(tempdir, "test_label2.nii.gz"),
"extra": os.path.join(tempdir, "test_extra2.nii.gz"),
},
]
dataset = CacheDataset(
data=test_data,
transform=Compose([LoadNiftid(keys=["image", "label", "extra"])]),
cache_rate=0.5)
data1 = dataset[0]
data2 = dataset[1]
shutil.rmtree(tempdir)
self.assertTupleEqual(data1["image"].shape, expected_shape)
self.assertTupleEqual(data1["label"].shape, expected_shape)
self.assertTupleEqual(data1["extra"].shape, expected_shape)
self.assertTupleEqual(data2["image"].shape, expected_shape)
self.assertTupleEqual(data2["label"].shape, expected_shape)
self.assertTupleEqual(data2["extra"].shape, expected_shape)
def __init__(
self,
root_dir: str,
task: str,
transform: Union[Sequence[Callable],
Callable] = LoadNiftid(["image", "label"]),
download: bool = False,
seed: int = 0,
nsplits: int = 5,
cache_num: int = sys.maxsize,
cache_rate: float = 1.0,
num_workers: int = 0,
) -> None:
self.root_dir = root_dir
self.task = task
self.transform = transform
self.download = download
self.seed = seed
if nsplits < 2:
raise ValueError(
"nplits must be greater than 1 for cross validation.")
self.nsplits = nsplits
self.cache_num = cache_num
self.cache_rate = cache_rate
self.num_workers = num_workers
def test_shape(self):
expected_shape = (128, 128, 128)
test_image = nib.Nifti1Image(
np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
test_data = list()
with tempfile.TemporaryDirectory() as tempdir:
for i in range(6):
nib.save(test_image,
os.path.join(tempdir, f"test_image{str(i)}.nii.gz"))
test_data.append({
"image":
os.path.join(tempdir, f"test_image{str(i)}.nii.gz")
})
test_transform = Compose([
LoadNiftid(keys="image"),
SimulateDelayd(keys="image", delay_time=1e-7),
])
test_stream = _Stream(data=test_data,
dbpath=os.path.join(tempdir, "countDB"))
dataset = IterableDataset(data=test_stream,
transform=test_transform)
for d in dataset:
self.assertTupleEqual(d["image"].shape, expected_shape)
test_stream.reset()
dataloader = DataLoader(dataset=dataset,
batch_size=3,
num_workers=2)
for d in dataloader:
self.assertTupleEqual(d["image"].shape[1:], expected_shape)
def test_load_spacingd_rotate(self, filename):
data = {"image": filename}
data_dict = LoadNiftid(keys="image")(data)
data_dict = AddChanneld(keys="image")(data_dict)
affine = data_dict["image.affine"]
data_dict["image.original_affine"] = data_dict["image.affine"] = (
np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]
]) @ affine)
t = time.time()
res_dict = Spacingd(keys="image",
pixdim=(1, 2, 3),
diagonal=True,
mode="zeros")(data_dict)
t1 = time.time()
print(f"time monai: {t1 - t}")
anat = nibabel.Nifti1Image(data_dict["image"][0],
data_dict["image.affine"])
ref = resample_to_output(anat, (1, 2, 3), order=1)
t2 = time.time()
print(f"time scipy: {t2 - t1}")
self.assertTrue(t2 >= t1)
np.testing.assert_allclose(data_dict["image.affine"],
res_dict["image.original_affine"])
np.testing.assert_allclose(res_dict["image.affine"], ref.affine)
if "anatomical" not in filename:
np.testing.assert_allclose(res_dict["image"].shape[1:], ref.shape)
np.testing.assert_allclose(ref.get_fdata(),
res_dict["image"][0],
atol=0.05)
else:
# different from the ref implementation (shape computed by round
# instead of ceil)
np.testing.assert_allclose(ref.get_fdata()[..., :-1],
res_dict["image"][0],
atol=0.05)
def test_load_spacingd_rotate(self, filename):
data = {"image": filename}
data_dict = LoadNiftid(keys="image")(data)
data_dict = AddChanneld(keys="image")(data_dict)
affine = data_dict["image.affine"]
data_dict["image.original_affine"] = data_dict["image.affine"] = (
np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]
]) @ affine)
res_dict = Spacingd(keys="image",
pixdim=(1, 2, 3),
diagonal=True,
mode="constant")(data_dict)
np.testing.assert_allclose(data_dict["image.affine"],
res_dict["image.original_affine"])
anat = nibabel.Nifti1Image(data_dict["image"][0],
data_dict["image.affine"])
ref = resample_to_output(anat, (1, 2, 3))
np.testing.assert_allclose(res_dict["image.affine"], ref.affine)
if "anatomical" not in filename:
np.testing.assert_allclose(res_dict["image"].shape[1:], ref.shape)
np.testing.assert_allclose(ref.get_fdata(), res_dict["image"][0])
else:
# different from the ref implementation (shape computed by round
# instead of ceil)
np.testing.assert_allclose(ref.get_fdata()[..., :-1],
res_dict["image"][0])
def prepare_data(self):
data_dir = self.hparams.data_dir
# Train imgs/masks
train_imgs = []
train_masks = []
with open(data_dir + 'train_imgs.txt', 'r') as f:
train_imgs = [data_dir + image.rstrip() for image in f.readlines()]
with open(data_dir + 'train_masks.txt', 'r') as f:
train_masks = [data_dir + mask.rstrip() for mask in f.readlines()]
train_dicts = [{'image': image, 'mask': mask} for (image, mask) in zip(train_imgs, train_masks)]
train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2)
# Basic transforms
data_keys = ["image", "mask"]
data_transforms = Compose(
[
LoadNiftid(keys=data_keys),
AddChanneld(keys=data_keys),
ScaleIntensityRangePercentilesd(
keys='image',
lower=25,
upper=75,
b_min=-0.5,
b_max=0.5
)
]
)
self.train_dataset = monai.data.CacheDataset(
data=train_dicts,
transform=Compose(
[
data_transforms,
RandCropByPosNegLabeld(
keys=data_keys,
label_key="mask",
spatial_size=self.hparams.patch_size,
num_samples=4,
image_key="image",
pos=0.8,
neg=0.2
),
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
self.val_dataset = monai.data.CacheDataset(
data=val_dicts,
transform=Compose(
[
data_transforms,
CenterSpatialCropd(keys=data_keys, roi_size=self.hparams.patch_size),
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
def test_shape(self, expected_shape):
test_image = nib.Nifti1Image(
np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
tempdir = tempfile.mkdtemp()
nib.save(test_image, os.path.join(tempdir, 'test_image1.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_label1.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_extra1.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_image2.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_label2.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_extra2.nii.gz'))
test_data = [{
'image': os.path.join(tempdir, 'test_image1.nii.gz'),
'label': os.path.join(tempdir, 'test_label1.nii.gz'),
'extra': os.path.join(tempdir, 'test_extra1.nii.gz')
}, {
'image': os.path.join(tempdir, 'test_image2.nii.gz'),
'label': os.path.join(tempdir, 'test_label2.nii.gz'),
'extra': os.path.join(tempdir, 'test_extra2.nii.gz')
}]
dataset = CacheDataset(
data=test_data,
transform=Compose([LoadNiftid(keys=['image', 'label', 'extra'])]),
cache_rate=0.5)
data1 = dataset[0]
data2 = dataset[1]
shutil.rmtree(tempdir)
self.assertTupleEqual(data1['image'].shape, expected_shape)
self.assertTupleEqual(data1['label'].shape, expected_shape)
self.assertTupleEqual(data1['extra'].shape, expected_shape)
self.assertTupleEqual(data2['image'].shape, expected_shape)
self.assertTupleEqual(data2['label'].shape, expected_shape)
self.assertTupleEqual(data2['extra'].shape, expected_shape)
def run_inference_test(root_dir, device=torch.device("cuda:0")):
images = sorted(glob(os.path.join(root_dir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose([
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
Spacingd(keys=["img", "seg"],
pixdim=[1.2, 0.8, 0.7],
mode=["bilinear", "nearest"],
dtype=np.float32),
ScaleIntensityd(keys=["img", "seg"]),
ToTensord(keys=["img", "seg"]),
])
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
# sliding window inferene need to input 1 image in every iteration
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
dice_metric = DiceMetric(include_background=True,
to_onehot_y=False,
sigmoid=True,
reduction="mean")
model = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
model_filename = os.path.join(root_dir, "best_metric_model.pth")
model.load_state_dict(torch.load(model_filename))
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
saver = NiftiSaver(output_dir=os.path.join(root_dir, "output"),
dtype=np.float32)
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(
device), val_data["seg"].to(device)
# define sliding window size and batch size for windows inference
sw_batch_size, roi_size = 4, (96, 96, 96)
val_outputs = sliding_window_inference(val_images, roi_size,
sw_batch_size, model)
value = dice_metric(y_pred=val_outputs, y=val_labels)
not_nans = dice_metric.not_nans.item()
metric_count += not_nans
metric_sum += value.item() * not_nans
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
saver.save_batch(val_outputs, val_data["img_meta_dict"])
metric = metric_sum / metric_count
return metric
def __init__(
self,
root_dir,
section,
transform=LoadNiftid(["image", "label"]),
cache_rate=1.0,
num_workers=0,
shuffle=False,
) -> None:
if not os.path.isdir(root_dir):
raise ValueError("Root directory root_dir must be a directory.")
self.section = section
self.shuffle = shuffle
self.val_frac = 0.2
self.set_random_state(seed=0)
dataset_dir = os.path.join(root_dir, "Task01_BrainTumour")
if not os.path.exists(dataset_dir):
raise RuntimeError(
f"Cannot find dataset directory: {dataset_dir}, please download it from Decathlon challenge."
)
data = self._generate_data_list(dataset_dir)
super(DecathlonDataset, self).__init__(data,
transform,
cache_rate=cache_rate,
num_workers=num_workers)
def get_seg_transforms(end_seg_axcodes):
seg_transforms = Compose([
LoadNiftid(keys=["image"], as_closest_canonical=False),
AddChanneld(keys=["image"]),
Orientationd(keys=["image"], axcodes=end_seg_axcodes),
])
return seg_transforms
def __init__(
self,
root_dir: str,
task: str,
section: str,
transform: Union[Sequence[Callable], Callable] = LoadNiftid(["image", "label"]),
download: bool = False,
seed: int = 0,
val_frac: float = 0.2,
cache_num: int = sys.maxsize,
cache_rate: float = 1.0,
num_workers: int = 0,
) -> None:
if not os.path.isdir(root_dir):
raise ValueError("Root directory root_dir must be a directory.")
self.section = section
self.val_frac = val_frac
self.set_random_state(seed=seed)
if task not in self.resource:
raise ValueError(f"Unsupported task: {task}, available options are: {list(self.resource.keys())}.")
dataset_dir = os.path.join(root_dir, task)
tarfile_name = f"{dataset_dir}.tar"
if download:
download_and_extract(self.resource[task], tarfile_name, root_dir, self.md5[task])
if not os.path.exists(dataset_dir):
raise RuntimeError(
f"Cannot find dataset directory: {dataset_dir}, please use download=True to download it."
)
data = self._generate_data_list(dataset_dir)
super().__init__(data, transform, cache_num=cache_num, cache_rate=cache_rate, num_workers=num_workers)
def test_shape(self, expected_shape):
test_image = nib.Nifti1Image(
np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
tempdir = tempfile.mkdtemp()
nib.save(test_image, os.path.join(tempdir, 'test_image1.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_label1.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_extra1.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_image2.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_label2.nii.gz'))
nib.save(test_image, os.path.join(tempdir, 'test_extra2.nii.gz'))
test_data = [{
'image': os.path.join(tempdir, 'test_image1.nii.gz'),
'label': os.path.join(tempdir, 'test_label1.nii.gz'),
'extra': os.path.join(tempdir, 'test_extra1.nii.gz')
}, {
'image': os.path.join(tempdir, 'test_image2.nii.gz'),
'label': os.path.join(tempdir, 'test_label2.nii.gz'),
'extra': os.path.join(tempdir, 'test_extra2.nii.gz')
}]
test_transform = Compose([
LoadNiftid(keys=['image', 'label', 'extra']),
SimulateDelayd(keys=['image', 'label', 'extra'],
delay_time=[1e-7, 1e-6, 1e-5])
])
dataset = Dataset(data=test_data, transform=test_transform)
data1 = dataset[0]
data2 = dataset[1]
self.assertTupleEqual(data1['image'].shape, expected_shape)
self.assertTupleEqual(data1['label'].shape, expected_shape)
self.assertTupleEqual(data1['extra'].shape, expected_shape)
self.assertTupleEqual(data2['image'].shape, expected_shape)
self.assertTupleEqual(data2['label'].shape, expected_shape)
self.assertTupleEqual(data2['extra'].shape, expected_shape)
dataset = Dataset(
data=test_data,
transform=LoadNiftid(keys=['image', 'label', 'extra']))
data1_simple = dataset[0]
data2_simple = dataset[1]
self.assertTupleEqual(data1_simple['image'].shape, expected_shape)
self.assertTupleEqual(data1_simple['label'].shape, expected_shape)
self.assertTupleEqual(data1_simple['extra'].shape, expected_shape)
self.assertTupleEqual(data2_simple['image'].shape, expected_shape)
self.assertTupleEqual(data2_simple['label'].shape, expected_shape)
self.assertTupleEqual(data2_simple['extra'].shape, expected_shape)
def run_inference_test(root_dir, device=torch.device("cuda:0")):
images = sorted(glob(os.path.join(root_dir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose([
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys=["img", "seg"]),
ToTensord(keys=["img", "seg"]),
])
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
# sliding window inferene need to input 1 image in every iteration
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=4,
collate_fn=list_data_collate,
pin_memory=torch.cuda.is_available())
model = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
model_filename = os.path.join(root_dir, "best_metric_model.pth")
model.load_state_dict(torch.load(model_filename))
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
saver = NiftiSaver(output_dir=os.path.join(root_dir, "output"),
dtype=int)
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(
device), val_data["seg"].to(device)
# define sliding window size and batch size for windows inference
sw_batch_size, roi_size = 4, (96, 96, 96)
val_outputs = sliding_window_inference(val_images, roi_size,
sw_batch_size, model)
value = compute_meandice(y_pred=val_outputs,
y=val_labels,
include_background=True,
to_onehot_y=False,
add_sigmoid=True)
metric_count += len(value)
metric_sum += value.sum().item()
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
saver.save_batch(
val_outputs, {
"filename_or_obj": val_data["img.filename_or_obj"],
"affine": val_data["img.affine"]
})
metric = metric_sum / metric_count
return metric
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# IXI dataset as a demo, downloadable from https://brain-development.org/ixi-dataset/
images = [
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI607-Guys-1097-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI175-HH-1570-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI385-HH-2078-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI344-Guys-0905-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI409-Guys-0960-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI584-Guys-1129-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI253-HH-1694-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI092-HH-1436-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI574-IOP-1156-T1.nii.gz"]),
os.sep.join(["workspace", "data", "medical", "ixi", "IXI-T1", "IXI585-Guys-1130-T1.nii.gz"]),
]
# 2 binary labels for gender classification: man and woman
labels = np.array([0, 0, 1, 0, 1, 0, 1, 0, 1, 0], dtype=np.int64)
val_files = [{"img": img, "label": label} for img, label in zip(images, labels)]
# Define transforms for image
val_transforms = Compose(
[
LoadNiftid(keys=["img"]),
AddChanneld(keys=["img"]),
ScaleIntensityd(keys=["img"]),
Resized(keys=["img"], spatial_size=(96, 96, 96)),
ToTensord(keys=["img"]),
]
)
# create a validation data loader
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
val_loader = DataLoader(val_ds, batch_size=2, num_workers=4, pin_memory=torch.cuda.is_available())
# Create DenseNet121
device = torch.device("cuda:0")
model = monai.networks.nets.densenet.densenet121(spatial_dims=3, in_channels=1, out_channels=2).to(device)
model.load_state_dict(torch.load("best_metric_model.pth"))
model.eval()
with torch.no_grad():
num_correct = 0.0
metric_count = 0
saver = CSVSaver(output_dir="./output")
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(device), val_data["label"].to(device)
val_outputs = model(val_images).argmax(dim=1)
value = torch.eq(val_outputs, val_labels)
metric_count += len(value)
num_correct += value.sum().item()
saver.save_batch(val_outputs, val_data["img_meta_dict"])
metric = num_correct / metric_count
print("evaluation metric:", metric)
saver.finalize()
def test_load_spacingd_rotate_non_diag(self):
data = {'image': FILES[0]}
data_dict = LoadNiftid(keys='image')(data)
data_dict = AddChanneld(keys='image')(data_dict)
res_dict = Spacingd(keys='image', pixdim=(1, 2, 3), diagonal=False, mode='nearest')(data_dict)
np.testing.assert_allclose(data_dict['image.affine'], res_dict['image.original_affine'])
np.testing.assert_allclose(
res_dict['image.affine'],
np.array([[-1., 0., 0., 32.], [0., 2., 0., -40.], [0., 0., 3., -16.], [0., 0., 0., 1.]]))
def test_values(self):
testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)),
"testing_data")
transform = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
ScaleIntensityd(keys="image"),
ToTensord(keys=["image", "label"]),
])
def _test_dataset(dataset):
self.assertEqual(len(dataset), 52)
self.assertTrue("image" in dataset[0])
self.assertTrue("label" in dataset[0])
self.assertTrue("image_meta_dict" in dataset[0])
self.assertTupleEqual(dataset[0]["image"].shape, (1, 33, 47, 34))
try: # will start downloading if testing_dir doesn't have the Decathlon files
data = DecathlonDataset(
root_dir=testing_dir,
task="Task04_Hippocampus",
transform=transform,
section="validation",
download=True,
)
except (ContentTooShortError, HTTPError, RuntimeError) as e:
print(str(e))
if isinstance(e, RuntimeError):
# FIXME: skip MD5 check as current downloading method may fail
self.assertTrue(str(e).startswith("MD5 check"))
return # skipping this test due the network connection errors
_test_dataset(data)
data = DecathlonDataset(root_dir=testing_dir,
task="Task04_Hippocampus",
transform=transform,
section="validation",
download=False)
_test_dataset(data)
data = DecathlonDataset(root_dir=testing_dir,
task="Task04_Hippocampus",
section="validation",
download=False)
self.assertTupleEqual(data[0]["image"].shape, (33, 47, 34))
shutil.rmtree(os.path.join(testing_dir, "Task04_Hippocampus"))
try:
data = DecathlonDataset(
root_dir=testing_dir,
task="Task04_Hippocampus",
transform=transform,
section="validation",
download=False,
)
except RuntimeError as e:
print(str(e))
self.assertTrue(str(e).startswith("Cannot find dataset directory"))
def get_transforms(self):
self.logger.info("Getting transforms...")
# Setup transforms of data sets
train_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
NormalizeIntensityd(keys=["image"]),
SpatialPadd(keys=["image", "label"],
spatial_size=self.pad_crop_shape),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandSpatialCropd(keys=["image", "label"],
roi_size=self.pad_crop_shape,
random_center=True,
random_size=False),
ToTensord(keys=["image", "label"]),
])
val_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
NormalizeIntensityd(keys=["image"]),
SpatialPadd(keys=["image", "label"],
spatial_size=self.pad_crop_shape),
RandSpatialCropd(
keys=["image", "label"],
roi_size=self.pad_crop_shape,
random_center=True,
random_size=False,
),
ToTensord(keys=["image", "label"]),
])
test_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
NormalizeIntensityd(keys=["image"]),
ToTensord(keys=["image", "label"]),
])
return train_transforms, val_transforms, test_transforms
def __init__(
self,
root_dir: str,
task: str,
section: str,
transform: Union[Sequence[Callable], Callable] = (),
download: bool = False,
seed: int = 0,
val_frac: float = 0.2,
cache_num: int = sys.maxsize,
cache_rate: float = 1.0,
num_workers: int = 0,
) -> None:
if not os.path.isdir(root_dir):
raise ValueError("Root directory root_dir must be a directory.")
self.section = section
self.val_frac = val_frac
self.set_random_state(seed=seed)
if task not in self.resource:
raise ValueError(
f"Unsupported task: {task}, available options are: {list(self.resource.keys())}."
)
dataset_dir = os.path.join(root_dir, task)
tarfile_name = f"{dataset_dir}.tar"
if download:
download_and_extract(self.resource[task], tarfile_name, root_dir,
self.md5[task])
if not os.path.exists(dataset_dir):
raise RuntimeError(
f"Cannot find dataset directory: {dataset_dir}, please use download=True to download it."
)
self.indices: np.ndarray = np.array([])
data = self._generate_data_list(dataset_dir)
# as `release` key has typo in Task04 config file, ignore it.
property_keys = [
"name",
"description",
"reference",
"licence",
"tensorImageSize",
"modality",
"labels",
"numTraining",
"numTest",
]
self._properties = load_decathlon_properties(
os.path.join(dataset_dir, "dataset.json"), property_keys)
if transform == ():
transform = LoadNiftid(["image", "label"])
super().__init__(data,
transform,
cache_num=cache_num,
cache_rate=cache_rate,
num_workers=num_workers)
def test_values(self):
tempdir = tempfile.mkdtemp()
transform = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
ScaleIntensityd(keys="image"),
ToTensord(keys=["image", "label"]),
])
def _test_dataset(dataset):
self.assertEqual(len(dataset), 52)
self.assertTrue("image" in dataset[0])
self.assertTrue("label" in dataset[0])
self.assertTrue("image_meta_dict" in dataset[0])
self.assertTupleEqual(dataset[0]["image"].shape, (1, 33, 47, 34))
try:
data = DecathlonDataset(root_dir=tempdir,
task="Task04_Hippocampus",
transform=transform,
section="validation",
download=True)
except RuntimeError as e:
if str(e).startswith(
"download failed due to network issue or permission denied."
):
shutil.rmtree(tempdir)
return
_test_dataset(data)
data = DecathlonDataset(root_dir=tempdir,
task="Task04_Hippocampus",
transform=transform,
section="validation",
download=False)
_test_dataset(data)
data = DecathlonDataset(root_dir=tempdir,
task="Task04_Hippocampus",
section="validation",
download=False)
self.assertTupleEqual(data[0]["image"].shape, (33, 47, 34))
shutil.rmtree(os.path.join(tempdir, "Task04_Hippocampus"))
try:
data = DecathlonDataset(root_dir=tempdir,
task="Task04_Hippocampus",
transform=transform,
section="validation",
download=False)
except RuntimeError as e:
print(str(e))
self.assertTrue(
str(e).startswith("can not find dataset directory"))
shutil.rmtree(tempdir)
def test_load_spacingd(self, filename):
data = {'image': filename}
data_dict = LoadNiftid(keys='image')(data)
data_dict = AddChanneld(keys='image')(data_dict)
res_dict = Spacingd(keys='image', pixdim=(1, 2, 3), diagonal=True, mode='constant')(data_dict)
np.testing.assert_allclose(data_dict['image.affine'], res_dict['image.original_affine'])
anat = nibabel.Nifti1Image(data_dict['image'][0], data_dict['image.affine'])
ref = resample_to_output(anat, (1, 2, 3))
np.testing.assert_allclose(res_dict['image.affine'], ref.affine)
np.testing.assert_allclose(res_dict['image'].shape[1:], ref.shape)
np.testing.assert_allclose(ref.get_fdata(), res_dict['image'][0])
def test_shape(self, input_param, expected_shape):
test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
tempdir = tempfile.mkdtemp()
test_data = dict()
with tempfile.TemporaryDirectory() as tempdir:
for key in KEYS:
nib.save(test_image, os.path.join(tempdir, key + '.nii.gz'))
test_data.update({key: os.path.join(tempdir, key + '.nii.gz')})
result = LoadNiftid(**input_param)(test_data)
for key in KEYS:
self.assertTupleEqual(result[key].shape, expected_shape)
def get_test_transforms(end_image_shape):
test_transforms = Compose([
LoadNiftid(keys=["image"]),
AddChanneld(keys=["image"]),
Orientationd(keys=["image"], axcodes="RAS"),
Winsorized(keys=["image"]),
NormalizeIntensityd(keys=["image"]),
ScaleIntensityd(keys=["image"]),
SpatialPadd(keys=["image"], spatial_size=end_image_shape),
ToTensord(keys=["image"]),
])
return test_transforms
def test_load_spacingd_non_diag(self):
data = {'image': FILES[1]}
data_dict = LoadNiftid(keys='image')(data)
data_dict = AddChanneld(keys='image')(data_dict)
affine = data_dict['image.affine']
data_dict['image.original_affine'] = data_dict['image.affine'] = \
np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]]) @ affine
res_dict = Spacingd(keys='image', pixdim=(1, 2, 3), diagonal=False, mode='constant')(data_dict)
np.testing.assert_allclose(data_dict['image.affine'], res_dict['image.original_affine'])
np.testing.assert_allclose(
res_dict['image.affine'],
np.array([[0., 0., 3., -27.599409], [0., 2., 0., -47.977585], [-1., 0., 0., 35.297897], [0., 0., 0., 1.]]))
def test_load_spacingd_rotate_non_diag(self):
data = {"image": FILES[0]}
data_dict = LoadNiftid(keys="image")(data)
data_dict = AddChanneld(keys="image")(data_dict)
res_dict = Spacingd(keys="image",
pixdim=(1, 2, 3),
diagonal=False,
mode="border")(data_dict)
np.testing.assert_allclose(
res_dict["image_meta_dict"]["affine"],
np.array([[-1.0, 0.0, 0.0, 32.0], [0.0, 2.0, 0.0, -40.0],
[0.0, 0.0, 3.0, -16.0], [0.0, 0.0, 0.0, 1.0]]),
)
def setup(self, stage):
data_dir = 'data/'
# Train imgs/masks
train_imgs = []
with open(data_dir + 'train_imgs.txt', 'r') as f:
train_imgs = [image.rstrip() for image in f.readlines()]
train_masks = []
with open(data_dir + 'train_masks.txt', 'r') as f:
train_masks = [mask.rstrip() for mask in f.readlines()]
train_dicts = [{'image': image, 'mask': mask} for (image, mask) in zip(train_imgs, train_masks)]
train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2)
# Basic transforms
data_keys = ["image", "mask"]
data_transforms = Compose(
[
LoadNiftid(keys=data_keys),
AddChanneld(keys=data_keys),
NormalizeIntensityd(keys="image"),
RandCropByPosNegLabeld(
keys=data_keys, label_key="mask", size=(256, 256, 16), num_samples=4, image_key="image"
),
]
)
self.train_dataset = monai.data.CacheDataset(
data=train_dicts,
transform=Compose(
[
data_transforms,
self.augmentations,
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
self.val_dataset = monai.data.CacheDataset(
data=val_dicts,
transform=Compose(
[
data_transforms,
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
def test_load_spacingd_non_diag_ornt(self):
data = {'image': FILES[1]}
data_dict = LoadNiftid(keys='image')(data)
data_dict = AddChanneld(keys='image')(data_dict)
affine = data_dict['image.affine']
data_dict['image.original_affine'] = data_dict['image.affine'] = \
np.array([[0, 0, 1, 0], [0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1]]) @ affine
res_dict = Spacingd(keys='image', pixdim=(1, 2, 3), diagonal=False, mode='constant')(data_dict)
res_dict = Orientationd(keys='image', axcodes='LPI')(res_dict)
np.testing.assert_allclose(data_dict['image.affine'], res_dict['image.original_affine'])
np.testing.assert_allclose(
res_dict['image.affine'],
np.array([[-3., 0., 0., 56.4005909], [0., -2., 0., 52.02241516], [0., 0., -1., 35.29789734],
[0., 0., 0., 1.]]))
def test_values(self):
testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)),
"testing_data")
transform = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
ScaleIntensityd(keys="image"),
ToTensord(keys=["image", "label"]),
])
def _test_dataset(dataset):
self.assertEqual(len(dataset), 52)
self.assertTrue("image" in dataset[0])
self.assertTrue("label" in dataset[0])
self.assertTrue("image_meta_dict" in dataset[0])
self.assertTupleEqual(dataset[0]["image"].shape, (1, 34, 49, 41))
cvdataset = CrossValidation(
dataset_cls=DecathlonDataset,
nfolds=5,
seed=12345,
root_dir=testing_dir,
task="Task04_Hippocampus",
section="validation",
transform=transform,
download=True,
)
try: # will start downloading if testing_dir doesn't have the Decathlon files
data = cvdataset.get_dataset(folds=0)
except (ContentTooShortError, HTTPError, RuntimeError) as e:
print(str(e))
if isinstance(e, RuntimeError):
# FIXME: skip MD5 check as current downloading method may fail
self.assertTrue(str(e).startswith("md5 check"))
return # skipping this test due the network connection errors
_test_dataset(data)
# test training data for fold 0 of 5 splits
data = cvdataset.get_dataset(folds=[1, 2, 3, 4])
self.assertTupleEqual(data[0]["image"].shape, (1, 35, 52, 33))
self.assertEqual(len(data), 208)
# test train / validation for fold 4 of 5 splits
data = cvdataset.get_dataset(folds=[4])
self.assertTupleEqual(data[0]["image"].shape, (1, 38, 53, 30))
self.assertEqual(len(data), 52)
data = cvdataset.get_dataset(folds=[0, 1, 2, 3])
self.assertTupleEqual(data[0]["image"].shape, (1, 34, 49, 41))
self.assertEqual(len(data), 208)
def get_xforms(args, mode="train", keys=("image", "label")):
"""returns a composed transform for train/val/infer."""
xforms = [
LoadNiftid(keys),
AddChanneld(keys),
Orientationd(keys, axcodes="LPS"),
Spacingd(keys,
pixdim=(1.25, 1.25, 5.0),
mode=("bilinear", "nearest")[:len(keys)]),
ScaleIntensityRanged(keys[0],
a_min=-1000.0,
a_max=500.0,
b_min=0.0,
b_max=1.0,
clip=True),
]
if mode == "train":
xforms.extend([
SpatialPadd(keys,
spatial_size=(args.patch_size, args.patch_size, -1),
mode="reflect"), # ensure at least 192x192
RandAffined(
keys,
prob=0.15,
rotate_range=(-0.05, 0.05),
scale_range=(-0.1, 0.1),
mode=("bilinear", "nearest"),
as_tensor_output=False,
),
RandCropByPosNegLabeld(keys,
label_key=keys[1],
spatial_size=(args.patch_size,
args.patch_size,
args.n_slice),
num_samples=3),
RandGaussianNoised(keys[0], prob=0.15, std=0.01),
RandFlipd(keys, spatial_axis=0, prob=0.5),
RandFlipd(keys, spatial_axis=1, prob=0.5),
RandFlipd(keys, spatial_axis=2, prob=0.5),
])
dtype = (np.float32, np.uint8)
if mode == "val":
dtype = (np.float32, np.uint8)
if mode == "infer":
dtype = (np.float32, )
xforms.extend([CastToTyped(keys, dtype=dtype), ToTensord(keys)])
return monai.transforms.Compose(xforms)
