| |
| """Classes for feeding data during training.""" |
| import numpy as np |
| import pandas as pd |
| from .helpers import img_extend |
| from .datahelpers import sequences_to_sparse |
|
|
|
|
| class BucketDataIterator(): |
| """Iterator for feeding CTC model during training.""" |
| def __init__(self, |
| images, |
| targets, |
| num_buckets=5, |
| slider=(60, 30), |
| augmentation=None, |
| dropout=0.0, |
| train=True): |
|
|
| self.train = train |
| self.slider = slider |
| self.augmentation = augmentation |
| self.dropout = dropout |
| for i in range(len(images)): |
| images[i] = img_extend( |
| images[i], |
| (self.slider[0], |
| max(images[i].shape[1], self.slider[1]))) |
| in_length = [image.shape[1] for image in images] |
| |
| |
| self.dataFrame = pd.DataFrame({ |
| 'in_length': in_length, |
| 'images': images, |
| 'targets': targets}).sort_values('in_length').reset_index(drop=True) |
|
|
| bsize = int(len(images) / num_buckets) |
| self.num_buckets = num_buckets |
| self.buckets = [] |
| for bucket in range(num_buckets-1): |
| self.buckets.append( |
| self.dataFrame.iloc[bucket * bsize: (bucket+1) * bsize]) |
| self.buckets.append(self.dataFrame.iloc[(num_buckets-1) * bsize:]) |
|
|
| self.buckets_size = [len(bucket) for bucket in self.buckets] |
| self.cursor = np.array([0] * num_buckets) |
| self.bucket_order = np.random.permutation(num_buckets) |
| self.bucket_cursor = 0 |
| self.shuffle() |
| print("Iterator created.") |
|
|
|
|
| def shuffle(self, idx=None): |
| """Shuffle idx bucket or each bucket separately.""" |
| for i in [idx] if idx is not None else range(self.num_buckets): |
| self.buckets[i] = self.buckets[i].sample(frac=1).reset_index(drop=True) |
| self.cursor[i] = 0 |
|
|
|
|
| def next_batch(self, batch_size): |
| """Creates next training batch of size. |
| Args: |
| batch_size: size of next batch |
| Retruns: |
| (images, labels, images lengths, labels lengths) |
| """ |
| i_bucket = self.bucket_order[self.bucket_cursor] |
| |
| self.bucket_cursor = (self.bucket_cursor + 1) % self.num_buckets |
| if self.bucket_cursor == 0: |
| self.bucket_order = np.random.permutation(self.num_buckets) |
|
|
| if self.cursor[i_bucket] + batch_size > self.buckets_size[i_bucket]: |
| self.shuffle(i_bucket) |
|
|
| |
| if (batch_size > self.buckets_size[i_bucket]): |
| batch_size = self.buckets_size[i_bucket] |
|
|
| res = self.buckets[i_bucket].iloc[self.cursor[i_bucket]: |
| self.cursor[i_bucket]+batch_size] |
| self.cursor[i_bucket] += batch_size |
|
|
| |
| input_max = max(res['in_length']) |
|
|
| input_imgs = np.zeros( |
| (batch_size, self.slider[0], input_max, 1), dtype=np.uint8) |
| for i, img in enumerate(res['images']): |
| input_imgs[i][:, :res['in_length'].values[i], 0] = img |
| |
| if self.train: |
| input_imgs = self.augmentation.augment_images(input_imgs) |
| input_imgs = input_imgs.astype(np.float32) |
|
|
| targets = sequences_to_sparse(res['targets'].values) |
| return input_imgs, targets, res['in_length'].values |
|
|
|
|
