def test_pad_collation(self, t_type, collate_method, transform):
if t_type == dict:
dataset = CacheDataset(self.dict_data, transform, progress=False)
else:
dataset = _Dataset(self.list_data, self.list_labels, transform)
# Default collation should raise an error
loader_fail = DataLoader(dataset, batch_size=10)
with self.assertRaises(RuntimeError):
for _ in loader_fail:
pass
# Padded collation shouldn't
loader = DataLoader(dataset, batch_size=10, collate_fn=collate_method)
# check collation in forward direction
for data in loader:
if t_type == dict:
shapes = []
decollated_data = decollate_batch(data)
for d in decollated_data:
output = PadListDataCollate.inverse(d)
shapes.append(output["image"].shape)
self.assertTrue(
len(set(shapes)) > 1
) # inverted shapes must be different because of random xforms
def test_loading_dict(self):
set_determinism(seed=1234)
# test sequence input data with dict
data = [
{
"image": np.arange(16, dtype=float).reshape(1, 4, 4),
"label": np.arange(16, dtype=float).reshape(1, 4, 4),
"metadata": "test string",
},
{
"image": np.arange(16, dtype=float).reshape(1, 4, 4),
"label": np.arange(16, dtype=float).reshape(1, 4, 4),
"metadata": "test string",
},
]
# image level
patch_intensity = RandShiftIntensityd(keys="image",
offsets=1.0,
prob=1.0)
patch_iter = PatchIterd(keys=["image", "label"],
patch_size=(2, 2),
start_pos=(0, 0))
ds = GridPatchDataset(data=data,
patch_iter=patch_iter,
transform=patch_intensity,
with_coordinates=True)
# use the grid patch dataset
for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=0):
np.testing.assert_equal(item[0]["image"].shape, (2, 1, 2, 2))
np.testing.assert_equal(item[0]["label"].shape, (2, 1, 2, 2))
self.assertListEqual(item[0]["metadata"],
["test string", "test string"])
np.testing.assert_allclose(
item[0]["image"],
np.array([[[[1.4965, 2.4965], [5.4965, 6.4965]]],
[[[11.3584, 12.3584], [15.3584, 16.3584]]]]),
rtol=1e-4,
)
np.testing.assert_allclose(item[1],
np.array([[[0, 1], [0, 2], [2, 4]],
[[0, 1], [2, 4], [2, 4]]]),
rtol=1e-5)
if sys.platform != "win32":
for item in DataLoader(ds,
batch_size=2,
shuffle=False,
num_workers=2):
np.testing.assert_equal(item[0]["image"].shape, (2, 1, 2, 2))
np.testing.assert_allclose(
item[0]["image"],
np.array([[[[1.2548, 2.2548], [5.2548, 6.2548]]],
[[[9.1106, 10.1106], [13.1106, 14.1106]]]]),
rtol=1e-3,
)
np.testing.assert_allclose(item[1],
np.array([[[0, 1], [0, 2], [2, 4]],
[[0, 1], [2, 4], [2, 4]]]),
rtol=1e-5)
def test_loading_array(self):
set_determinism(seed=1234)
# image dataset
images = [
np.arange(16, dtype=float).reshape(1, 4, 4),
np.arange(16, dtype=float).reshape(1, 4, 4)
]
# image patch sampler
n_samples = 8
sampler = RandSpatialCropSamples(roi_size=(3, 3),
num_samples=n_samples,
random_center=True,
random_size=False)
# image level
patch_intensity = RandShiftIntensity(offsets=1.0, prob=1.0)
image_ds = Dataset(images, transform=patch_intensity)
# patch level
ds = PatchDataset(dataset=image_ds,
patch_func=sampler,
samples_per_image=n_samples,
transform=patch_intensity)
np.testing.assert_equal(len(ds), n_samples * len(images))
# use the patch dataset, length: len(images) x samplers_per_image
for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=0):
np.testing.assert_equal(tuple(item.shape), (2, 1, 3, 3))
np.testing.assert_allclose(
item[0],
np.array([[[1.338681, 2.338681, 3.338681],
[5.338681, 6.338681, 7.338681],
[9.338681, 10.338681, 11.338681]]]),
rtol=1e-5,
)
if sys.platform != "win32":
for item in DataLoader(ds,
batch_size=2,
shuffle=False,
num_workers=2):
np.testing.assert_equal(tuple(item.shape), (2, 1, 3, 3))
np.testing.assert_allclose(
item[0],
np.array([[
[4.957847, 5.957847, 6.957847],
[8.957847, 9.957847, 10.957847],
[12.957847, 13.957847, 14.957847],
]]),
rtol=1e-5,
)
set_determinism(seed=None)
def test_dataloader(self):
dataset = Dataset(
data=[{
"img": np.array([[[0.0, 1.0], [2.0, 3.0]]])
}, {
"img": np.array([[[0.0, 1.0], [2.0, 3.0]]])
}],
transform=IntensityStatsd(keys="img",
ops=["max", "mean"],
key_prefix="orig"),
)
# set num workers = 0 for mac / win
num_workers = 2 if sys.platform == "linux" else 0
dataloader = DataLoader(dataset=dataset,
num_workers=num_workers,
batch_size=2)
orig_method = mp.get_start_method()
mp.set_start_method("spawn", force=True)
for d in dataloader:
meta = d["img_meta_dict"]
np.testing.assert_allclose(meta["orig_max"], [3.0, 3.0], atol=1e-3)
np.testing.assert_allclose(meta["orig_mean"], [1.5, 1.5],
atol=1e-3)
# restore the mp method
mp.set_start_method(orig_method, force=True)
def test_collation(self, _, transform, collate_fn, ndim):
data = self.data_3d if ndim == 3 else self.data_2d
if collate_fn:
modified_transform = transform
else:
modified_transform = Compose(
[transform,
ResizeWithPadOrCropd(KEYS, 100),
ToTensord(KEYS)])
# num workers = 0 for mac or gpu transforms
num_workers = 0 if sys.platform != "linux" or torch.cuda.is_available(
) else 2
dataset = CacheDataset(data,
transform=modified_transform,
progress=False)
loader = DataLoader(dataset,
num_workers,
batch_size=self.batch_size,
collate_fn=collate_fn)
for item in loader:
np.testing.assert_array_equal(
item["image_transforms"][0]["do_transforms"],
item["label_transforms"][0]["do_transforms"])
def test_set_data(self):
data_list1 = list(range(10))
transform = Compose([
Lambda(func=lambda x: np.array([x * 10])),
RandLambda(func=lambda x: x + 1)
])
dataset = CacheDataset(
data=data_list1,
transform=transform,
cache_rate=1.0,
num_workers=4,
progress=True,
copy_cache=not sys.platform == "linux",
)
num_workers = 2 if sys.platform == "linux" else 0
dataloader = DataLoader(dataset=dataset,
num_workers=num_workers,
batch_size=1)
for i, d in enumerate(dataloader):
np.testing.assert_allclose([[data_list1[i] * 10 + 1]], d)
# simulate another epoch, the cache content should not be modified
for i, d in enumerate(dataloader):
np.testing.assert_allclose([[data_list1[i] * 10 + 1]], d)
# update the datalist and fill the cache content
data_list2 = list(range(-10, 0))
dataset.set_data(data=data_list2)
# rerun with updated cache content
for i, d in enumerate(dataloader):
np.testing.assert_allclose([[data_list2[i] * 10 + 1]], d)
def test_exception(self, datalist):
dataset = Dataset(data=datalist, transform=None)
dataloader = DataLoader(dataset=dataset, batch_size=2, num_workers=0)
with self.assertRaisesRegex((TypeError, RuntimeError),
"Collate error on the key"):
for _ in dataloader:
pass
def show_element(dataset, number_of_examples=1):
check_loader = DataLoader(dataset)
for i, batch_data in enumerate(check_loader):
image, label = (batch_data["image"][0][0], batch_data["label"][0][0])
show3(np.append(image, label * torch.max(image), axis=1))
if i + 1 == number_of_examples:
break
def test_decollation(self, *transforms):
batch_size = 2
num_workers = 2
t_compose = Compose(
[AddChanneld(KEYS),
Compose(transforms),
ToTensord(KEYS)])
# If nibabel present, read from disk
if has_nib:
t_compose = Compose([LoadImaged("image"), t_compose])
dataset = CacheDataset(self.data, t_compose, progress=False)
loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
for b, batch_data in enumerate(loader):
decollated_1 = decollate_batch(batch_data)
decollated_2 = Decollated()(batch_data)
for decollated in [decollated_1, decollated_2]:
for i, d in enumerate(decollated):
self.check_match(dataset[b * batch_size + i], d)
def test_values(self):
datalist = [
{
"image": "spleen_19.nii.gz",
"label": "spleen_label_19.nii.gz"
},
{
"image": "spleen_31.nii.gz",
"label": "spleen_label_31.nii.gz"
},
]
transform = Compose([
DataStatsd(keys=["image", "label"],
data_shape=False,
value_range=False,
data_value=True),
SimulateDelayd(keys=["image", "label"], delay_time=0.1),
])
dataset = CacheDataset(data=datalist,
transform=transform,
cache_rate=0.5,
cache_num=1)
dataloader = DataLoader(dataset=dataset, batch_size=2, num_workers=2)
for d in dataloader:
self.assertEqual(d["image"][0], "spleen_19.nii.gz")
self.assertEqual(d["image"][1], "spleen_31.nii.gz")
self.assertEqual(d["label"][0], "spleen_label_19.nii.gz")
self.assertEqual(d["label"][1], "spleen_label_31.nii.gz")
def test_decollation(self, batch_size=2, num_workers=2):
im = create_test_image_2d(100, 101)[0]
data = [{
"image": make_nifti_image(im) if has_nib else im
} for _ in range(6)]
transforms = Compose([
AddChanneld("image"),
SpatialPadd("image", 150),
RandFlipd("image", prob=1.0, spatial_axis=1),
ToTensord("image"),
])
# If nibabel present, read from disk
if has_nib:
transforms = Compose([LoadImaged("image"), transforms])
dataset = CacheDataset(data, transforms, progress=False)
loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
for b, batch_data in enumerate(loader):
decollated_1 = decollate_batch(batch_data)
decollated_2 = Decollated()(batch_data)
for decollated in [decollated_1, decollated_2]:
for i, d in enumerate(decollated):
self.check_match(dataset[b * batch_size + i], d)
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 = []
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{i}.nii.gz")})
test_transform = Compose([
LoadImaged(keys="image"),
SimulateDelayd(keys="image", delay_time=1e-7)
])
data_iterator = _Stream(test_data)
with self.assertRaises(TypeError): # Dataset doesn't work
dataset = Dataset(data=data_iterator, transform=test_transform)
for _ in dataset:
pass
dataset = IterableDataset(data=data_iterator,
transform=test_transform)
for d in dataset:
self.assertTupleEqual(d["image"].shape, expected_shape)
num_workers = 2 if sys.platform == "linux" else 0
dataloader = DataLoader(dataset=dataset,
batch_size=3,
num_workers=num_workers)
for d in dataloader:
self.assertTupleEqual(d["image"].shape[1:], expected_shape)
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([
LoadImaged(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 main():
root_path = Path(
"../CoronaryVesselGeneration/AngioGenAppNew/output/testing/")
keys = range(100)
values = [str(root_path / f"{(i+1):04}" / "mesh.stl") for i in keys]
meshes = dict(zip(keys, values))
input()
dataset = MeshDataset(meshes)
input()
dataloader = DataLoader(dataset, batch_size=100, num_workers=0)
# Evaluate how long this all takes
import time
from PIL import Image
for batch in dataloader:
print(batch.shape)
for i in range(3):
im = Image.fromarray(batch.numpy()[i, :, :, 0].astype(np.uint8),
'L')
file_name = f'{i:03d}.png'
im.save(file_name)
t1 = time.time()
for batch in dataloader:
t2 = time.time()
input(t2 - t1) # Allows memory to be checked externally
t1 = time.time()
def test_time(self):
dataset = Dataset(
data=self.datalist * 2,
transform=self.transform) # contains data for 2 batches
dataloader = DataLoader(dataset=dataset, batch_size=2, num_workers=0)
tbuffer = ThreadBuffer(dataloader)
with PerfContext() as pc:
for _ in dataloader:
time.sleep(
0.5
) # each batch takes 0.8 s to generate on top of this time
unbuffered_time = pc.total_time
with PerfContext() as pc:
for _ in tbuffer:
time.sleep(
0.5
) # while "computation" is happening the next batch is being generated, saving 0.4 s
buffered_time = pc.total_time
if sys.platform == "darwin": # skip macOS measure
print(
f"darwin: Buffered time {buffered_time} vs unbuffered time {unbuffered_time}"
)
else:
self.assertTrue(
buffered_time < unbuffered_time,
f"Buffered time {buffered_time} should be less than unbuffered time {unbuffered_time}",
)
def test_container(self):
net = torch.nn.Conv2d(1, 1, 3, padding=1)
opt = torch.optim.Adam(net.parameters())
img = torch.rand(1, 16, 16)
data = {CommonKeys.IMAGE: img, CommonKeys.LABEL: img}
loader = DataLoader([data for _ in range(10)])
trainer = SupervisedTrainer(
device=torch.device("cpu"),
max_epochs=1,
train_data_loader=loader,
network=net,
optimizer=opt,
loss_function=torch.nn.L1Loss(),
)
con = ThreadContainer(trainer)
con.start()
time.sleep(1) # wait for trainer to start
self.assertTrue(con.is_alive)
self.assertIsNotNone(con.status())
self.assertTrue(len(con.status_dict) > 0)
con.join()
def test_epistemic_scoring(self):
input_size = (20, 20, 20)
device = "cuda" if torch.cuda.is_available() else "cpu"
keys = ["image", "label"]
num_training_ims = 10
train_data = self.get_data(num_training_ims, input_size)
test_data = self.get_data(1, input_size)
transforms = Compose([
AddChanneld(keys),
CropForegroundd(keys, source_key="image"),
DivisiblePadd(keys, 4),
])
infer_transforms = Compose([
AddChannel(),
CropForeground(),
DivisiblePad(4),
])
train_ds = CacheDataset(train_data, transforms)
# output might be different size, so pad so that they match
train_loader = DataLoader(train_ds,
batch_size=2,
collate_fn=pad_list_data_collate)
model = UNet(3, 1, 1, channels=(6, 6), strides=(2, 2)).to(device)
loss_function = DiceLoss(sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-3)
num_epochs = 10
for _ in trange(num_epochs):
epoch_loss = 0
for batch_data in train_loader:
inputs, labels = batch_data["image"].to(
device), batch_data["label"].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_loss /= len(train_loader)
entropy_score = EpistemicScoring(model=model,
transforms=infer_transforms,
roi_size=[20, 20, 20],
num_samples=10)
# Call Individual Infer from Epistemic Scoring
ip_stack = [test_data["image"], test_data["image"], test_data["image"]]
ip_stack = np.array(ip_stack)
score_3d = entropy_score.entropy_3d_volume(ip_stack)
score_3d_sum = np.sum(score_3d)
# Call Entropy Metric from Epistemic Scoring
self.assertEqual(score_3d.shape, input_size)
self.assertIsInstance(score_3d_sum, np.float32)
self.assertGreater(score_3d_sum, 3.0)
def test_randomize(self):
dataset = _RandomDataset()
dataloader = DataLoader(dataset, batch_size=2, num_workers=3)
output = []
for _ in range(2):
for batch in dataloader:
output.extend(batch.data.numpy().flatten().tolist())
self.assertListEqual(output, [594, 170, 524, 778, 370, 906, 292, 589, 762, 763, 156, 886, 42, 405, 221, 166])
def test_loading_array(self):
set_determinism(seed=1234)
# image dataset
images = [
np.arange(16, dtype=float).reshape(1, 4, 4),
np.arange(16, dtype=float).reshape(1, 4, 4)
]
# image level
patch_intensity = RandShiftIntensity(offsets=1.0, prob=1.0)
patch_iter = PatchIter(patch_size=(2, 2), start_pos=(0, 0))
ds = GridPatchDataset(dataset=images,
patch_iter=patch_iter,
transform=patch_intensity)
# use the grid patch dataset
for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=0):
np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2))
np.testing.assert_allclose(
item[0],
np.array([[[[1.7413, 2.7413], [5.7413, 6.7413]]],
[[[9.1419, 10.1419], [13.1419, 14.1419]]]]),
rtol=1e-5,
)
np.testing.assert_allclose(
item[1],
np.array([[[0, 1], [0, 2], [2, 4]], [[0, 1], [2, 4], [2, 4]]]),
rtol=1e-5,
)
if sys.platform != "win32":
for item in DataLoader(ds,
batch_size=2,
shuffle=False,
num_workers=2):
np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2))
np.testing.assert_allclose(
item[0],
np.array([[[[2.3944, 3.3944], [6.3944, 7.3944]]],
[[[10.6551, 11.6551], [14.6551, 15.6551]]]]),
rtol=1e-3,
)
np.testing.assert_allclose(
item[1],
np.array([[[0, 1], [0, 2], [2, 4]], [[0, 1], [2, 4], [2, 4]]]),
rtol=1e-5,
)
def test_inverse_inferred_seg(self):
test_data = []
for _ in range(20):
image, label = create_test_image_2d(100, 101)
test_data.append({
"image": image,
"label": label.astype(np.float32)
})
batch_size = 10
# num workers = 0 for mac
num_workers = 2 if sys.platform != "darwin" else 0
transforms = Compose([
AddChanneld(KEYS),
SpatialPadd(KEYS, (150, 153)),
CenterSpatialCropd(KEYS, (110, 99))
])
num_invertible_transforms = sum(1 for i in transforms.transforms
if isinstance(i, InvertibleTransform))
dataset = CacheDataset(test_data, transform=transforms, progress=False)
loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = UNet(
dimensions=2,
in_channels=1,
out_channels=1,
channels=(2, 4),
strides=(2, ),
).to(device)
data = first(loader)
labels = data["label"].to(device)
segs = model(labels).detach().cpu()
label_transform_key = "label" + InverseKeys.KEY_SUFFIX.value
segs_dict = {
"label": segs,
label_transform_key: data[label_transform_key]
}
segs_dict_decollated = decollate_batch(segs_dict)
# inverse of individual segmentation
seg_dict = first(segs_dict_decollated)
with allow_missing_keys_mode(transforms):
inv_seg = transforms.inverse(seg_dict)["label"]
self.assertEqual(len(data["label_transforms"]),
num_invertible_transforms)
self.assertEqual(len(seg_dict["label_transforms"]),
num_invertible_transforms)
self.assertEqual(inv_seg.shape[1:], test_data[0]["label"].shape)
def test_inverse_inferred_seg(self, extra_transform):
test_data = []
for _ in range(20):
image, label = create_test_image_2d(100, 101)
test_data.append({
"image": image,
"label": label.astype(np.float32)
})
batch_size = 10
# num workers = 0 for mac
num_workers = 2 if sys.platform == "linux" else 0
transforms = Compose([
AddChanneld(KEYS),
SpatialPadd(KEYS, (150, 153)), extra_transform
])
dataset = CacheDataset(test_data, transform=transforms, progress=False)
loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = UNet(spatial_dims=2,
in_channels=1,
out_channels=1,
channels=(2, 4),
strides=(1, )).to(device)
data = first(loader)
self.assertEqual(data["image"].shape[0], batch_size * NUM_SAMPLES)
labels = data["label"].to(device)
self.assertIsInstance(labels, MetaTensor)
segs = model(labels).detach().cpu()
segs_decollated = decollate_batch(segs)
self.assertIsInstance(segs_decollated[0], MetaTensor)
# inverse of individual segmentation
seg_metatensor = first(segs_decollated)
# test to convert interpolation mode for 1 data of model output batch
convert_applied_interp_mode(seg_metatensor.applied_operations,
mode="nearest",
align_corners=None)
# manually invert the last crop samples
xform = seg_metatensor.applied_operations.pop(-1)
shape_before_extra_xform = xform["orig_size"]
resizer = ResizeWithPadOrCrop(spatial_size=shape_before_extra_xform)
with resizer.trace_transform(False):
seg_metatensor = resizer(seg_metatensor)
with allow_missing_keys_mode(transforms):
inv_seg = transforms.inverse({"label": seg_metatensor})["label"]
self.assertEqual(inv_seg.shape[1:], test_data[0]["label"].shape)
def test_pad_collation(self, t_type, transform):
if t_type == dict:
dataset = CacheDataset(self.dict_data, transform, progress=False)
else:
dataset = _Dataset(self.list_data, self.list_labels, transform)
# Default collation should raise an error
loader_fail = DataLoader(dataset, batch_size=10)
with self.assertRaises(RuntimeError):
for _ in loader_fail:
pass
# Padded collation shouldn't
loader = DataLoader(dataset,
batch_size=2,
collate_fn=pad_list_data_collate)
for _ in loader:
pass
def test_loading_array(self):
set_determinism(seed=1234)
# test sequence input data with images
images = [
np.arange(16, dtype=float).reshape(1, 4, 4),
np.arange(16, dtype=float).reshape(1, 4, 4)
]
# image level
patch_intensity = RandShiftIntensity(offsets=1.0, prob=1.0)
patch_iter = PatchIter(patch_size=(2, 2), start_pos=(0, 0))
ds = GridPatchDataset(data=images,
patch_iter=patch_iter,
transform=patch_intensity)
# use the grid patch dataset
for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=0):
np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2))
np.testing.assert_allclose(
item[0],
np.array([[[[1.4965, 2.4965], [5.4965, 6.4965]]],
[[[11.3584, 12.3584], [15.3584, 16.3584]]]]),
rtol=1e-4,
)
np.testing.assert_allclose(item[1],
np.array([[[0, 1], [0, 2], [2, 4]],
[[0, 1], [2, 4], [2, 4]]]),
rtol=1e-5)
if sys.platform != "win32":
for item in DataLoader(ds,
batch_size=2,
shuffle=False,
num_workers=2):
np.testing.assert_equal(tuple(item[0].shape), (2, 1, 2, 2))
np.testing.assert_allclose(
item[0],
np.array([[[[1.2548, 2.2548], [5.2548, 6.2548]]],
[[[9.1106, 10.1106], [13.1106, 14.1106]]]]),
rtol=1e-3,
)
np.testing.assert_allclose(item[1],
np.array([[[0, 1], [0, 2], [2, 4]],
[[0, 1], [2, 4], [2, 4]]]),
rtol=1e-5)
def test_values(self):
dataset = Dataset(data=self.datalist, transform=self.transform)
dataloader = DataLoader(dataset=dataset, batch_size=2, num_workers=0)
tbuffer = ThreadBuffer(dataloader)
for d in tbuffer:
self.assertEqual(d["image"][0], "spleen_19.nii.gz")
self.assertEqual(d["image"][1], "spleen_31.nii.gz")
self.assertEqual(d["label"][0], "spleen_label_19.nii.gz")
self.assertEqual(d["label"][1], "spleen_label_31.nii.gz")
def _define_prediction_data_loaders(
self,
prediction_folder_path: Union[Path, str]
) -> bool:
"""Initialize prediction datasets and data loaders.
@Note: in Windows, it is essential to set `persistent_workers=True` in the data loaders!
@return True if datasets and data loaders could be instantiated, False otherwise.
"""
# Check that the path exists
prediction_folder_path = Path(prediction_folder_path)
if not prediction_folder_path.is_dir():
return False
# Scan for images
self._prediction_image_names = natsorted(
glob(str(Path(prediction_folder_path) / "*.tif"))
)
# Optimize arguments
if sys.platform == 'win32':
persistent_workers = True
pin_memory = False
else:
persistent_workers = False
pin_memory = torch.cuda.is_available()
if len(self._prediction_image_names) == 0:
self._prediction_dataset = None
self._prediction_dataloader = None
return False
# Define the transforms
self._define_prediction_transforms()
# Prediction
self._prediction_dataset = Dataset(
self._prediction_image_names,
self._prediction_image_transforms
)
self._prediction_dataloader = DataLoader(
self._prediction_dataset,
batch_size=self._test_batch_size,
shuffle=False,
num_workers=self._test_num_workers,
persistent_workers=persistent_workers,
pin_memory=pin_memory
)
return True
def test_data_loader_2(self):
set_determinism(seed=123)
xform_2 = Compose([_RandXform(), _RandXform()])
train_ds = Dataset([1], transform=xform_2)
out_2 = train_ds[0]
self.assertAlmostEqual(out_2, 0.4092510)
train_loader = DataLoader(train_ds, num_workers=0)
out_2 = next(iter(train_loader))
self.assertAlmostEqual(out_2.cpu().item(), 0.7858843729)
if sys.platform != "win32": # skip multi-worker tests on win32
train_loader = DataLoader(train_ds, num_workers=1)
out_2 = next(iter(train_loader))
self.assertAlmostEqual(out_2.cpu().item(), 0.305763411)
train_loader = DataLoader(train_ds, num_workers=2)
out_1 = next(iter(train_loader))
self.assertAlmostEqual(out_1.cpu().item(), 0.131966779)
set_determinism(None)
def test_pad_collation(self, t_type, collate_method, transform):
if t_type == dict:
dataset = CacheDataset(self.dict_data, transform, progress=False)
else:
dataset = _Dataset(self.list_data, self.list_labels, transform)
# Default collation should raise an error
loader_fail = DataLoader(dataset, batch_size=10)
with self.assertRaises(RuntimeError):
for _ in loader_fail:
pass
# Padded collation shouldn't
loader = DataLoader(dataset, batch_size=10, collate_fn=collate_method)
# check collation in forward direction
for data in loader:
if t_type == dict:
decollated_data = decollate_batch(data)
for d in decollated_data:
PadListDataCollate.inverse(d)
def test_shape(self):
test_dataset = ["vwxyz", "hello", "world"]
n_per_image = len(test_dataset[0])
result = PatchDataset(dataset=test_dataset, patch_func=identity, samples_per_image=n_per_image)
output = []
n_workers = 0 if sys.platform == "win32" else 2
for item in DataLoader(result, batch_size=3, num_workers=n_workers):
output.append("".join(item))
expected = ["vwx", "yzh", "ell", "owo", "rld"]
self.assertEqual(output, expected)
def test_decollate(self, dtype):
batch_size = 3
ims = [self.get_im(dtype=dtype)[0] for _ in range(batch_size * 2)]
ds = Dataset(ims)
dl = DataLoader(ds, num_workers=batch_size, batch_size=batch_size)
batch = next(iter(dl))
decollated = decollate_batch(batch)
self.assertIsInstance(decollated, list)
self.assertEqual(len(decollated), batch_size)
for elem, im in zip(decollated, ims):
self.assertIsInstance(elem, MetaTensor)
self.check(elem, im, ids=False)
def test_endianness(self, endianness, use_array, image_only):
hdr = nib.Nifti1Header(endianness=endianness)
nii = nib.Nifti1Image(self.im, np.eye(4), header=hdr)
nib.save(nii, self.fname)
data = [self.fname] if use_array else [{"image": self.fname}]
tr = LoadImage(image_only=image_only) if use_array else LoadImaged(
"image", image_only=image_only)
check_ds = Dataset(data, tr)
check_loader = DataLoader(check_ds, batch_size=1)
_ = next(iter(check_loader))
