forked from yiskw713/fastfcn
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataloader.py
83 lines (55 loc) · 2.35 KB
/
dataloader.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
from torch.utils.data.dataset import Dataset
import torch
from torchvision.transforms import transforms
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data import DataLoader
from torch.optim.rmsprop import RMSprop
import numpy as np
import pytorch_lightning as pl
from configs import Configs
configs = Configs()
class CustomDataset(Dataset):
def __init__(self, configs: Configs):
self.device = configs.device
self.path = configs.datasetPath
self.size = (configs.image_size, configs.image_size)
self.data = torch.load(self.path)
self.length = len(self.data)
print(self.length, 'images in', self.path)
self.transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(self.size),
transforms.ToTensor(),
])
def __getitem__(self, index):
if index > self.length:
raise Exception(
f"Dataloader out of index. Max Index:{self.length - self.n_images}, Index asked:{index}.")
image = self.transform(self.data[index]['front']['rgb']).float()
seg = self.transform(self.data[index]['front']['seg']).squeeze().long()
# images.shape = [3,128,128]
return {'input': image, 'target': seg}
def __len__(self):
return self.length
class lit_custom_data(pl.LightningDataModule):
def setup(self, stage):
self.configs = Configs()
self.dataset = CustomDataset(self.configs)
dataset_size = len(self.dataset)
indices = list(range(dataset_size))
split = int(np.floor(self.configs.valSplit * dataset_size))
self.trainIndices, self.valIndices = indices[split:], indices[:split]
def train_dataloader(self):
return DataLoader(self.dataset, batch_size=self.configs.batchSize,
num_workers=0, sampler=SubsetRandomSampler(self.trainIndices))
def val_dataloader(self):
return DataLoader(self.dataset, batch_size=self.configs.batchSize,
num_workers=0, sampler=SubsetRandomSampler(self.valIndices))
if __name__ == "__main__":
device = 'cuda' if torch.cuda.is_available() else 'cpu'
path = r"dataset.pt"
# cd = CustomDataset(config)
# print(cd[0])
data_module = lit_custom_data()
print("python")
# print(cd[0]['front']['seg'].shape)