| 101 | class TestTupleDataset(unittest.TestCase): |
| 102 | @parameterized.expand([TEST_CASE_1]) |
| 103 | def test_shape(self, expected_shape): |
| 104 | test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]).astype(float), np.eye(4)) |
| 105 | with tempfile.TemporaryDirectory() as tempdir: |
| 106 | nib.save(test_image, os.path.join(tempdir, "test_image1.nii.gz")) |
| 107 | nib.save(test_image, os.path.join(tempdir, "test_label1.nii.gz")) |
| 108 | nib.save(test_image, os.path.join(tempdir, "test_image2.nii.gz")) |
| 109 | nib.save(test_image, os.path.join(tempdir, "test_label2.nii.gz")) |
| 110 | test_data = [ |
| 111 | (os.path.join(tempdir, "test_image1.nii.gz"), os.path.join(tempdir, "test_label1.nii.gz")), |
| 112 | (os.path.join(tempdir, "test_image2.nii.gz"), os.path.join(tempdir, "test_label2.nii.gz")), |
| 113 | ] |
| 114 | |
| 115 | test_transform = Compose([LoadImage(), SimulateDelay(delay_time=1e-5)]) |
| 116 | |
| 117 | # Here test_transform is applied element by element for the tuple. |
| 118 | dataset = Dataset(data=test_data, transform=test_transform) |
| 119 | data1 = dataset[0] |
| 120 | data2 = dataset[1] |
| 121 | |
| 122 | # Output is a list/tuple |
| 123 | self.assertTrue(isinstance(data1, (list, tuple))) |
| 124 | self.assertTrue(isinstance(data2, (list, tuple))) |
| 125 | |
| 126 | # Number of elements are 2 |
| 127 | self.assertEqual(len(data1), 2) |
| 128 | self.assertEqual(len(data2), 2) |
| 129 | |
| 130 | # Output shapes are as expected |
| 131 | self.assertTupleEqual(data1[0].shape, expected_shape) |
| 132 | self.assertTupleEqual(data1[1].shape, expected_shape) |
| 133 | self.assertTupleEqual(data2[0].shape, expected_shape) |
| 134 | self.assertTupleEqual(data2[1].shape, expected_shape) |
| 135 | |
| 136 | # Here test_transform is applied to the tuple as a whole. |
| 137 | test_transform = Compose( |
| 138 | [ |
| 139 | # LoadImage creates a channel-stacked image when applied to a tuple |
| 140 | LoadImage(), |
| 141 | # Get the channel-stacked image and the label |
| 142 | Lambda(func=lambda x: (x[0].permute(2, 1, 0), x[1])), |
| 143 | ], |
| 144 | map_items=False, |
| 145 | ) |
| 146 | |
| 147 | dataset = Dataset(data=test_data, transform=test_transform) |
| 148 | data1 = dataset[0] |
| 149 | data2 = dataset[1] |
| 150 | |
| 151 | # Output is a list/tuple |
| 152 | self.assertTrue(isinstance(data1, (list, tuple))) |
| 153 | self.assertTrue(isinstance(data2, (list, tuple))) |
| 154 | |
| 155 | # Number of elements are 2 |
| 156 | self.assertEqual(len(data1), 2) |
| 157 | self.assertEqual(len(data2), 2) |
| 158 | |
| 159 | # Output shapes are as expected |
| 160 | self.assertTupleEqual(data1[0].shape, expected_shape) |