torchio.data.sampler.label

Complexity

Total Complexity

3

Size/Duplication

Total Lines	32
Duplicated Lines	0 %

Importance

Changes

0

2 Methods

Rating	Name	Duplication	Size	Complexity
A	LabelSampler.__init__()	0	2	1
A	LabelSampler.get_probability_map()	0	11	2

from .weighted import WeightedSampler


class LabelSampler(WeightedSampler):
    r"""Extract random patches containing labeled voxels.

    This iterable dataset yields patches whose center value is greater than 0
    in the :py:attr:`label_name`.

    Args:
        patch_size: Tuple of integers :math:`(d, h, w)` to generate patches
            of size :math:`d \times h \times w`.
            If a single number :math:`n` is provided,
            :math:`d = h = w = n`.
        label_name: Name of the label image in the sample that will be used to
            generate the sampling probability map.
    """
    def __init__(self, patch_size, label_name):
        super().__init__(patch_size, probability_map=label_name)

    def get_probability_map(self, sample):
        """Return binarized image for sampling."""
        if self.probability_map_name in sample:
            data = sample[self.probability_map_name].data > 0.5
        else:
            message = (
                f'Image "{self.probability_map_name}"'
                f' not found in subject sample: {sample}'
            )
            raise KeyError(message)
        return data