Spaces:
Runtime error
Runtime error
File size: 5,832 Bytes
2daf4c3 |
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 |
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'])
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() |