add f3 folder
Showing
2 changed files
with
378 additions
and
0 deletions
1 | import os | ||
2 | import time | ||
3 | import json | ||
4 | import shutil | ||
5 | import base64 | ||
6 | import signal | ||
7 | import requests | ||
8 | import traceback | ||
9 | from django.core.management import BaseCommand | ||
10 | from multiprocessing import Process | ||
11 | from openpyxl import load_workbook, Workbook | ||
12 | |||
13 | from settings import conf | ||
14 | from common.mixins import LoggerMixin | ||
15 | from common.tools.pdf_to_img import PDFHandler, PDFBuild | ||
16 | from apps.doc import consts | ||
17 | from apps.doc.exceptions import OCR1Exception, OCR2Exception | ||
18 | |||
19 | |||
20 | class Command(BaseCommand, LoggerMixin): | ||
21 | |||
22 | def __init__(self): | ||
23 | super().__init__() | ||
24 | self.log_base = '[folder f3 process]' | ||
25 | # input folder | ||
26 | self.input_dir = conf.F3_DIR | ||
27 | # 处理文件开关 | ||
28 | self.switch = True | ||
29 | # 睡眠时间 | ||
30 | self.sleep_time = float(conf.SLEEP_SECOND_FOLDER) | ||
31 | # 输出结果 | ||
32 | self.wb_name = 'result.xlsx' | ||
33 | |||
34 | self.field_map = { | ||
35 | # sheet_name, key_field, side_field_order, src_field_order | ||
36 | consts.IC_CLASSIFY: (consts.IC_CN_NAME, '有效期限', consts.IC_FIELD_ORDER_3, consts.IC_FIELD_ORDER_2), | ||
37 | } | ||
38 | |||
39 | self.pdf_name_map = { | ||
40 | consts.IC_CLASSIFY: consts.IC_CN_NAME, | ||
41 | consts.BL_CLASSIFY: consts.BL_CN_NAME, | ||
42 | } | ||
43 | |||
44 | # ocr url | ||
45 | self.ocr_url = conf.OCR_URL_FOLDER | ||
46 | self.ocr_url_2 = conf.OCR2_URL_FOLDER | ||
47 | |||
48 | # 优雅退出信号:15 | ||
49 | signal.signal(signal.SIGTERM, self.signal_handler) | ||
50 | |||
51 | def signal_handler(self, sig, frame): | ||
52 | self.switch = False # 停止处理文件 | ||
53 | |||
54 | @staticmethod | ||
55 | def license1_process(ocr_data, ocr_res, classify): | ||
56 | # 类别:'0'身份证, '1'居住证 | ||
57 | license_data = ocr_data.get('data') | ||
58 | if not license_data: | ||
59 | return | ||
60 | if isinstance(license_data, dict): | ||
61 | license_data.pop('base64_img', '') | ||
62 | if classify == consts.IC_CLASSIFY: | ||
63 | id_card_dict = {} | ||
64 | card_type = license_data.get('type', '') | ||
65 | is_ic = card_type.startswith('身份证') | ||
66 | is_info_side = card_type.endswith('信息面') | ||
67 | id_card_dict['类别'] = '0' if is_ic else '1' | ||
68 | if is_ic: | ||
69 | field_map = consts.IC_MAP_0 if is_info_side else consts.IC_MAP_1 | ||
70 | else: | ||
71 | field_map = consts.RP_MAP_0 if is_info_side else consts.RP_MAP_1 | ||
72 | for write_field, search_field in field_map: | ||
73 | id_card_dict[write_field] = license_data.get('words_result', {}).get(search_field, {}).get('words', '') | ||
74 | if not is_info_side: | ||
75 | start_time = license_data.get('words_result', {}).get('签发日期', {}).get('words', '') | ||
76 | end_time = license_data.get('words_result', {}).get('失效日期', {}).get('words', '') | ||
77 | id_card_dict['有效期限'] = '{0}-{1}'.format(start_time, end_time) | ||
78 | # for id_card_dict in license_data: | ||
79 | # try: | ||
80 | # id_card_dict.pop('base64_img') | ||
81 | # except Exception as e: | ||
82 | # continue | ||
83 | ocr_res.setdefault(classify, []).append(id_card_dict) | ||
84 | |||
85 | def license2_process(self, ocr_data, ocr_res, classify, img_path): | ||
86 | if classify != consts.BL_CLASSIFY: | ||
87 | return | ||
88 | |||
89 | pid, _, _, _, _, _ = consts.LICENSE_CLASSIFY_MAPPING.get(classify) | ||
90 | file_data = ocr_data.get('section_img') | ||
91 | if file_data is None: | ||
92 | with open(img_path, 'rb') as f: | ||
93 | base64_data = base64.b64encode(f.read()) | ||
94 | # 获取解码后的base64值 | ||
95 | file_data = base64_data.decode() | ||
96 | json_data_2 = { | ||
97 | "pid": str(pid), | ||
98 | "filedata": file_data | ||
99 | } | ||
100 | |||
101 | for times in range(consts.RETRY_TIMES): | ||
102 | try: | ||
103 | start_time = time.time() | ||
104 | ocr_2_response = requests.post(self.ocr_url_2, data=json_data_2) | ||
105 | if ocr_2_response.status_code != 200: | ||
106 | raise OCR2Exception('ocr_2 status code: {0}'.format(ocr_2_response.status_code)) | ||
107 | except Exception as e: | ||
108 | self.folder_log.warn( | ||
109 | '{0} [ocr_2 failed] [times={1}] [img_path={2}] [error={3}]'.format( | ||
110 | self.log_base, times, img_path, traceback.format_exc())) | ||
111 | else: | ||
112 | ocr_res_2 = json.loads(ocr_2_response.text) | ||
113 | end_time = time.time() | ||
114 | speed_time = int(end_time - start_time) | ||
115 | self.folder_log.info( | ||
116 | '{0} [ocr_2 success] [img={1}] [speed_time={2}]'.format( | ||
117 | self.log_base, img_path, speed_time)) | ||
118 | |||
119 | if ocr_res_2.get('ErrorCode') in consts.SUCCESS_CODE_SET: | ||
120 | if pid == consts.BC_PID: | ||
121 | ocr_res.append(ocr_res_2) | ||
122 | else: | ||
123 | # 营业执照等 | ||
124 | for result_dict in ocr_res_2.get('ResultList', []): | ||
125 | res_dict = {} | ||
126 | for field_dict in result_dict.get('FieldList', []): | ||
127 | res_dict[field_dict.get('chn_key', '')] = field_dict.get('value', '') | ||
128 | ocr_res.append(res_dict) | ||
129 | break | ||
130 | |||
131 | def wb_process(self, ocr_res, output_dir): | ||
132 | excel_path = os.path.join(output_dir, self.wb_name) | ||
133 | |||
134 | try: | ||
135 | if os.path.exists(excel_path): | ||
136 | wb = load_workbook(excel_path) | ||
137 | else: | ||
138 | wb = Workbook() | ||
139 | |||
140 | for c, res_list in ocr_res.items(): | ||
141 | sheet_name, key_field, side_field_order, src_field_order = self.field_map.get(c) | ||
142 | if sheet_name in wb.sheetnames: | ||
143 | ws = wb.get_sheet_by_name(sheet_name) | ||
144 | else: | ||
145 | ws = wb.create_sheet(sheet_name) | ||
146 | |||
147 | for res in res_list: | ||
148 | if key_field is not None and key_field in res: | ||
149 | field_order = side_field_order | ||
150 | else: | ||
151 | field_order = src_field_order | ||
152 | for search_field, write_field in field_order: | ||
153 | field_value = res.get(search_field, '') | ||
154 | if isinstance(field_value, list): | ||
155 | ws.append((write_field, *field_value)) | ||
156 | else: | ||
157 | ws.append((write_field, field_value)) | ||
158 | ws.append((None,)) | ||
159 | |||
160 | if 'Sheet' in wb.sheetnames and len(wb.sheetnames) > 1: | ||
161 | wb.remove(wb.get_sheet_by_name('Sheet')) | ||
162 | wb.save(excel_path) | ||
163 | except Exception as e: | ||
164 | self.folder_log.error('{0} [wb build error] [path={1}] [error={2}]'.format( | ||
165 | self.log_base, excel_path, traceback.format_exc())) | ||
166 | |||
167 | def ocr_process(self, img_path, all_ocr_res, img_res): | ||
168 | if os.path.exists(img_path): | ||
169 | with open(img_path, 'rb') as f: | ||
170 | base64_data = base64.b64encode(f.read()) | ||
171 | # 获取解码后的base64值 | ||
172 | file_data = base64_data.decode() | ||
173 | json_data = { | ||
174 | "file": file_data, | ||
175 | "channel": consts.AFC_PREFIX, | ||
176 | } | ||
177 | |||
178 | for times in range(consts.RETRY_TIMES): | ||
179 | try: | ||
180 | start_time = time.time() | ||
181 | ocr_response = requests.post(self.ocr_url, json=json_data) | ||
182 | if ocr_response.status_code != 200: | ||
183 | raise OCR1Exception('{0} ocr status code: {1}'.format(self.log_base, ocr_response.status_code)) | ||
184 | except Exception as e: | ||
185 | self.folder_log.warn('{0} [ocr failed] [times={1}] [img_path={2}] [error={3}]'.format( | ||
186 | self.log_base, times, img_path, traceback.format_exc())) | ||
187 | else: | ||
188 | ocr_res = ocr_response.json() | ||
189 | end_time = time.time() | ||
190 | speed_time = int(end_time - start_time) | ||
191 | self.folder_log.info('{0} [ocr success] [img={1}] [speed_time={2}]'.format( | ||
192 | self.log_base, img_path, speed_time)) | ||
193 | |||
194 | if isinstance(ocr_res, dict): | ||
195 | if ocr_res.get('code') == 1: | ||
196 | data_list = ocr_res.get('data', []) | ||
197 | if isinstance(data_list, list): | ||
198 | for ocr_data in data_list: | ||
199 | classify = ocr_data.get('classify') | ||
200 | img_res.setdefault(classify, set()).add(img_path) | ||
201 | if classify in consts.LICENSE_CLASSIFY_SET_1: | ||
202 | self.license1_process(ocr_data, all_ocr_res, classify) | ||
203 | elif classify in consts.LICENSE_CLASSIFY_SET_2: | ||
204 | self.license2_process(ocr_data, all_ocr_res, classify, img_path) | ||
205 | break | ||
206 | else: | ||
207 | self.folder_log.warn('{0} [ocr failed] [img_path={1}]'.format(self.log_base, img_path)) | ||
208 | |||
209 | def img_process(self, img_res, output_dir): | ||
210 | for classify, img_path_set in img_res.items(): | ||
211 | pdf_path = os.path.join(output_dir, '{0}.pdf'.format(self.pdf_name_map.get(classify, '其他'))) | ||
212 | pdf_build = PDFBuild(pdf_path) | ||
213 | pdf_build.insert_img(img_path_set) | ||
214 | |||
215 | def images_process(self, img_path_list, output_dir): | ||
216 | ocr_res = dict() | ||
217 | img_res = dict() | ||
218 | for img_path in img_path_list: | ||
219 | self.ocr_process(img_path, ocr_res, img_res) | ||
220 | self.wb_process(ocr_res, output_dir) | ||
221 | self.img_process(img_res, output_dir) | ||
222 | |||
223 | def pdf_process(self, name, path, img_output_dir, output_dir): | ||
224 | pdf_handler = PDFHandler(path, os.path.join(img_output_dir, name)) | ||
225 | |||
226 | try: | ||
227 | self.folder_log.info('{0} [pdf to img start] [path={1}]'.format(self.log_base, path)) | ||
228 | pdf_handler.extract_image() | ||
229 | self.folder_log.info('{0} [pdf to img end] [path={1}]'.format(self.log_base, path)) | ||
230 | except Exception as e: | ||
231 | self.folder_log.error('{0} [pdf to img error] [path={1}] [error={2}]'.format( | ||
232 | self.log_base, path, traceback.format_exc())) | ||
233 | raise e | ||
234 | else: | ||
235 | self.images_process(pdf_handler.img_path_list, output_dir) | ||
236 | |||
237 | def folder_process(self, unique_folder_name, folder_path, main_output_dir): | ||
238 | output_dir = os.path.join(main_output_dir, unique_folder_name) | ||
239 | |||
240 | img_output_dir = os.path.join(output_dir, 'image') | ||
241 | failed_output_dir = os.path.join(output_dir, 'failed') | ||
242 | os.makedirs(output_dir, exist_ok=True) | ||
243 | os.makedirs(img_output_dir, exist_ok=True) | ||
244 | os.makedirs(failed_output_dir, exist_ok=True) | ||
245 | |||
246 | os_error_filename_set = set() | ||
247 | |||
248 | list_dir = os.listdir(folder_path) | ||
249 | if not list_dir and len(os_error_filename_set) == 0: | ||
250 | self.folder_log.info('{0} [folder empty, completed] [path={1}]'.format(self.log_base, folder_path)) | ||
251 | return | ||
252 | |||
253 | all_file_set = set(list_dir) | ||
254 | true_file_set = all_file_set - os_error_filename_set | ||
255 | if len(true_file_set) == 0 and len(os_error_filename_set) > 0: | ||
256 | true_file_set.add(os_error_filename_set.pop()) | ||
257 | for name in true_file_set: | ||
258 | time.sleep(5) | ||
259 | path = os.path.join(folder_path, name) | ||
260 | |||
261 | try: | ||
262 | if not os.path.exists(path): | ||
263 | self.folder_log.info('{0} [path is not exists] [path={1}]'.format(self.log_base, path)) | ||
264 | continue | ||
265 | elif os.path.isfile(path): | ||
266 | self.folder_log.info('{0} [file start] [path={1}]'.format(self.log_base, path)) | ||
267 | self.pdf_process(name, path, img_output_dir, output_dir) | ||
268 | self.folder_log.info('{0} [file end] [path={1}]'.format(self.log_base, path)) | ||
269 | else: | ||
270 | self.folder_log.info('{0} [path is dir] [path={1}]'.format(self.log_base, path)) | ||
271 | failed_path = os.path.join(failed_output_dir, name) | ||
272 | shutil.copyfile(path, failed_path) | ||
273 | except OSError: | ||
274 | os_error_filename_set.add(name) | ||
275 | self.folder_log.error('{0} [os error] [path={1}] [error={2}]'.format( | ||
276 | self.log_base, path, traceback.format_exc())) | ||
277 | except Exception as e: | ||
278 | try: | ||
279 | self.folder_log.error('{0} [file error] [path={1}] [error={2}]'.format(self.log_base, path, | ||
280 | traceback.format_exc())) | ||
281 | failed_path = os.path.join(failed_output_dir, name) | ||
282 | shutil.copyfile(path, failed_path) | ||
283 | except Exception as e: | ||
284 | os_error_filename_set.add(name) | ||
285 | self.folder_log.error('{0} [file move error] [path={1}] [error={2}]'.format( | ||
286 | self.log_base, path, traceback.format_exc())) | ||
287 | |||
288 | def main_folder_process(self, input_dir): | ||
289 | while not os.path.isdir(input_dir): | ||
290 | self.folder_log.info('{0} [input dir is not dir] [input_dir={1}]'.format(self.log_base, input_dir)) | ||
291 | if self.switch: | ||
292 | time.sleep(self.sleep_time) | ||
293 | continue | ||
294 | else: | ||
295 | return | ||
296 | |||
297 | output_dir = os.path.join(os.path.dirname(input_dir), 'Output') | ||
298 | completed_output_dir = os.path.join(output_dir, 'Completed') | ||
299 | failed_output_dir = os.path.join(output_dir, 'Failed') | ||
300 | os.makedirs(output_dir, exist_ok=True) | ||
301 | os.makedirs(completed_output_dir, exist_ok=True) | ||
302 | os.makedirs(failed_output_dir, exist_ok=True) | ||
303 | |||
304 | os_error_filename_set = set() | ||
305 | while self.switch: | ||
306 | # 1. 从input dir获取pdf or image | ||
307 | list_dir = os.listdir(input_dir) | ||
308 | if not list_dir and len(os_error_filename_set) == 0: | ||
309 | self.folder_log.info('{0} [input dir empty] [input_dir={1}]'.format(self.log_base, input_dir)) | ||
310 | time.sleep(self.sleep_time) | ||
311 | continue | ||
312 | all_file_set = set(list_dir) | ||
313 | true_file_set = all_file_set - os_error_filename_set | ||
314 | if len(true_file_set) == 0 and len(os_error_filename_set) > 0: | ||
315 | true_file_set.add(os_error_filename_set.pop()) | ||
316 | for name in true_file_set: | ||
317 | time.sleep(10) | ||
318 | unique_folder_name = '{0}_{1}'.format(time.time(), name) | ||
319 | path = os.path.join(input_dir, name) | ||
320 | |||
321 | try: | ||
322 | if not os.path.exists(path): | ||
323 | self.folder_log.info('{0} [path is not exists] [path={1}]'.format(self.log_base, path)) | ||
324 | continue | ||
325 | elif os.path.isdir(path): | ||
326 | self.folder_log.info('{0} [dir start] [path={1}]'.format(self.log_base, name)) | ||
327 | self.folder_process(unique_folder_name, path, output_dir) | ||
328 | completed_path = os.path.join(completed_output_dir, unique_folder_name) | ||
329 | shutil.move(path, completed_path) | ||
330 | self.folder_log.info('{0} [dir end] [path={1}]'.format(self.log_base, name)) | ||
331 | else: | ||
332 | self.folder_log.info('{0} [path is not dir] [path={1}]'.format(self.log_base, path)) | ||
333 | failed_path = os.path.join(failed_output_dir, unique_folder_name) | ||
334 | shutil.move(path, failed_path) | ||
335 | except OSError: | ||
336 | os_error_filename_set.add(name) | ||
337 | self.folder_log.error('{0} [os error] [path={1}] [error={2}]'.format( | ||
338 | self.log_base, path, traceback.format_exc())) | ||
339 | except Exception as e: | ||
340 | try: | ||
341 | self.folder_log.error('{0} [file error] [path={1}] [error={2}]'.format(self.log_base, path, | ||
342 | traceback.format_exc())) | ||
343 | failed_path = os.path.join(failed_output_dir, unique_folder_name) | ||
344 | shutil.move(path, failed_path) | ||
345 | except Exception as e: | ||
346 | os_error_filename_set.add(name) | ||
347 | self.folder_log.error('{0} [file move error] [path={1}] [error={2}]'.format( | ||
348 | self.log_base, path, traceback.format_exc())) | ||
349 | |||
350 | def handle(self, *args, **kwargs): | ||
351 | process_list = [] | ||
352 | |||
353 | process = Process(target=self.main_folder_process, args=(self.input_dir, )) | ||
354 | process_list.append(process) | ||
355 | |||
356 | for p in process_list: | ||
357 | p.start() | ||
358 | for p in process_list: | ||
359 | p.join() | ||
360 | |||
361 | self.folder_log.info('{0} [stop safely]'.format(self.log_base)) |
... | @@ -24,6 +24,23 @@ WH_COUPLE_4 = (100, 300) | ... | @@ -24,6 +24,23 @@ WH_COUPLE_4 = (100, 300) |
24 | WH_COUPLE_5 = (100, 200) | 24 | WH_COUPLE_5 = (100, 200) |
25 | 25 | ||
26 | 26 | ||
27 | class PDFBuild: | ||
28 | |||
29 | def __init__(self, path): | ||
30 | self.path = path | ||
31 | |||
32 | def insert_img(self, img_path_list): | ||
33 | if os.path.exists(self.path): | ||
34 | pdf = fitz.Document(self.path) | ||
35 | else: | ||
36 | pdf = fitz.Document() | ||
37 | for img_path in img_path_list: | ||
38 | new_page = pdf.newPage() | ||
39 | new_page.insertImage(new_page.rect, img_path) | ||
40 | pdf.save(self.path) | ||
41 | pdf.close() | ||
42 | |||
43 | |||
27 | class PDFHandler: | 44 | class PDFHandler: |
28 | 45 | ||
29 | def __init__(self, path, img_dir_path, document_name=None): | 46 | def __init__(self, path, img_dir_path, document_name=None): | ... | ... |
-
Please register or sign in to post a comment