afc_contract_ocr.py
1.39 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
# -*- coding: utf-8 -*-
# @Author : lk
# @Email : 9428.al@gmail.com
# @Created Date : 2021-06-29 17:43:46
# @Last Modified : 2021-09-07 14:11:25
# @Description :
from .get_char import Finder
import numpy as np
def predict(pdf_info):
ocr_results = {}
for pno in pdf_info:
ocr_results[pno] = {}
ocr_result = []
for key, block in enumerate(pdf_info[pno]['blocks']):
if block['type'] != 0:
continue
for line in block['lines']:
for span in line['spans']:
bbox, text = span['bbox'], span['text']
if len(text) == 0:
continue
# print(text)
xmin, ymin, xmax, ymax = bbox
polygon = [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax]
polygon = np.array(polygon, dtype=np.int32).tolist()
text = text.replace("：", ":").replace(" ", "")
ocr_result.append([polygon, text])
ocr_result = sorted(ocr_result, key=lambda x: x[0][1], reverse=False) # 按 y0 从小到大排
keys = list(range(len(ocr_result)))
ocr_result = dict(zip(keys, ocr_result))
ocr_results[pno] = ocr_result
# 输入是整个 PDF 中的信息
f = Finder(pdf_info, ocr_results=ocr_results)
results = f.get_info()
return results