Method test_shape

tests/data/test_smartcachedataset.py:42–88

(self, replace_rate, num_replace_workers, transform)

class TestSmartCacheDataset(unittest.TestCase):
    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
    def test_shape(self, replace_rate, num_replace_workers, transform):
        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[8, 8, 8]).astype(float), np.eye(4))
        with tempfile.TemporaryDirectory() as tempdir:
            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"))
            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"),
                }
            ] * 20
            dataset = SmartCacheDataset(
                data=test_data,
                transform=transform,
                replace_rate=replace_rate,
                cache_num=16,
                num_init_workers=4,
                num_replace_workers=num_replace_workers,
            )
            if transform is None:
                # Check without providing transform
                dataset2 = SmartCacheDataset(
                    data=test_data,
                    replace_rate=replace_rate,
                    cache_num=16,
                    num_init_workers=4,
                    num_replace_workers=num_replace_workers,
                )
                for k in ["image", "label", "extra"]:
                    self.assertEqual(dataset[0][k], dataset2[0][k])

            self.assertEqual(len(dataset._cache), dataset.cache_num)
            for i in range(dataset.cache_num):
                self.assertIsNotNone(dataset._cache[i])

            for _ in range(2):
                dataset.start()
                for _ in range(3):
                    dataset.update_cache()
                    self.assertIsNotNone(dataset[15])
                    if isinstance(dataset[15]["image"], (np.ndarray, torch.Tensor)):
                        assert_allclose(dataset[15]["image"], dataset[15]["label"])
                    else:
                        self.assertIsInstance(dataset[15]["image"], str)
                dataset.shutdown()

    def test_update_cache(self):
        # Given