def run(input_dir, output_dir, model_dir):
images_dir = input_dir
images = get_images(images_dir)
predict_fn = predictor.from_saved_model(model_dir)
for image_file_path in images:
im, img_resized, ratio_h, ratio_w = read_image(image_file_path)
result = predict_fn({'images': img_resized})
get_text_segmentation_pb(img_mat=im,
result=result,
output_dir=output_dir,
file_name=os.path.basename(image_file_path),
ratio_h=ratio_h,
ratio_w=ratio_w)
def east_flow_predictions(input_dir=demo_config.IMAGE_ROOT_DIR,
output_dir=demo_config.EAST_OUT_DIR,
model_dir=demo_config.EAST_MODEL_DIR):
print(">>>>>>>>>>>>>>>>>>>>>", input_dir)
images_dir = input_dir
images = get_images(images_dir)
predict_fn = tf.contrib.predictor.from_saved_model(model_dir)
for image_file_path in images:
im, img_resized, ratio_h, ratio_w = read_image(image_file_path)
result = predict_fn({'images': img_resized})
get_text_segmentation_pb(img_mat=im,
result=result,
output_dir=output_dir,
file_name=os.path.basename(image_file_path),
ratio_h=ratio_h,
ratio_w=ratio_w)
def run_on_images(host, port, images_dir, model, signature_name, output_dir):
images = get_images(images_dir)
for image_file_path in images:
im, img_resized, ratio_h, ratio_w = read_image(image_file_path)
result = get_predictions(host=host,
port=port,
model_name=model,
signature_name=signature_name,
input_placeholder_name="images",
mat=img_resized)
get_text_segmentation(img_mat=im,
result=result,
output_dir=output_dir,
file_name=os.path.basename(image_file_path),
ratio_h=ratio_h,
ratio_w=ratio_w)
def run(save_checkpoints_steps=10,
keep_checkpoint_max=5,
save_summary_steps=5,
log_step_count_steps=5,
num_epochs=5,
test_iterator=False,
test_images_dir="",
output_dir=gin.REQUIRED):
"""
"""
model = EASTModel()
data_iterator = CIDARIterator()
run_config = tf.ConfigProto()
run_config.gpu_options.allow_growth = True
# run_config.gpu_options.per_process_gpu_memory_fraction = 0.50
run_config.allow_soft_placement = True
run_config.log_device_placement = False
model_dir = model.model_dir
run_config = tf.estimator.RunConfig(
session_config=run_config,
save_checkpoints_steps=save_checkpoints_steps,
keep_checkpoint_max=keep_checkpoint_max,
save_summary_steps=save_summary_steps,
model_dir=model_dir,
log_step_count_steps=log_step_count_steps)
executor = Executor(model=model,
data_iterator=data_iterator,
config=run_config,
train_hooks=None,
eval_hooks=None,
session_config=None)
if test_iterator:
executor.test_iterator()
num_samples = data_iterator._num_train_examples
batch_size = data_iterator._batch_size
if not FLAGS.predict:
for current_epoch in tqdm(range(num_epochs), desc="Epoch"):
current_max_steps = (num_samples // batch_size) * (current_epoch +
1)
print("\n\n Training for epoch {} with steps {}\n\n".format(
current_epoch, current_max_steps))
# executor.train(max_steps=None)
print("\n\n Evaluating for epoch\n\n", current_epoch)
# executor.evaluate(steps=None)
executor.train_and_evaluate()
executor.export_model(model_dir + "/exported/")
else:
estimator = executor._estimator
images = get_images(test_images_dir)
for image_file_path in images:
print("================> Text segmentation on :", image_file_path)
im = cv2.imread(image_file_path)[:, :, ::-1]
start_time = time.time()
im_resized, (ratio_h, ratio_w) = resize_image(im)
im_resized = np.expand_dims(im_resized, axis=0).astype(np.float32)
def get_dataset():
dataset = tf.data.Dataset.from_tensor_slices(({
"images":
im_resized
}, np.ones_like(im_resized)))
dataset = dataset.batch(batch_size=1)
print(dataset.output_shapes)
return dataset
start = time.time()
timer = {'net': 0, 'restore': 0, 'nms': 0}
predict_fn = estimator.predict(input_fn=lambda: get_dataset())
for prediction in predict_fn:
score = prediction["f_score"]
geometry = prediction["f_geometry"]
score = np.expand_dims(score, axis=0)
geometry = np.expand_dims(geometry, axis=0)
print("===============================")
print(score.shape)
print(geometry.shape)
print("===============================")
print(score)
print(geometry)
timer['net'] = time.time() - start
boxes, timer = detect(score_map=score,
geo_map=geometry,
timer=timer)
print('{} : net {:.0f}ms, restore {:.0f}ms, nms {:.0f}ms'.format(
image_file_path, timer['net'] * 1000, timer['restore'] * 1000,
timer['nms'] * 1000))
if boxes is not None:
boxes = boxes[:, :8].reshape((-1, 4, 2))
boxes[:, :, 0] /= ratio_w
boxes[:, :, 1] /= ratio_h
duration = time.time() - start_time
print('[timing] {}'.format(duration))
# save to file
if boxes is not None:
res_file = os.path.join(
output_dir, '{}.txt'.format(
os.path.basename(image_file_path).split('.')[0]))
with open(res_file, 'w') as f:
for box in boxes:
# to avoid submitting errors
box = sort_poly(box.astype(np.int32))
if np.linalg.norm(box[0] -
box[1]) < 5 or np.linalg.norm(
box[3] - box[0]) < 5:
continue
f.write('{},{},{},{},{},{},{},{}\r\n'.format(
box[0, 0],
box[0, 1],
box[1, 0],
box[1, 1],
box[2, 0],
box[2, 1],
box[3, 0],
box[3, 1],
))
cv2.polylines(
im[:, :, ::-1],
[box.astype(np.int32).reshape((-1, 1, 2))],
True,
color=(255, 255, 0),
thickness=1)
# if not FLAGS.no_write_images:
img_path = os.path.join(output_dir,
os.path.basename(image_file_path))
cv2.imwrite(img_path, im[:, :, ::-1])