forked from MontaEllis/Pytorch-Medical-Classification
-
Notifications
You must be signed in to change notification settings - Fork 0
/
data_function.py
executable file
·145 lines (107 loc) · 3.61 KB
/
data_function.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
from glob import glob
from os.path import dirname, join, basename, isfile
import sys
sys.path.append('./')
import csv
import torch
from medpy.io import load
import numpy as np
from PIL import Image
from torch import nn
import torch.nn.functional as F
import random
import torchio as tio
from torchio import AFFINE, DATA
import torchio
from torchio import ScalarImage, LabelMap, Subject, SubjectsDataset, Queue
from torchio.data import UniformSampler
from torchio.transforms import (
RandomFlip,
RandomAffine,
RandomElasticDeformation,
RandomNoise,
RandomMotion,
RandomBiasField,
RescaleIntensity,
Resample,
ToCanonical,
ZNormalization,
CropOrPad,
HistogramStandardization,
OneOf,
Compose,
)
from pathlib import Path
from hparam import hparams as hp
class MedData_train(torch.utils.data.Dataset):
def __init__(self, images_dir_0, images_dir_1):
self.subjects = []
images_dir_0 = Path(images_dir_0)
self.image_paths_0 = sorted(images_dir_0.glob(hp.fold_arch))
images_dir_1 = Path(images_dir_1)
self.image_paths_1 = sorted(images_dir_1.glob(hp.fold_arch))
for (image_path) in zip(self.image_paths_0):
subject = tio.Subject(
source=tio.ScalarImage(image_path),
label= 0,
)
self.subjects.append(subject)
for (image_path) in zip(self.image_paths_1):
subject = tio.Subject(
source=tio.ScalarImage(image_path),
label= 1,
)
self.subjects.append(subject)
self.transforms = self.transform()
self.training_set = tio.SubjectsDataset(self.subjects, transform=self.transforms)
# one_subject = self.training_set[0]
# one_subject.plot()
def transform(self):
if hp.aug:
training_transform = Compose([
CropOrPad((hp.crop_or_pad_size), padding_mode='reflect'),
# ToCanonical(),
RandomBiasField(),
ZNormalization(),
RandomNoise(),
RandomFlip(axes=(0,)),
OneOf({
RandomAffine(): 0.8,
RandomElasticDeformation(): 0.2,
}),
])
else:
training_transform = Compose([
CropOrPad((hp.crop_or_pad_size), padding_mode='reflect'),
ZNormalization(),
])
return training_transform
class MedData_test(torch.utils.data.Dataset):
def __init__(self, images_dir_0, images_dir_1):
self.subjects = []
images_dir_0 = Path(images_dir_0)
self.image_paths_0 = sorted(images_dir_0.glob(hp.fold_arch))
images_dir_1 = Path(images_dir_1)
self.image_paths_1 = sorted(images_dir_1.glob(hp.fold_arch))
for (image_path) in zip(self.image_paths_0):
subject = tio.Subject(
source=tio.ScalarImage(image_path),
label= 0,
)
self.subjects.append(subject)
for (image_path) in zip(self.image_paths_1):
subject = tio.Subject(
source=tio.ScalarImage(image_path),
label= 1,
)
self.subjects.append(subject)
self.transforms = self.transform()
self.testing_set = tio.SubjectsDataset(self.subjects, transform=self.transforms)
# one_subject = self.training_set[0]
# one_subject.plot()
def transform(self):
testing_transform = Compose([
CropOrPad((hp.crop_or_pad_size), padding_mode='reflect'),
ZNormalization(),
])
return testing_transform