torchio.data.sampler.uniform

Complexity

Total Complexity

4

Size/Duplication

Total Lines	35
Duplicated Lines	0 %

Importance

Changes

0

3 Methods

Rating	Name	Duplication	Size	Complexity
A	UniformSampler.__init__()	0	2	1
A	UniformSampler.__call__()	0	14	2
A	UniformSampler.get_probability_map()	0	2	1

import torch
from ...data.subject import Subject
from ...torchio import TypePatchSize
from .sampler import RandomSampler
from typing import Optional, Tuple, Generator
import numpy as np


class UniformSampler(RandomSampler):
    """Randomly extract patches from a volume with uniform probability.

    Args:
        patch_size: See :py:class:`~torchio.data.PatchSampler`.
    """
    def __init__(self, patch_size: TypePatchSize):
        super().__init__(patch_size)

    def get_probability_map(self, sample: Subject) -> torch.Tensor:
        return torch.ones(1, *sample.spatial_shape)

    def __call__(self, sample: Subject) -> Generator[Subject, None, None]:

        sample.check_consistent_spatial_shape()

        if np.any(self.patch_size > sample.spatial_shape):
            message = (
                f'Patch size {tuple(self.patch_size)} cannot be'
                f' larger than image size {tuple(sample.spatial_shape)}'
            )
            raise RuntimeError(message)

        valid_range = sample.spatial_shape - self.patch_size
        corners = np.random.randint(valid_range + 1)
        yield self.extract_patch(sample, corners)