tests.transforms.preprocessing.test_keep_largest

Complexity

Total Complexity

3

Size/Duplication

Total Lines	18
Duplicated Lines	0 %

Importance

Changes

0

2 Methods

Rating	Name	Duplication	Size	Complexity
A	TestKeepLargestComponent.test_multichannel()	0	4	2
A	TestKeepLargestComponent.test_one_hot()	0	5	1

import torch
import torchio as tio
from ...utils import TorchioTestCase


class TestKeepLargestComponent(TorchioTestCase):
    """Tests for `KeepLargestComponent`."""
    def test_one_hot(self):
        tensor = torch.as_tensor([1, 0, 1, 1, 0, 1]).reshape(1, 1, 1, 6)
        label_map = tio.LabelMap(tensor=tensor)
        largest = tio.KeepLargestComponent()(label_map)
        assert largest.data.sum() == 2

    def test_multichannel(self):
        label_map = tio.LabelMap(tensor=torch.rand(2, 3, 3, 3) > 1)
        with self.assertRaises(RuntimeError):
            tio.KeepLargestComponent()(label_map)