-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdatasets.py
129 lines (112 loc) · 5.83 KB
/
datasets.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
import os
from typing import Optional
import pandas as pd
import pytorch_lightning as pl
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from torchvision.transforms import Compose, Resize, ToTensor, Grayscale, RandomRotation, RandomApply, \
GaussianBlur, CenterCrop
class KaggleHandwrittenNames(Dataset):
def __init__(self, data, transforms, label_to_index, img_path):
self.data = data
self.transforms = transforms
self.img_path = img_path
self.label_to_index = label_to_index
def __len__(self):
return self.data.shape[0]
def __getitem__(self, index):
row = self.data.iloc[index]
file_name = row['FILENAME']
image_label = row['IDENTITY']
the_image = Image.open(os.path.join(self.img_path, file_name))
transformed_image = self.transforms(the_image)
target_len = len(image_label)
label_chars = list(image_label)
image_label = torch.tensor([self.label_to_index[char] for char in label_chars])
return {
'transformed_image': transformed_image,
'label': image_label,
'target_len': target_len
}
class KaggleHandwritingDataModule(pl.LightningDataModule):
def __init__(self, train_data, val_data, hparams, label_to_index):
super().__init__()
self.train_data = train_data
self.val_data = val_data
self.train_batch_size = hparams['train_batch_size']
self.val_batch_size = hparams['val_batch_size']
self.transforms = Compose([Resize((hparams['input_height'], hparams['input_width'])), Grayscale(),
ToTensor()])
applier1 = RandomApply(transforms=[RandomRotation(10), GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5))], p=0.5)
applier2 = RandomApply(transforms=[CenterCrop((hparams['input_height'] - 1, hparams['input_width'] - 2))], p=0.5)
self.train_transforms = Compose([applier2, Resize((hparams['input_height'], hparams['input_width'])), Grayscale(),
applier1, ToTensor()])
self.train_img_path = hparams['train_img_path']
self.val_img_path = hparams['val_img_path']
self.label_to_index = label_to_index
def setup(self, stage: Optional[str] = None):
if stage == 'fit' or stage is None:
self.train = KaggleHandwrittenNames(self.train_data, self.train_transforms, self.label_to_index, self.train_img_path)
self.val = KaggleHandwrittenNames(self.val_data, self.transforms, self.label_to_index, self.val_img_path)
def custom_collate(data):
'''
To handle variable max seq length batch size
'''
transformed_images = []
labels = []
target_lens = []
for d in data:
transformed_images.append(d['transformed_image'])
labels.append(d['label'])
target_lens.append(d['target_len'])
# pad_sequence pads all batch items to the longest length sequence in the batch because labels
# can be of variable length
batch_labels = pad_sequence(labels, batch_first=True, padding_value=-1)
transformed_images = torch.stack(transformed_images)
target_lens = torch.tensor(target_lens)
return {
'transformed_images': transformed_images,
'labels': batch_labels,
'target_lens': target_lens
}
def train_dataloader(self):
return DataLoader(self.train, batch_size=self.train_batch_size, shuffle=True, pin_memory=True,
num_workers=8, collate_fn=KaggleHandwritingDataModule.custom_collate)
def val_dataloader(self):
return DataLoader(self.val, batch_size=self.val_batch_size, shuffle=False, pin_memory=True,
num_workers=8, collate_fn=KaggleHandwritingDataModule.custom_collate)
def test_kaggle_handwritting():
pl.seed_everything(267)
hparams = {
'train_img_path': './data/kaggle-handwriting-recognition/train_v2/train/',
'lr': 1e-3, 'val_img_path': './data/kaggle-handwriting-recognition/validation_v2/validation/',
'test_img_path': './data/kaggle-handwriting-recognition/test_v2/test/',
'data_path': './data/kaggle-handwriting-recognition', 'gru_input_size': 256,
'train_batch_size': 64, 'val_batch_size': 256, 'input_height': 36, 'input_width': 324, 'gru_hidden_size': 128,
'gru_num_layers': 1, 'num_classes': 28
}
label_to_index = {' ': 0, '-': 1, 'A': 2, 'B': 3, 'C': 4, 'D': 5, 'E': 6, 'F': 7, 'G': 8, 'H': 9, 'I': 10, 'J': 11,
'K': 12, 'L': 13, 'M': 14, 'N': 15, 'O': 16, 'P': 17, 'Q': 18, 'R': 19, 'S': 20, 'T': 21, 'U': 22,
'V': 23, 'W': 24, 'X': 25, 'Y': 26, 'Z': 27}
train_df = pd.read_csv(os.path.join(hparams['data_path'], 'train_new.csv'))
train_df = train_df[train_df.word_type == 'normal_word']
train_df = train_df.sample(frac=1).reset_index(drop=True)
val_df = pd.read_csv(os.path.join(hparams['data_path'], 'val_new.csv'))
val_df = val_df[val_df.word_type == 'normal_word']
val_df = val_df.sample(frac=1).reset_index(drop=True)
sample_module = KaggleHandwritingDataModule(train_df, val_df, hparams, label_to_index)
sample_module.setup()
sample_train_module = sample_module.train_dataloader()
sample_val_module = sample_module.val_dataloader()
sample_train_batch = next(iter(sample_train_module))
sample_val_batch = next(iter(sample_val_module))
print(sample_train_batch['transformed_images'].shape)
print(sample_val_batch['transformed_images'].shape)
print(sample_train_batch['labels'].shape)
print(sample_val_batch['labels'].shape)
print(sample_train_batch['target_lens'].shape)
print(sample_val_batch['target_lens'].shape)
if __name__ == '__main__':
test_kaggle_handwritting()