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