def predict(self, sess, img_file, out_dir='./results_prediction'):
#
# input data
# img_data, feat_size, target_cls, target_ver, target_hor = \
# model_detect_data.getImageAndTargets(img_file, meta.anchor_heights)
#
img = Image.open(img_file)
img_data = np.array(img, dtype=np.float32) / 255
# height, width, channel
#
img_data = [img_data[:, :, 0:3]] # rgba
img_size = model_detect_data.getImageSize(img_file) # width, height
height_feat = int(floor(ceil(ceil(img_size[1] / 2.0) / 2.0) / 3.0) - 2)
#
w_arr = np.ones((height_feat, ), dtype=np.int32) * img_size[0]
#
# predication_result save-path
if not os.path.exists(out_dir): os.mkdir(out_dir)
#
with self.graph.as_default():
#
feed_dict = {self.x: img_data, self.w: w_arr}
#
r_cls, r_ver, r_hor = sess.run(
[self.rnn_cls, self.rnn_ver, self.rnn_hor], feed_dict)
#
#
# filename = os.path.basename(img_file)
# arr_str = os.path.splitext(filename)
# #
# # image
# r = Image.fromarray(img_data[0][:,:,0] *255).convert('L')
# g = Image.fromarray(img_data[0][:,:,1] *255).convert('L')
# b = Image.fromarray(img_data[0][:,:,2] *255).convert('L')
# #
# file_target = os.path.join(out_dir, arr_str[0] + '_predict.png')
# img_target = Image.merge("RGB", (r, g, b))
# img_target.save(file_target)
#
# trans
text_bbox,text_score = model_detect_data.transResults(r_cls, r_ver, r_hor, \
meta.anchor_heights, meta.threshold)
#
#model_detect_data.drawTextBox(file_target, text_bbox)
#
return text_bbox, text_score
def predict(self, sess, img_file, out_dir='./results_prediction'):
#
# input data
img_data, feat_size, target_cls, target_ver, target_hor = \
model_detect_data.getImageAndTargets(img_file, meta.anchor_heights)
#
img_size = model_detect_data.getImageSize(img_file) # width, height
#
w_arr = np.ones((feat_size[0], ), dtype=np.int32) * img_size[0]
#
# predication_result save-path
if not os.path.exists(out_dir): os.mkdir(out_dir)
#
with self.graph.as_default():
#
feed_dict = {self.x: img_data, self.w: w_arr}
#
r_cls, r_ver, r_hor = sess.run(
[self.rnn_cls, self.rnn_ver, self.rnn_hor], feed_dict)
#
#
filename = os.path.basename(img_file)
arr_str = os.path.splitext(filename)
#
# image
r = Image.fromarray(img_data[0][:, :, 0] * 255).convert('L')
g = Image.fromarray(img_data[0][:, :, 1] * 255).convert('L')
b = Image.fromarray(img_data[0][:, :, 2] * 255).convert('L')
#
file_target = os.path.join(out_dir, arr_str[0] + '_predict.png')
img_target = Image.merge("RGB", (r, g, b))
img_target.save(file_target)
#
# trans
text_bbox = model_detect_data.transResults(r_cls, r_ver, r_hor, \
meta.anchor_heights, meta.threshold)
#
model_detect_data.drawTextBox(file_target, text_bbox)
def train_and_valid(self):
#
# get training images
list_images_train = model_detect_data.getFilesInDirect(meta.dir_images_train, meta.str_dot_img_ext)
#
# model save-path
if not os.path.exists(meta.model_detect_dir): os.mkdir(meta.model_detect_dir)
#
# training graph
self.z_graph = tf.Graph()
#
self.z_define_graph_all(self.z_graph, True)
#
with self.z_graph.as_default():
#
saver = tf.train.Saver()
#
# gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = 0.95, allocator_type = 'BFC')
# sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
#
with tf.Session(config = self.z_sess_config) as sess:
#
tf.global_variables_initializer().run()
#
# restore with saved data
ckpt = tf.train.get_checkpoint_state(meta.model_detect_dir)
#
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
#
#
# variables
#
x = self.z_graph.get_tensor_by_name('x-input:0')
w = self.z_graph.get_tensor_by_name('w-input:0')
#
t_cls = self.z_graph.get_tensor_by_name('c-input:0')
t_ver = self.z_graph.get_tensor_by_name('v-input:0')
t_hor = self.z_graph.get_tensor_by_name('h-input:0')
#
loss = self.z_graph.get_tensor_by_name('loss:0')
#
global_step = self.z_graph.get_tensor_by_name('global_step:0')
learning_rate = self.z_graph.get_tensor_by_name('learning_rate:0')
train_op = self.z_graph.get_tensor_by_name('train_op/control_dependency:0')
#
#
print('begin to train ...')
#
# start training
start_time = time.time()
begin_time = start_time
#
for i in range(TRAINING_STEPS):
#
print(list_images_train)
img_file = random.choice(list_images_train)
#
# print(img_file)
#
# input data
print(meta.anchor_heights)
img_data, feat_size, target_cls, target_ver, target_hor = \
model_detect_data.getImageAndTargets(img_file, meta.anchor_heights)
#
img_size = model_detect_data.getImageSize(img_file) # width, height
#
print(feat_size[0])
w_arr = np.ones((int(feat_size[0]),), dtype = np.int32) * img_size[0]
#
#
feed_dict = {x: img_data, w: w_arr, \
t_cls: target_cls, t_ver: target_ver, t_hor: target_hor}
#
#print(img_data.size)
#print(feat_size)
#print(w_arr)
#
#rnn_cls_v = sess.run(seq_len, feed_dict)
#print(len(rnn_cls_v))
#
#loss_value = sess.run(loss, feed_dict)
#print('sess.run, loss = %g' % loss_value)
#
_, loss_value, step, lr = sess.run([train_op, loss, global_step, learning_rate], \
feed_dict)
#
if i % 1 == 0:
#
curr_time = time.time()
#
print('step: %d, loss: %g, lr: %g, sect_time: %.1f, total_time: %.1f, %s' %
(step, loss_value, lr,
curr_time - begin_time,
curr_time - start_time,
os.path.basename(img_file)))
#
begin_time = curr_time
#
#
# validation
if step % self.z_valid_freq == 0:
#
# ckpt
saver.save(sess, os.path.join(meta.model_detect_dir, meta.model_detect_name), \
global_step = step)
#
self.validate(step, self.z_valid_option)
def validate(self, step, training):
#
# get validation images
list_images_valid = model_detect_data.getFilesInDirect(meta.dir_images_valid, meta.str_dot_img_ext)
#
# valid_result save-path
if not os.path.exists(meta.dir_results_valid): os.mkdir(meta.dir_results_valid)
#
# if os.path.exists(dir_results): shutil.rmtree(dir_results)
# time.sleep(0.1)
# os.mkdir(dir_results)
#
# validation graph
self.graph = tf.Graph()
#
self.z_define_graph_all(self.graph, training)
#
with self.graph.as_default():
#
saver = tf.train.Saver()
#
# gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = 0.95)
# sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
#
with tf.Session(config = self.z_sess_config) as sess:
#
tf.global_variables_initializer().run()
#
# restore with saved data
ckpt = tf.train.get_checkpoint_state(meta.model_detect_dir)
#
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
#
# pb
constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names = \
['rnn_cls','rnn_ver','rnn_hor'])
with tf.gfile.FastGFile(self.z_pb_file, mode='wb') as f:
f.write(constant_graph.SerializeToString())
#
# variables
#
x = self.graph.get_tensor_by_name('x-input:0')
w = self.graph.get_tensor_by_name('w-input:0')
#
rnn_cls = self.graph.get_tensor_by_name('rnn_cls:0')
rnn_ver = self.graph.get_tensor_by_name('rnn_ver:0')
rnn_hor = self.graph.get_tensor_by_name('rnn_hor:0')
#
t_cls = self.graph.get_tensor_by_name('c-input:0')
t_ver = self.graph.get_tensor_by_name('v-input:0')
t_hor = self.graph.get_tensor_by_name('h-input:0')
#
loss = self.graph.get_tensor_by_name('loss:0')
#
# test
NumImages = len(list_images_valid)
curr = 0
for img_file in list_images_valid:
#
# input data
img_data, feat_size, target_cls, target_ver, target_hor = \
model_detect_data.getImageAndTargets(img_file, meta.anchor_heights)
#
img_size = model_detect_data.getImageSize(img_file) # width, height
#
w_arr = np.ones((int(feat_size[0]),), dtype = np.int32) * img_size[0]
#
feed_dict = {x: img_data, w: w_arr, \
t_cls: target_cls, t_ver: target_ver, t_hor: target_hor}
#
r_cls, r_ver, r_hor, loss_value = sess.run([rnn_cls, rnn_ver, rnn_hor, loss], feed_dict)
#
#
curr += 1
print('curr: %d / %d, loss: %f' % (curr, NumImages, loss_value))
#
filename = os.path.basename(img_file)
arr_str = os.path.splitext(filename)
#
# image
r = Image.fromarray(img_data[0][:,:,0] *255).convert('L')
g = Image.fromarray(img_data[0][:,:,1] *255).convert('L')
b = Image.fromarray(img_data[0][:,:,2] *255).convert('L')
#
file_target = os.path.join(meta.dir_results_valid, str(step) + '_' +arr_str[0] + '.png')
img_target = Image.merge("RGB", (r, g, b))
img_target.save(file_target)
#
# trans
text_bbox = model_detect_data.transResults(r_cls, r_ver, r_hor, \
meta.anchor_heights, meta.threshold)
#
model_detect_data.drawTextBox(file_target, text_bbox)
#
#
print('validation finished')
#
# start training
start_time = time.time()
begin_time = start_time
#
for i in range(TRAINING_STEPS):
#
img_file = random.choice(list_images_train)
#
#print(img_file)
#
# input data
img_data, feat_size, target_cls, target_ver, target_hor = \
model_detect_data.getImageAndTargets(img_file, anchor_heights)
#
img_size = model_detect_data.getImageSize(img_file) # width, height
#
w_arr = np.ones((feat_size[0], ), dtype=np.int32) * img_size[0]
#
#
feed_dict = {x: img_data, w: w_arr, \
t_cls: target_cls, t_ver: target_ver, t_hor: target_hor}
#
_, loss_value, step, lr = sess.run(
[train_op, loss, global_step, learning_rate], feed_dict)
#
if i % 1 == 0:
#
saver.save(sess,
os.path.join(model_dir, model_name),
global_step=step)