PatternAnalysis-2023/dataset.py at ccdc0e57ae5e9054d6246bad9788a24899fe25eb · shakes76/PatternAnalysis-2023 · GitHub

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
import numpy as np
import torch
from PIL import Image
import os
import torchvision.transforms as transforms
from torch.utils.data import TensorDataset, DataLoader, random_split

transform = transforms.Compose([
    transforms.PILToTensor()
])

xtrain = []
xtest = []
ytrain = []
ytest = []
slicemax = 20 #20 images per patient

ntrainimgs_AD = 0
patient = []
slice = 0
for filename in sorted(os.listdir('../ADNI_AD_NC_2D/AD_NC/train/AD/')):
    f = os.path.join('../ADNI_AD_NC_2D/AD_NC/train/AD/', filename)
    img = Image.open(f)
    imgtorch = transform(img).float()
    imgtorch.require_grad = True
    patient.append(imgtorch/255) #go from 0,255 to 0,1
    slice  = (slice+1) % slicemax
    if slice == 0:
        xtrain.append(torch.stack(patient))
        patient = []
        ntrainimgs_AD += 1
    pass
ntrainimgs_NC = 0
patient = []
slice = 0
for filename in sorted(os.listdir('../ADNI_AD_NC_2D/AD_NC/train/NC')):
    f = os.path.join('../ADNI_AD_NC_2D/AD_NC/train/NC', filename)
    img = Image.open(f)
    imgtorch = transform(img).float()
    imgtorch.require_grad = True
    patient.append(imgtorch/255) #go from 0,255 to 0,1
    slice  = (slice+1) % slicemax
    if slice == 0:
        xtrain.append(torch.stack(patient))
        patient = []
        ntrainimgs_NC += 1
    pass
ntestimgs_AD = 0
patient = []
slice = 0
for filename in sorted(os.listdir('../ADNI_AD_NC_2D/AD_NC/test/AD')):
    f = os.path.join('../ADNI_AD_NC_2D/AD_NC/test/AD', filename)
    img = Image.open(f)
    imgtorch = transform(img).float()
    imgtorch.require_grad = True
    patient.append(imgtorch/255) #go from 0,255 to 0,1
    slice  = (slice+1) % slicemax
    if slice == 0:
        xtest.append(torch.stack(patient))
        patient = []
        ntestimgs_AD += 1
    pass
ntestimgs_NC = 0
patient = []
slice = 0
for filename in sorted(os.listdir('../ADNI_AD_NC_2D/AD_NC/test/NC')):
    f = os.path.join('../ADNI_AD_NC_2D/AD_NC/test/NC', filename)
    img = Image.open(f)
    imgtorch = transform(img).float()
    imgtorch.require_grad = True
    patient.append(imgtorch/255) #go from 0,255 to 0,1
    slice  = (slice+1) % slicemax
    if slice == 0:
        xtest.append(torch.stack(patient))
        patient = []
        ntestimgs_NC += 1
    pass
xtrain = torch.stack(xtrain)
xtest = torch.stack(xtest)
ytrain = torch.from_numpy(np.concatenate((np.ones(ntrainimgs_AD), np.zeros(ntrainimgs_NC)), axis=0)).type(torch.LongTensor)
ytest = torch.from_numpy(np.concatenate((np.ones(ntestimgs_AD), np.zeros(ntestimgs_NC)), axis=0)).type(torch.LongTensor)

data_val, data_test = random_split(TensorDataset(xtest, ytest), [0.7,0.3])
dataloader_train = DataLoader(TensorDataset(xtrain, ytrain), batch_size=32, shuffle=True)
dataloader_test = DataLoader(data_test, batch_size=32, shuffle=True)
dataloader_val = DataLoader(data_val, batch_size=32, shuffle=True)

def returnDataLoaders():
    return dataloader_train, dataloader_test, dataloader_val