-
Notifications
You must be signed in to change notification settings - Fork 0
/
7.py
118 lines (95 loc) · 3.2 KB
/
7.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
import ads3 as ads3
import torch
import torchvision.transforms as transforms
from pathlib import Path
from PIL import Image
from torch.utils import data as D
torch.manual_seed(0)
INPUT_SIZE = 224
root = Path("data")
file_train = root / "train.txt"
folder_images = root / "image"
class CarDataset(D.Dataset):
def __init__(self, labels: list):
self.filenames = []
self.labels = labels
"""Read the dataset index file"""
with open(file_train, newline="\n") as trainfile:
for line in trainfile:
self.filenames.append(folder_images / line.strip())
def __getitem__(self, index: int):
"""Get a sample from the dataset"""
image = Image.open(str(self.filenames[index]))
labelStr = self.filenames[index].parts[-3]
label = self.labels.index(labelStr)
return image, label
def __len__(self):
"""
Total number of samples in the dataset
"""
return len(self.filenames)
# In order to apply transformations specific to either training or validation data
# I use the following class. Inspired from https://stackoverflow.com/a/59615584
class DatasetFromSubset(D.Dataset):
def __init__(self, subset, transform=None):
self.subset = subset
self.transform = transform
def __getitem__(self, index):
x, y = self.subset[index]
if self.transform:
x = self.transform(x)
return x, y
def __len__(self):
return len(self.subset)
train_transforms = transforms.Compose(
[
transforms.Resize((INPUT_SIZE, INPUT_SIZE)),
transforms.RandomVerticalFlip(p=0.5),
transforms.RandomPerspective(p=0.5),
transforms.RandomApply(torch.nn.ModuleList([
transforms.ColorJitter(contrast=0.5, saturation=0.5, hue=0.5),
]), p=0.5),
transforms.RandomApply(torch.nn.ModuleList([
transforms.Grayscale(num_output_channels=3),
]), p=0.5),
transforms.ToTensor(),
]
)
valid_transforms = transforms.Compose(
[
transforms.Resize((INPUT_SIZE, INPUT_SIZE)),
transforms.ToTensor(),
]
)
if __name__ == "__main__":
"""Initialise dataset"""
labels = ads3.get_labels()
dataset = CarDataset(labels=labels)
log_name = "results/7.csv"
"""Split train and test"""
train_len = int(0.7 * len(dataset))
valid_len = len(dataset) - train_len
train, valid = D.random_split(dataset, lengths=[train_len, valid_len])
train = DatasetFromSubset(train, train_transforms)
valid = DatasetFromSubset(valid, valid_transforms)
# When running image augmentation you should define seperate training and validation!
print("train size: %d, valid size %d" % (len(train), len(valid)))
loader_train = D.DataLoader(
train,
batch_size=80,
shuffle=True,
num_workers=1,
pin_memory=True,
prefetch_factor=2,
)
loader_valid = D.DataLoader(
valid,
batch_size=80,
shuffle=True,
num_workers=1,
pin_memory=True,
prefetch_factor=2,
)
ads3.run_experiment(
loader_train, loader_valid, log_name,
) # For profiling feel free to lower epoch count via epoch=X