ordinal_regression/DataModule.py at main · paddyOnGithub/ordinal_regression · 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
from torch.utils.data import DataLoader
import pytorch_lightning as pl
import CONFIG
from CXRDataset import CXRDataset

class DataModule(pl.LightningDataModule):

    def __init__(
            self,
            numWorker,
            batchSize,
            targetFunc,
            fiveFold,
            runOnHpc,
        ):
        super().__init__()
        self.numWorker = numWorker
        self.batchSize = batchSize
        self.targetFunc = targetFunc
        self.fiveFold = fiveFold

        if runOnHpc:
            self.metaFilePath = CONFIG.HPC_META_FILE_PATH
            self.imgFolderPath = CONFIG.HPC_IMG_FOLDER_PATH
        else:
            self.metaFilePath = CONFIG.LOCAL_META_FILE_PATH
            self.imgFolderPath = CONFIG.LOCAL_IMG_FOLDER_PATH

    def setup(self,stage):
        self.trainDataset = CXRDataset(self.metaFilePath,self.imgFolderPath,"train",self.targetFunc,self.fiveFold)

        self.validDataset = CXRDataset(self.metaFilePath,self.imgFolderPath,"valid",self.targetFunc,self.fiveFold)

        self.testDataset = CXRDataset(self.metaFilePath,self.imgFolderPath,"test",self.targetFunc,self.fiveFold)

    def train_dataloader(self):
        trainLoader = DataLoader(self.trainDataset,batch_size=self.batchSize,num_workers=self.numWorker,shuffle=True)
        return trainLoader

    def val_dataloader(self):
        validLoader = DataLoader(self.validDataset,batch_size=self.batchSize,num_workers=self.numWorker,shuffle=False)
        return validLoader

    def test_dataloader(self):
        testLoader = DataLoader(self.testDataset,batch_size=self.batchSize,num_workers=self.numWorker,shuffle=False)
        return testLoader