def test_shape(self, img_transform, label_transform, indexes, expected_shape):
test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4))
tempdir = tempfile.mkdtemp()
test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
test_seg1 = os.path.join(tempdir, "test_seg1.nii.gz")
test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
test_seg2 = os.path.join(tempdir, "test_seg2.nii.gz")
nib.save(test_image, test_image1)
nib.save(test_image, test_seg1)
nib.save(test_image, test_image2)
nib.save(test_image, test_seg2)
test_images = [test_image1, test_image2]
test_segs = [test_seg1, test_seg2]
test_labels = [1, 1]
dataset = ArrayDataset(test_images, img_transform, test_segs, label_transform, test_labels, None)
self.assertEqual(len(dataset), 2)
dataset.set_random_state(1234)
data1 = dataset[0]
data2 = dataset[1]
self.assertTupleEqual(data1[indexes[0]].shape, expected_shape)
self.assertTupleEqual(data1[indexes[1]].shape, expected_shape)
np.testing.assert_allclose(data1[indexes[0]], data1[indexes[1]])
self.assertTupleEqual(data2[indexes[0]].shape, expected_shape)
self.assertTupleEqual(data2[indexes[1]].shape, expected_shape)
np.testing.assert_allclose(data2[indexes[0]], data2[indexes[0]])
dataset = ArrayDataset(test_images, img_transform, test_segs, label_transform, test_labels, None)
dataset.set_random_state(1234)
_ = dataset[0]
data2_new = dataset[1]
np.testing.assert_allclose(data2[indexes[0]], data2_new[indexes[0]], atol=1e-3)
shutil.rmtree(tempdir)
def test_dataloading_img_label(self, img_transform, expected_shape):
test_image = nib.Nifti1Image(
np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4))
with tempfile.TemporaryDirectory() as tempdir:
test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
test_label1 = os.path.join(tempdir, "test_label1.nii.gz")
test_label2 = os.path.join(tempdir, "test_label2.nii.gz")
nib.save(test_image, test_image1)
nib.save(test_image, test_image2)
nib.save(test_image, test_label1)
nib.save(test_image, test_label2)
test_images = [test_image1, test_image2]
test_labels = [test_label1, test_label2]
dataset = ArrayDataset(test_images, img_transform, test_labels,
img_transform)
self.assertEqual(len(dataset), 2)
dataset.set_random_state(1234)
loader = DataLoader(dataset, batch_size=10, num_workers=1)
data = next(iter(loader)) # test batching
np.testing.assert_allclose(data[0].shape,
[2] + list(expected_shape))
dataset.set_random_state(1234)
new_data = next(iter(loader)) # test batching
np.testing.assert_allclose(data[0], new_data[0], atol=1e-3)
def test_dataloading(self, img_transform, expected_shape):
test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4))
tempdir = tempfile.mkdtemp()
test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
nib.save(test_image, test_image1)
nib.save(test_image, test_image2)
test_images = [test_image1, test_image2]
dataset = ArrayDataset(test_images, img_transform)
self.assertEqual(len(dataset), 2)
dataset.set_random_state(1234)
loader = DataLoader(dataset, batch_size=10, num_workers=1)
imgs = next(iter(loader)) # test batching
np.testing.assert_allclose(imgs.shape, [2] + list(expected_shape))
dataset.set_random_state(1234)
new_imgs = next(iter(loader)) # test batching
np.testing.assert_allclose(imgs, new_imgs, atol=1e-3)
def test_default_none(self, img_transform, expected_shape):
test_image = nib.Nifti1Image(np.random.randint(0, 2, size=(128, 128, 128)), np.eye(4))
with tempfile.TemporaryDirectory() as tempdir:
test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
nib.save(test_image, test_image1)
nib.save(test_image, test_image2)
test_images = [test_image1, test_image2]
dataset = ArrayDataset(test_images, img_transform)
self.assertEqual(len(dataset), 2)
dataset.set_random_state(1234)
data1 = dataset[0]
data2 = dataset[1]
self.assertTupleEqual(data1.shape, expected_shape)
self.assertTupleEqual(data2.shape, expected_shape)
dataset = ArrayDataset(test_images, img_transform)
dataset.set_random_state(1234)
_ = dataset[0]
data2_new = dataset[1]
np.testing.assert_allclose(data2, data2_new, atol=1e-3)
def test_dataloading_img(self, img_transform, expected_shape):
test_image = nib.Nifti1Image(
np.random.randint(0, 2, size=(128, 128, 128)).astype(float),
np.eye(4))
with tempfile.TemporaryDirectory() as tempdir:
test_image1 = os.path.join(tempdir, "test_image1.nii.gz")
test_image2 = os.path.join(tempdir, "test_image2.nii.gz")
nib.save(test_image, test_image1)
nib.save(test_image, test_image2)
test_images = [test_image1, test_image2]
dataset = ArrayDataset(test_images, img_transform)
self.assertEqual(len(dataset), 2)
dataset.set_random_state(1234)
n_workers = 0 if sys.platform == "win32" else 2
loader = DataLoader(dataset, batch_size=10, num_workers=n_workers)
imgs = next(iter(loader)) # test batching
np.testing.assert_allclose(imgs.shape, [2] + list(expected_shape))
dataset.set_random_state(1234)
new_imgs = next(iter(loader)) # test batching
np.testing.assert_allclose(imgs, new_imgs, atol=1e-3)
