使用Python和OCR进行文档解析的完整代码演示( 二 ) _OCR

文章插图

我将要使用的模型只能检测4个对象(文本、标题、列表、表格、图形) 。因此，如果你需要识别其他东西(如方程)，你就必须使用其他模型。
## load pre-trained model model = lp.Detectron2LayoutModel( "lp://PubLa.NET/mask_rcnn_X_101_32x8d_FPN_3x/config", extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8], label_map={0:"Text", 1:"Title", 2:"List", 3:"Table", 4:"Figure"}) ## turn img into array i = 21 img = np.asarray(doc[i]) ## predict detected = model.detect(img) ## plot lp.draw_box(img, detected, box_width=5, box_alpha=0.2, show_element_type=True)

文章插图

结果包含每个检测到的布局的细节，例如边界框的坐标。根据页面上显示的顺序对输出进行排序是很有用的：

## sort new_detected = detected.sort(key=lambda x: x.coordinates[1]) ## assign ids detected = lp.Layout([block.set(id=idx) for idx,block in enumerate(new_detected)])## check for block in detected: print("---", str(block.id)+":", block.type, "---") print(block, end='nn')

文章插图

完成OCR的下一步是正确提取检测到内容中的有用信息。
提取我们已经对图像完成了分割，然后就需要使用另外一个模型处理分段的图像，并将提取的输出保存到字典中。
由于有不同类型的输出（文本，标题，图形，表格），所以这里准备了一个函数用来显示结果。

''' {'0-Title': '...', '1-Text': '...', '2-Figure': array([[ [0,0,0], ...]]), '3-Table': pd.DataFrame, } ''' def parse_doc(dic): for k,v in dic.items(): if "Title" in k: print('x1b[1;31m'+ v +'x1b[0m') elif "Figure" in k: plt.figure(figsize=(10,5)) plt.imshow(v) plt.show() else: print(v) print(" ")

首先看看文字：

# load model model = lp.TesseractAgent(languages='eng') dic_predicted = {} for block in [block for block in detected if block.type in ["Title","Text"]]: ## segmentation segmented = block.pad(left=15, right=15, top=5, bottom=5).crop_image(img) ## extraction extracted = model.detect(segmented) ## save dic_predicted[str(block.id)+"-"+block.type] = extracted.replace('n',' ').strip() # check parse_doc(dic_predicted)

文章插图

再看看图形报表

for block in [block for block in detected if block.type == "Figure"]: ## segmentation segmented = block.pad(left=15, right=15, top=5, bottom=5).crop_image(img) ## save dic_predicted[str(block.id)+"-"+block.type] = segmented # check parse_doc(dic_predicted)

文章插图

上面两个看着很不错，那是因为这两种类型相对简单，但是表格就要复杂得多。尤其是我们上看看到的的这个，因为它的行和列都是进行了合并后产生的。

for block in [block for block in detected if block.type == "Table"]: ## segmentation segmented = block.pad(left=15, right=15, top=5, bottom=5).crop_image(img) ## extraction extracted = model.detect(segmented) ## save dic_predicted[str(block.id)+"-"+block.type] = pd.read_csv( io.StringIO(extracted) ) # check parse_doc(dic_predicted)

文章插图

正如我们的预料提取的表格不是很好。好在Python有专门处理表格的包，我们可以直接处理而不将其转换为图像。这里使用TabulaPy 包：
import tabula tables = tabula.read_pdf("doc_apple.pdf", pages=i+1) tables[0]