File size: 3,272 Bytes
ba3d672 |
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 |
import torchvision
import torch
from torch.utils.data import Dataset
from torch.nn.utils.rnn import pad_sequence
class AddGaussianNoise(object):
def __init__(self, mean=0., std=1., thresh=0.2):
self.mean = mean
self.std = std
self.thresh = thresh
def __call__(self, tensor):
noise = torch.zeros_like(tensor)
noise[tensor>self.thresh] = 1
noise *= torch.randn(tensor.size()) * self.std + self.mean
return tensor + noise
def __repr__(self):
return self.__class__.__name__ + f'(mean={self.mean}, std={self.std})'
class TextProcessor:
def __init__(self, alphabet):
self.alphabet = alphabet
self.pad_token = "[PAD]"
self.stoi = {s: i for i, s in enumerate(self.alphabet,1)}
self.stoi[self.pad_token] = 0
self.itos = {i: s for s, i in self.stoi.items()}
def encode(self, label):
return [self.stoi[s] for s in label]
def decode(self, ids):
return ''.join([self.itos[i] for i in ids])
def __len__(self):
return len(self.alphabet) + 1
transform_train = torchvision.transforms.Compose(
[
torchvision.transforms.Grayscale(),
torchvision.transforms.ToTensor(),
torchvision.transforms.RandomApply([
torchvision.transforms.RandomAdjustSharpness(sharpness_factor=80),
AddGaussianNoise(mean=1, std=0.005, thresh=0.3),
])
]
)
transform_eval = torchvision.transforms.Compose(
[
torchvision.transforms.Grayscale(),
torchvision.transforms.ToTensor()
]
)
class CRNNDataset(Dataset):
def __init__(
self,
height,
text_processor:TextProcessor,
transforms:torchvision.transforms,
dataset=None
) -> None:
super().__init__()
self.height = height
self.transform = transforms
self.dataset = dataset
self.text_processor = text_processor
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
dset = self.dataset[idx]
image, text = dset['image'], dset['text']
label = torch.tensor(self.text_processor.encode(text), dtype=torch.long)
original_width, original_height = image.size
new_width = int(self.height * original_width / original_height) # Calculate width to preserve aspect ratio
image = image.resize((new_width, self.height))
image = self.transform(image)
return image, label
def collate_fn(batch):
images, labels = zip(*batch)
max_h = max(img.size(1) for img in images)
max_w = max(img.size(2) for img in images)
padded_images = []
for img in images:
h, w = img.size(1), img.size(2)
padding = (0, max_w - w, 0, max_h - h) # left, right, top, bottom
padded_img = torch.nn.functional.pad(img, padding, mode='constant', value=0)
padded_images.append(padded_img)
images = torch.stack(padded_images, 0)
target_lengths = torch.tensor([len(label) for label in labels]).long()
labels = pad_sequence(labels, batch_first=True, padding_value=0)
return images, labels, target_lengths |