retriever.py
11.1 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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
import re
class HMHRetriever:
def __init__(self):
self.words_str = 'words'
self.position_str = 'location'
self.fix_hava_str = '有'
self.default_position = [0, 0, 0, 0]
self.search_fields_list = [
('借款/承租人姓名', ''),
('证件号码', ''),
('渠道', ''),
('合同编号', ''),
('借款人签字/盖章', '无'),
]
def get_target_fields(self, pdf_text_list):
result = dict()
is_find_name_id_company, is_find_application_no, is_find_name_date = False, False, False
for bbox, text in pdf_text_list.pop(str(0), []):
# print(text)
if not is_find_name_id_company:
name_id_company_list = re.findall(r'姓名(.*)证件号码(.*)与(.*公司)', text)
for name_id_company_tuple in name_id_company_list:
if len(name_id_company_tuple) == 3:
result[self.search_fields_list[0][0]] = {
self.words_str: name_id_company_tuple[0].replace('\u3000', '').strip(),
self.position_str: bbox
}
result[self.search_fields_list[1][0]] = {
self.words_str: name_id_company_tuple[1].replace('\u3000', '').replace(')', '').replace('）', '').strip(),
self.position_str: bbox
}
result[self.search_fields_list[2][0]] = {
self.words_str: name_id_company_tuple[2],
self.position_str: bbox
}
is_find_name_id_company = True
break
if not is_find_application_no:
application_no_list = re.findall(r'合同编号.*(CH-B\d*-\d*).*', text)
if len(application_no_list) == 1:
result[self.search_fields_list[3][0]] = {
self.words_str: application_no_list[0],
self.position_str: bbox
}
is_find_application_no = True
if not is_find_name_date:
name_date_list = re.findall(r'(.*).*签署日期.*(\d{4}-\d{2}-\d{2})', text)
for name_date_tuple in name_date_list:
if len(name_date_tuple) == 2:
result[self.search_fields_list[4][0]] = {
# self.words_str: '{0} {1}'.format(name_date_tuple[0].replace('\u3000', '').strip(), name_date_tuple[1]),
self.words_str: self.fix_hava_str,
self.position_str: bbox
}
is_find_name_date = True
break
for find_key, default_value in self.search_fields_list:
if find_key not in result:
result[find_key] = {
self.words_str: default_value,
self.position_str: self.default_position,
}
# simple_result = []
# for key, value_dict in result.items():
# simple_result.append((key, value_dict[self.words_str]))
# return simple_result
return {"words_result": result}
class Retriever:
def __init__(self, target_fields):
self.keys_str = 'keys'
self.value_str = 'value'
self.text_str = 'text'
self.words_str = 'words'
self.position_str = 'position'
self.default_position = [-1, -1, -1, -1]
self.target_fields = target_fields
self.replace_map = {
'int': {
'(': '0'
}
}
@staticmethod
def key_top1(coordinates_list, key_coordinates):
# 关键词查找方向：最上面
coordinates_list.sort(key=lambda x: x[1])
return coordinates_list[0]
def key_right(self, coordinates_list, key_coordinates, offset_tuple, rigorous=False):
# 关键词查找方向：右侧
if len(coordinates_list) == 1:
return coordinates_list[0]
# 没有上一层关键词的坐标时，返回最上面的坐标
if key_coordinates is None:
return self.key_top1(coordinates_list, key_coordinates)
x_min, y_min, x_max, y_max = self.get_target_bbox(key_coordinates, offset_tuple)
x_min_find, find_key_coordinates = None, None
for x0, y0, x1, y1 in coordinates_list:
if rigorous:
is_eligible = x_min < x0 and x1 < x_max and y_min < y0 and y1 < y_max
else:
cent_x = x0 + ((x1 - x0) / 2)
cent_y = y0 + ((y1 - y0) / 2)
is_eligible = x_min < cent_x < x_max and y_min < cent_y < y_max
if is_eligible:
if x_min_find is None or x0 < x_min_find:
x_min_find = x0
find_key_coordinates = (x0, y0, x1, y1)
if find_key_coordinates is None:
return self.key_top1(coordinates_list, key_coordinates)
else:
return find_key_coordinates
def value_right(self, search_list, key_coordinates, offset_tuple, value_type=None, rigorous=False):
# 字段值查找方向：右侧
x_min, y_min, x_max, y_max = self.get_target_bbox(key_coordinates, offset_tuple)
x_min_find, value, coordinates = None, None, None
for (x0, y0, x1, y1), text in search_list:
if rigorous:
is_eligible = x_min < x0 and x1 < x_max and y_min < y0 and y1 < y_max
else:
cent_x = x0 + ((x1 - x0) / 2)
cent_y = y0 + ((y1 - y0) / 2)
is_eligible = x_min < cent_x < x_max and y_min < cent_y < y_max
if is_eligible:
if x_min_find is None or x0 < x_min_find:
if len(text.strip()) > 0:
x_min_find = x0
value = text
coordinates = (x0, y0, x1, y1)
if isinstance(value_type, str) and value_type in self.replace_map and isinstance(value, str):
new_value = value.translate(str.maketrans(self.replace_map.get(value_type, {})))
return new_value, coordinates
return value, coordinates
def value_under(self, search_list, key_coordinates, offset_tuple, value_type=None, append=False, rigorous=False):
# 字段值查找方向：下方
x_min, y_min, x_max, y_max = self.get_target_bbox(key_coordinates, offset_tuple)
find_list = []
for (x0, y0, x1, y1), text in search_list:
if rigorous:
is_eligible = x_min < x0 and x1 < x_max and y_min < y0 and y1 < y_max
else:
cent_x = x0 + ((x1 - x0) / 2)
cent_y = y0 + ((y1 - y0) / 2)
is_eligible = x_min < cent_x < x_max and y_min < cent_y < y_max
if is_eligible:
if len(text.strip()) > 0:
find_list.append((x0, y0, x1, y1, text))
if len(find_list) == 0:
return None, None
else:
find_list.sort(key=lambda x: (x[1], x[0]))
coordinates = find_list[0][:-1]
if append:
value = ''.join([text for _, _, _, _, text in find_list])
else:
value = find_list[0][-1]
if isinstance(value_type, str) and value_type in self.replace_map and isinstance(value, str):
new_value = value.translate(str.maketrans(self.replace_map.get(value_type, {})))
return new_value, coordinates
return value, coordinates
@staticmethod
def get_target_bbox(key_coordinates, offset_tuple):
offset_xmin, offset_xmax, offset_ymin, offset_ymax = offset_tuple
width = key_coordinates[2] - key_coordinates[0]
height = key_coordinates[-1] - key_coordinates[1]
x_min = key_coordinates[0] - (width * offset_xmin) # -1
x_max = key_coordinates[2] + (width * offset_xmax)
y_min = key_coordinates[1] - (height * offset_ymin) # -1
y_max = key_coordinates[-1] + (height * offset_ymax)
return x_min, y_min, x_max, y_max
def get_target_fields(self, pdf_text_list, pdf_img_list):
pdf_result = dict()
for pno_str, fields_dict in self.target_fields.items():
is_last_pno = False
if pno_str == '-1':
is_last_pno = True
pno_int_list = [int(pno_str) for pno_str in pdf_text_list.keys()]
pno_str = str(max(pno_int_list))
# 搜索关键词
key_text_info = dict()
for key_text_list in fields_dict[self.keys_str].values():
for key_text, key_re_tuple, _, _ in key_text_list:
for (x0, y0, x1, y1), text in pdf_text_list.get(pno_str, []):
for key_re in key_re_tuple:
if re.match(key_re, text):
key_text_info.setdefault(key_text, list()).append((x0, y0, x1, y1))
# 搜索关键词
key_coordinates_info = dict()
for field, key_text_list in fields_dict[self.keys_str].items():
last_key_coordinates = None
for key_text, _, direction, kwargs in key_text_list:
if key_text not in key_text_info:
last_key_coordinates = None
continue
last_key_coordinates = getattr(self, 'key_{0}'.format(direction))(
key_text_info[key_text],
last_key_coordinates,
**kwargs)
key_coordinates_info[field] = last_key_coordinates
# 搜索字段值
page_result = dict()
for field, (source, direction, kwargs, default_value) in fields_dict[self.value_str].items():
if not isinstance(key_coordinates_info.get(field), tuple):
page_result[field] = {
self.words_str: default_value,
self.position_str: self.default_position,
}
continue
value, coordinates = getattr(self, 'value_{0}'.format(direction))(
pdf_text_list.get(pno_str, []) if source == self.text_str else pdf_img_list.get(pno_str, []),
key_coordinates_info[field],
**kwargs
)
if not isinstance(value, str):
page_result[field] = {
self.words_str: default_value,
self.position_str: self.default_position,
}
else:
page_result[field] = {
self.words_str: value,
self.position_str: list(coordinates),
}
# page_key = 'page_12' if is_last_pno else 'page_{0}'.format(int(pno_str) + 1)
page_key = 'page_{0}'.format(int(pno_str) + 1)
pdf_result[page_key] = page_result
return pdf_result