def check_tags(tags_file, segment_task, text_type):
with open(tags_file, "r", encoding="utf8") as fr:
lines = [line.strip() for line in fr.readlines()]
save_path = os.path.join(SEGMENT_BOOK_PAGE_ROOT_DIR, "samples")
check_or_makedirs(save_path)
for i, line in enumerate(lines):
np_img, split_pos = get_image_and_split_pos(line, segment_task="book_page")
text_type = text_type[0].lower()
if (segment_task, text_type) in (("book_page", "h"), ("double_line", "h"), ("text_line", "v"), ("mix_line", "v")):
np_img, split_pos = rotate_90_degrees(np_img, split_pos)
np_img = draw_split_lines(np_img, split_pos)
PIL_img = Image.fromarray(np_img)
PIL_img.save(os.path.join(save_path, str(i) + ".jpg"))
def segment_predict(images=None,
img_paths=None,
dest_dir=None,
segment_model=None,
segment_task="book_page",
text_type="horizontal",
model_struc="densenet_gru",
weights=""):
# images
if images is not None:
np_img_list = convert_images(images)
img_name_list = [str(i)+".jpg" for i in range(len(np_img_list))]
else:
assert img_paths is not None
np_img_list, img_name_list = load_images(img_paths)
# book page pre-processing
if segment_task == "book_page":
np_img_list = book_page_pre_processing(np_img_list)
# model
if segment_model is None:
K.set_learning_phase(False)
weights_path = model_weights_path(weights, segment_task, model_struc)
# 加载模型
segment_model = work_net(stage="predict", segment_task=segment_task, text_type=text_type, model_struc=model_struc)
segment_model.load_weights(weights_path, by_name=True)
print("\nLoad model weights from %s\n" % weights_path)
# segment_model.summary()
# predict
batch_size, fixed_h, feat_stride = get_segment_task_params(segment_task)
text_type = text_type[0].lower()
split_positions_list, scores_list = [], []
for i in range(0, len(np_img_list), batch_size):
_images_list, _scale_ratio_list = [], []
for np_img in np_img_list[i:i+batch_size]:
np_img, _, scale_ratio = adjust_img_to_fixed_height(np_img, None, fixed_h, segment_task, text_type)
_images_list.append(np_img)
_scale_ratio_list.append(scale_ratio)
batch_images, real_images_width, _ = pack_a_batch(_images_list, None, feat_stride, background="white")
nms_split_positions, nms_scores = segment_model.predict(x=[batch_images, real_images_width]) # 模型预测
for j in range(len(batch_images)):
scores = remove_pad_np(nms_scores[j])[:, 0]
split_positions = remove_pad_np(nms_split_positions[j])
split_positions = split_positions / _scale_ratio_list[j]
if (segment_task, text_type) in (("book_page", "h"), ("double_line", "h"), ("text_line", "v"), ("mix_line", "v")):
_, split_positions = restore_original_angle(np_img=None, pred_split_positions=split_positions)
split_positions_list.append(split_positions)
scores_list.append(scores)
# draw
if dest_dir is not None:
check_or_makedirs(dest_dir)
for i in range(len(np_img_list)):
if (segment_task, text_type) in (("book_page", "h"), ("double_line", "h"), ("text_line", "v"), ("mix_line", "v")):
np_img, split_positions = rotate_90_degrees(np_img_list[i], split_positions_list[i])
else:
np_img, split_positions = np_img_list[i], split_positions_list[i]
np_img = visualize.draw_split_lines(np_img, split_positions, scores_list[i]) # 可视化
if (segment_task, text_type) in (("book_page", "h"), ("double_line", "h"), ("text_line", "v"), ("mix_line", "v")):
np_img, _ = restore_original_angle(np_img)
PIL_img = Image.fromarray(np_img)
dest_path = os.path.join(dest_dir, os.path.splitext(img_name_list[i])[0] + ".jpg")
PIL_img.save(dest_path, format="jpeg")
print(i, "Finished: " + dest_path)
return np_img_list, img_name_list, split_positions_list, scores_list