Passed
Push — master ( 8bc331...f26751 )
by Fernando
01:12
created

torchio.datasets.bite.BITE._download()   A

Complexity

Conditions 1

Size

Total Lines 2
Code Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 2
dl 0
loc 2
rs 10
c 0
b 0
f 0
cc 1
nop 2
1
import abc
2
from pathlib import Path
3
from typing import Optional
4
5
from ..typing import TypePath
6
from ..transforms import Transform
7
from ..download import download_and_extract_archive
8
from .. import SubjectsDataset, Subject, ScalarImage, LabelMap
9
10
11
class BITE(SubjectsDataset, abc.ABC):
12
    base_url = 'http://www.bic.mni.mcgill.ca/uploads/Services/'
13
14
    def __init__(
15
            self,
16
            root: TypePath,
17
            transform: Optional[Transform] = None,
18
            download: bool = False,
19
            **kwargs,
20
            ):
21
        root = Path(root).expanduser().absolute()
22
        if download:
23
            self._download(root)
24
        subjects_list = self._get_subjects_list(root)
25
        self.kwargs = kwargs
26
        super().__init__(subjects_list, transform=transform, **kwargs)
27
28
    def _download(self, root: Path):
29
        raise NotImplementedError
30
31
    def _get_subjects_list(self, root: Path):
32
        raise NotImplementedError
33
34
35
class BITE3(BITE):
36
    dirname = 'group3'
37
    """Pre- and post-resection MR images in BITE.
38
39
    *The goal of BITE is to share in vivo medical images of patients wtith
40
    brain tumors to facilitate the development and validation of new image
41
    processing algorithms.*
42
43
    Please check the `BITE website`_ for more information and
44
    acknowledgments instructions.
45
46
    .. _BITE website: http://nist.mni.mcgill.ca/?page_id=672
47
48
    Args:
49
        root: Root directory to which the dataset will be downloaded.
50
        transform: An instance of
51
            :class:`~torchio.transforms.transform.Transform`.
52
        download: If set to ``True``, will download the data into :attr:`root`.
53
    """
54
    def _download(self, root: Path):
55
        if (root / self.dirname).is_dir():
56
            return
57
        root.mkdir(exist_ok=True, parents=True)
58
        filename = f'{self.dirname}.tar.gz'
59
        url = self.base_url + filename
60
        download_and_extract_archive(
61
            url,
62
            download_root=root,
63
            md5='e415b63887c40b727c45552614b44634',
64
        )
65
        (root / filename).unlink()  # cleanup
66
67
    def _get_subjects_list(self, root: Path):
68
        subjects_dir = root / self.dirname
69
        subjects = []
70
        for i in range(1, 15):
71
            if i == 13:
72
                continue  # no MRI for this subject
73
            subject_id = f'{i:02d}'
74
            subject_dir = subjects_dir / subject_id
75
            preop_path = subject_dir / f'{subject_id}_preop_mri.mnc'
76
            postop_path = subject_dir / f'{subject_id}_postop_mri.mnc'
77
            images_dict = {}
78
            images_dict['preop'] = ScalarImage(preop_path)
79
            images_dict['postop'] = ScalarImage(postop_path)
80
            for fp in subject_dir.glob('*tumor*'):
81
                images_dict[fp.stem[3:]] = LabelMap(fp)
82
            subject = Subject(images_dict)
83
            subjects.append(subject)
84
        return subjects
85