tests.data.test_queue

Complexity

Total Complexity

3

Size/Duplication

Total Lines	36
Duplicated Lines	0 %

Importance

Changes

0

2 Methods

Rating	Name	Duplication	Size	Complexity
A	TestQueue.test_queue()	0	16	2
A	TestQueue.setUp()	0	8	1

from torch.utils.data import DataLoader
from torchio.data import ImageSampler
from torchio import ImagesDataset, Queue, DATA
from torchio.utils import create_dummy_dataset
from ..utils import TorchioTestCase


class TestQueue(TorchioTestCase):
    """Tests for `queue` module."""
    def setUp(self):
        super().setUp()
        self.subjects_list = create_dummy_dataset(
            num_images=10,
            size_range=(10, 20),
            directory=self.dir,
            suffix='.nii',
            force=False,
        )

    def test_queue(self):
        subjects_dataset = ImagesDataset(self.subjects_list)
        queue_dataset = Queue(
            subjects_dataset,
            max_length=6,
            samples_per_volume=2,
            patch_size=10,
            sampler_class=ImageSampler,
            num_workers=0,
            verbose=True,
        )
        _ = str(queue_dataset)
        batch_loader = DataLoader(queue_dataset, batch_size=4)
        for batch in batch_loader:
            _ = batch['one_modality'][DATA]
            _ = batch['segmentation'][DATA]