diff --git a/photo_mask/__init__.py b/photo_mask/__init__.py index 0e67900..d0fc862 100644 --- a/photo_mask/__init__.py +++ b/photo_mask/__init__.py @@ -1,4 +1,4 @@ -from paddleocr import PaddleOCR +from paddlenlp.utils.doc_parser import DocParser """ 项目配置 @@ -25,4 +25,5 @@ NAME_KEYS = [ ] ID_CARD_NUM_KEYS = [{"key": "身份证号", "length": 19}, ] -OCR = PaddleOCR(use_angle_cls=False, show_log=False, gpu_id=1) +# OCR = PaddleOCR(use_angle_cls=False, show_log=False, gpu_id=1) +OCR = DocParser(use_gpu=True, layout_analysis=False) diff --git a/photo_mask/photo_mask.py b/photo_mask/photo_mask.py index 63e1a3b..4ee4681 100644 --- a/photo_mask/photo_mask.py +++ b/photo_mask/photo_mask.py @@ -88,7 +88,8 @@ def get_mask_layout(image, name, id_card_num): result = [] try: try: - layouts = util.get_ocr_layout(OCR, temp_file.name) + # layouts = util.get_ocr_layout(OCR, temp_file.name) + layouts = OCR.parse({"doc": temp_file.name})["layout"] except TypeError: # 如果是类型错误,大概率是没识别到文字 layouts = [] @@ -126,6 +127,14 @@ def get_mask_layout(image, name, id_card_num): util.delete_temp_file(temp_file.name) +def handle_image_for_mask(split_result): + expand_img, offset_x, offset_y = image_util.expand_to_a4_size(split_result["img"], True) + split_result["x_offset"] -= offset_x + split_result["y_offset"] -= offset_y + gray_image = cv2.cvtColor(expand_img, cv2.COLOR_BGR2GRAY) + return gray_image, split_result["x_offset"], split_result["y_offset"] + + def photo_mask(pk_phhd, name, id_card_num): session = MysqlSession() phrecs = session.query(ZxPhrec.cfjaddress).filter(and_( @@ -143,33 +152,18 @@ def photo_mask(pk_phhd, name, id_card_num): image = image_util.read(img_url) split_results = image_util.split(image) for split_result in split_results: - angles = image_util.parse_rotation_angles(split_result["img"]) - angle = int(angles[0]) - rotated_img = image_util.rotate(split_result["img"], angle) - rotated_img, offset_x, offset_y = image_util.expand_to_a4_size(rotated_img, True) - split_result["x_offset"] -= offset_x - split_result["y_offset"] -= offset_y - results = get_mask_layout(rotated_img, name, id_card_num) - if not results: - angle = int(angles[1]) - rotated_img = image_util.rotate(split_result["img"], angle) - results = get_mask_layout(rotated_img, name, id_card_num) - if not results and "0" not in angles: - angle = 0 - results = get_mask_layout(split_result["img"], name, id_card_num) + to_mask_img, x_offset, y_offset = handle_image_for_mask(split_result) + results = get_mask_layout(to_mask_img, name, id_card_num) if results: is_masked = True for result in results: - height, width = split_result["img"].shape[:2] - center = (width / 2, height / 2) - result = image_util.invert_rotate_rectangle(result, center, angle) result = ( - result[0] + split_result["x_offset"], - result[1] + split_result["y_offset"], - result[2] + split_result["x_offset"], - result[3] + split_result["y_offset"], + result[0] + x_offset, + result[1] + y_offset, + result[2] + x_offset, + result[3] + y_offset, ) cv2.rectangle(image, (int(result[0]), int(result[1])), (int(result[2]), int(result[3])), (255, 255, 255), -1, 0) @@ -212,7 +206,7 @@ def main(): logging.info(f"开始涂抹:{pk_phhd}") photo_mask(pk_phhd, phhd.cXm, phhd.cSfzh) - # 识别完成更新标识 + # 涂抹完成更新标识 session = MysqlSession() update_flag = (update(ZxPhhd).where(ZxPhhd.pk_phhd == pk_phhd).values(paint_flag="8")) session.execute(update_flag)