-
Notifications
You must be signed in to change notification settings - Fork 256
/
data_gen.py
104 lines (80 loc) · 3.29 KB
/
data_gen.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
from __future__ import absolute_import
import sys
import numpy as np
import tensorflow as tf
from PIL import Image
from six import BytesIO as IO
from .bucketdata import BucketData
try:
TFRecordDataset = tf.data.TFRecordDataset # pylint: disable=invalid-name
except AttributeError:
TFRecordDataset = tf.contrib.data.TFRecordDataset # pylint: disable=invalid-name
class DataGen(object):
GO_ID = 1
EOS_ID = 2
IMAGE_HEIGHT = 32
CHARMAP = ['', '', ''] + list('0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ')
@staticmethod
def set_full_ascii_charmap():
DataGen.CHARMAP = ['', '', ''] + [chr(i) for i in range(32, 127)]
def __init__(self,
annotation_fn,
buckets,
epochs=1000,
max_width=None):
"""
:param annotation_fn:
:param lexicon_fn:
:param valid_target_len:
:param img_width_range: only needed for training set
:param word_len:
:param epochs:
:return:
"""
self.epochs = epochs
self.max_width = max_width
self.bucket_specs = buckets
self.bucket_data = BucketData()
dataset = TFRecordDataset([annotation_fn])
dataset = dataset.map(self._parse_record)
dataset = dataset.shuffle(buffer_size=10000)
self.dataset = dataset.repeat(self.epochs)
def clear(self):
self.bucket_data = BucketData()
def gen(self, batch_size):
dataset = self.dataset.batch(batch_size)
iterator = tf.compat.v1.data.make_one_shot_iterator(dataset)
images, labels, comments = iterator.get_next()
with tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True)) as sess:
while True:
try:
raw_images, raw_labels, raw_comments = sess.run([images, labels, comments])
for img, lex, comment in zip(raw_images, raw_labels, raw_comments):
if self.max_width and (Image.open(IO(img)).size[0] <= self.max_width):
word = self.convert_lex(lex)
bucket_size = self.bucket_data.append(img, word, lex, comment)
if bucket_size >= batch_size:
bucket = self.bucket_data.flush_out(
self.bucket_specs,
go_shift=1)
yield bucket
except tf.errors.OutOfRangeError:
break
self.clear()
def convert_lex(self, lex):
if sys.version_info >= (3,):
lex = lex.decode('iso-8859-1')
assert len(lex) < self.bucket_specs[-1][1]
return np.array(
[self.GO_ID] + [self.CHARMAP.index(char) for char in lex] + [self.EOS_ID],
dtype=np.int32)
@staticmethod
def _parse_record(example_proto):
features = tf.io.parse_single_example(
example_proto,
features={
'image': tf.io.FixedLenFeature([], tf.string),
'label': tf.io.FixedLenFeature([], tf.string),
'comment': tf.io.FixedLenFeature([], tf.string, default_value=''),
})
return features['image'], features['label'], features['comment']