def segment_and_pred(source_path, print_path, img_type): print('Start process pic:' + source_path) splitter = Splitter() if 'school' in img_type: segment_flag = school_flag attr = school_attr else: segment_flag = degree_flag attr = degree_attr result = {} image_color = cv2.imread(source_path) image = cv2.cvtColor(image_color, cv2.COLOR_BGR2GRAY) ret, adaptive_threshold = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY) splitter.show_img('adaptive_threshold', adaptive_threshold) ret, at = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY_INV) splitter.show_img('at', at) first_column_img = adaptive_threshold[0:image_color.shape[1], 120:350] second_column_img = adaptive_threshold[0:image_color.shape[1], 400:659] first_column_img_w = first_column_img.shape[0] second_column_img_w = second_column_img.shape[0] # 计算换行内容索引 horizontal_sum = np.sum(at, axis=1) peek_ranges = splitter.extract_peek_ranges_from_array(horizontal_sum) last_pr = None # 内容换行 line_feed_index_dict = {} # 内容为空 line_empty_index_dict = {} j = 0 k = 0 for i, pr in enumerate(peek_ranges): if last_pr is not None and pr[0] - last_pr[1] < 15: line_feed_index_dict[i] = j j += 1 elif last_pr is not None and pr[0] - last_pr[1] >= 42: line_empty_index_dict[i] = k k += 1 last_pr = pr # 内容包括换行的行数 line_feed_count = 0 # 内容为空的行数 line_empty_count = 0 i = 0 while i < (len(segment_flag) + len(line_feed_index_dict)): if i in line_feed_index_dict: line_feed_count += 1 elif i in line_empty_index_dict: line_empty_count += 1 result[attr[i - line_feed_count][0]] = '' result[attr[i - line_feed_count][1]] = '' i += 1 continue tmp1 = first_column_img[ peek_ranges[i - line_empty_count][0]:peek_ranges[i - line_empty_count][1], 0:first_column_img_w - 1] splitter.show_img('first image', tmp1) kv0_path = print_path + str(i) + '/' if not os.path.exists(kv0_path): os.makedirs(kv0_path) cv2.imwrite(kv0_path + 'kv0.png', tmp1) # cv2.waitKey(0) if segment_flag[i - line_feed_count][0] == 1: min_width = 12 else: min_width = 3 kv0_print_path = print_path + str(i) + '/0/' if os.path.exists(kv0_print_path): shutil.rmtree(kv0_print_path) if not os.path.exists(kv0_print_path): os.makedirs(kv0_print_path) splitter.process_by_path(kv0_path + 'kv0.png', kv0_print_path, minimun_range=min_width) attr_name = attr[i - line_feed_count][0] if attr_name != '': pred_result, pred_val_list = chinese_ocr.pred(kv0_print_path) if resegment(pred_val_list): print('re segment:' + kv0_print_path) shutil.rmtree(kv0_print_path) os.makedirs(kv0_print_path) splitter.process_by_path(kv0_path + 'kv0.png', kv0_print_path, minimun_range=min_width, pred_val_list=pred_val_list) pred_result, pred_val_list = chinese_ocr.pred(kv0_print_path) if attr_name in result: # 内容换行 result[attr_name] = result[attr_name] + pred_result else: result[attr_name] = pred_result tmp2 = second_column_img[ peek_ranges[i - line_empty_count][0]:peek_ranges[i - line_empty_count][1], 0:second_column_img_w - 1] splitter.show_img('second image', tmp2) kv1_path = print_path + str(i) + '/' cv2.imwrite(kv1_path + 'kv1.png', tmp2) if segment_flag[i - line_feed_count][1] == 1: min_width = 12 else: min_width = 3 kv1_print_path = print_path + str(i) + '/1/' if os.path.exists(kv1_print_path): shutil.rmtree(kv1_print_path) if not os.path.exists(kv1_print_path): os.makedirs(kv1_print_path) splitter.process_by_path(kv1_path + 'kv1.png', kv1_print_path, minimun_range=min_width) attr_name = attr[i - line_feed_count][1] if attr_name != '': pred_result, pred_val_list = chinese_ocr.pred(kv1_print_path) if resegment(pred_val_list): print('re segment:' + kv1_print_path) shutil.rmtree(kv1_print_path) os.makedirs(kv1_print_path) splitter.process_by_path(kv1_path + 'kv1.png', kv1_print_path, minimun_range=min_width, pred_val_list=pred_val_list) pred_result, pred_val_list = chinese_ocr.pred(kv1_print_path) if attr_name in result: # 内容换行 result[attr_name] = result[attr_name] + pred_result else: result[attr_name] = pred_result i += 1 return result