tests.data.sampler.test_uniform_sampler

Complexity

Total Complexity

2

Size/Duplication

Total Lines	26
Duplicated Lines	0 %

Importance

Changes

0

2 Methods

Rating	Name	Duplication	Size	Complexity
A	TestUniformSampler.test_processed_uniform_probabilities()	0	9	1
A	TestUniformSampler.test_uniform_probabilities()	0	5	1

import torch
import torchio
import numpy as np
from torchio.data import UniformSampler
from ...utils import TorchioTestCase


class TestUniformSampler(TorchioTestCase):
    """Tests for `UniformSampler` class."""

    def test_uniform_probabilities(self):
        sampler = UniformSampler(5)
        probabilities = sampler.get_probability_map(self.sample)
        fixtures = torch.ones_like(probabilities)
        assert torch.all(probabilities.eq(fixtures))

    def test_processed_uniform_probabilities(self):
        sampler = UniformSampler(5)
        probabilities = sampler.get_probability_map(self.sample)
        probabilities = sampler.process_probability_map(probabilities)
        fixtures = np.zeros_like(probabilities)
        # Other positions cannot be patch centers
        fixtures[2:-2, 2:-2, 2:-2] = probabilities[2, 2, 2]
        self.assertAlmostEqual(probabilities.sum(), 1)
        assert np.equal(probabilities, fixtures).all()