Passed
Push — master ( d5bb42...1d7e64 )
by Fernando
58s
created

AorticValve.get_kwargs()   A

Complexity

Conditions 1

Size

Total Lines 10
Code Lines 9

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 1
eloc 9
nop 1
dl 0
loc 10
rs 9.95
c 0
b 0
f 0
1
import urllib.parse
2
from torchvision.datasets.utils import download_and_extract_archive
3
from ...utils import get_torchio_cache_dir
4
from ... import Subject, Image, LABEL
5
6
7
class SubjectITKSNAP(Subject):
8
    """ITK-SNAP Image Data Downloads.
9
10
    See `the itk-SNAP website <http://www.itksnap.org/pmwiki/pmwiki.php?n=Downloads.Data>`_
11
    for more information.
12
    """
13
    url_base = 'https://www.nitrc.org/frs/download.php/'
14
15
    def __init__(self, name, code):
16
        self.name = name
17
        self.url_dir = urllib.parse.urljoin(self.url_base, f'{code}/')
18
        self.filename = f'{self.name}.zip'
19
        self.url = urllib.parse.urljoin(self.url_dir, self.filename)
20
        self.download_root = get_torchio_cache_dir() / self.name
21
        if self.download_root.is_dir():
22
            print(f'Using cache found in {self.download_root}')
23
        else:
24
            download_and_extract_archive(
25
                self.url,
26
                download_root=self.download_root,
27
                filename=self.filename,
28
            )
29
        super().__init__(**self.get_kwargs())
30
31
    def get_kwargs(self):
32
        raise NotImplementedError
33
34
35
class BrainTumor(SubjectITKSNAP):
36
    def __init__(self):
37
        super().__init__('braintumor', '6161')
38
39
    def get_kwargs(self):
40
        t1, t1c, t2, flair, seg = [
41
            self.download_root / self.name / f'BRATS_HG0015_{name}.mha'
42
            for name in ('T1', 'T1C', 'T2', 'FLAIR', 'truth')
43
        ]
44
        return dict(
45
            t1=Image(t1),
46
            t1c=Image(t1c),
47
            t2=Image(t2),
48
            flair=Image(flair),
49
            seg=Image(seg, type=LABEL),
50
        )
51
52
53
class T1T2(SubjectITKSNAP):
54
    def __init__(self):
55
        super().__init__('ashs_test', '10983')
56
57
    def get_kwargs(self):
58
        mprage = self.download_root / self.name / 'mprage_3T_bet_dr.nii'
59
        tse = self.download_root / self.name / 'tse_3t_dr.nii'
60
        return dict(
61
            mprage=Image(mprage),
62
            tse=Image(tse),
63
        )
64
65
66
class AorticValve(SubjectITKSNAP):
67
    def __init__(self):
68
        super().__init__('bav_example', '11021')
69
70
    def get_kwargs(self):
71
        b14, b14_seg, b25, b25_seg = [
72
            self.download_root / self.name / f'bav_frame_{name}.nii.gz'
73
            for name in ('14', '14_manseg', '25', '25_manseg')
74
        ]
75
        return dict(
76
            b14=Image(b14),
77
            b14_seg=Image(b14_seg, type=LABEL),
78
            b25=Image(b25),
79
            b25_seg=Image(b25_seg, type=LABEL),
80
        )
81