fcb_photo_review/photo_mask/photo_mask.py

import logging.config
import tempfile
from time import sleep

import cv2
from sqlalchemy import update, and_

from db import MysqlSession
from db.mysql import ZxPhrec, ZxPhhd
from photo_mask import OCR, PHHD_BATCH_SIZE, SLEEP_MINUTES, COPY_TRY_TIMES, UPLOAD_TRY_TIMES, NAME_KEYS, \
    ID_CARD_NUM_KEYS
from ucloud import BUCKET, ucloud
from util import image_util, util


def find_box(content, layout, offset=0, length=None, improve=False, image_path=None):
    full_box = layout[0]
    x_len = full_box[2] - full_box[0]
    y_len = full_box[3] - full_box[1]
    if x_len >= y_len:
        # 横向排布
        box_len = x_len
        direction = "x"
    else:
        # 纵向排布
        box_len = y_len
        direction = "y"
    text = layout[1]
    text_len = len(text)
    char_len = box_len / text_len
    index = text.index(content)

    if not length:
        length = len(content) + 1
    if direction == "x":
        # 横向排布
        box = [
            full_box[0] + (index + offset) * char_len,
            full_box[1],
            full_box[0] + (index + offset + length) * char_len,
            full_box[3],
        ]
    else:
        # 纵向排布
        box = [
            full_box[0],
            full_box[1] + (index + offset) * char_len,
            full_box[2],
            full_box[1] + (index + offset + length) * char_len,
        ]

    if improve:
        # 再次识别，提高精度
        image = cv2.imread(image_path)
        # 截图时偏大一点
        capture_box = util.zoom_rectangle(box, 0.2)
        captured_image = image_util.capture(image, capture_box)
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
            captured_image, offset_x, offset_y = image_util.expand_to_a4_size(captured_image, True)
            cv2.imwrite(temp_file.name, captured_image)
        try:
            layouts = OCR.parse({"doc": temp_file.name})["layout"]
        except TypeError:
            # 如果是类型错误，大概率是没识别到文字
            layouts = []
        except Exception as e:
            # 如果出现其他错误，抛出
            raise e
        for layout in layouts:
            if content in layout[1]:
                temp_box = find_box(content, layout)
                if temp_box:
                    box = [
                        temp_box[0] + capture_box[0] - offset_x,
                        temp_box[1] + capture_box[1] - offset_y,
                        temp_box[2] + capture_box[0] - offset_x,
                        temp_box[3] + capture_box[1] - offset_y,
                    ]
                    break
        util.delete_temp_file(temp_file.name)
    return box


def get_mask_layout(image, name, id_card_num):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
        cv2.imwrite(temp_file.name, image)

    result = []
    try:
        try:
            # layouts = util.get_ocr_layout(OCR, temp_file.name)
            layouts = OCR.parse({"doc": temp_file.name})["layout"]
        except TypeError:
            # 如果是类型错误，大概率是没识别到文字
            layouts = []
        except Exception as e:
            # 如果出现其他错误，抛出
            raise e

        if not layouts:
            # 无识别结果
            return result
        else:
            # 涂抹
            for layout in layouts:
                find_name_by_key = True
                find_id_card_num_by_key = True
                if name in layout[1]:
                    result.append(find_box(name, layout, improve=True, image_path=temp_file.name))
                    find_name_by_key = False
                if id_card_num in layout[1]:
                    result.append(find_box(id_card_num, layout, improve=True, image_path=temp_file.name))
                    find_id_card_num_by_key = False

                keys = []
                if find_name_by_key:
                    keys += NAME_KEYS
                if find_id_card_num_by_key:
                    keys += ID_CARD_NUM_KEYS
                for key in keys:
                    if key["key"] in layout[1]:
                        result.append(find_box(key["key"], layout, offset=len(key["key"]), length=key["length"]))
            return result
    except Exception as e:
        logging.error("涂抹时出错！", exc_info=e)
    finally:
        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_(
        ZxPhrec.pk_phhd == pk_phhd,
        ZxPhrec.cRectype.in_(["3", "4"])
    )).all()
    session.close()
    for phrec in phrecs:
        img_url = ucloud.get_private_url(phrec.cfjaddress)
        if not img_url:
            continue
        # 是否有涂抹
        is_masked = False
        # 打开图片
        image = image_util.read(img_url)
        split_results = image_util.split(image)
        for split_result in split_results:
            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:
                result = (
                    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)

        # 如果涂抹了要备份以及更新
        if is_masked:
            for i in range(COPY_TRY_TIMES):
                is_copy_success = ucloud.copy_file(BUCKET, phrec.cfjaddress, "drg2015", phrec.cfjaddress)
                if is_copy_success:
                    break

            with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
                cv2.imwrite(temp_file.name, image)
            try:
                for i in range(UPLOAD_TRY_TIMES):
                    is_upload_success = ucloud.upload_file(phrec.cfjaddress, temp_file.name)
                    if is_upload_success:
                        break
            except Exception as e:
                logging.error("上传图片出错", exc_info=e)
            finally:
                util.delete_temp_file(temp_file.name)


def main():
    while 1:
        session = MysqlSession()
        phhds = session.query(ZxPhhd.pk_phhd, ZxPhhd.cXm, ZxPhhd.cSfzh).filter(
            ZxPhhd.paint_flag == "1"
        ).limit(PHHD_BATCH_SIZE).all()
        # 将状态改为正在涂抹中
        pk_phhd_values = [phhd.pk_phhd for phhd in phhds]
        update_flag = (update(ZxPhhd).where(ZxPhhd.pk_phhd.in_(pk_phhd_values)).values(paint_flag="2"))
        session.execute(update_flag)
        session.commit()
        session.close()
        if phhds:
            for phhd in phhds:
                pk_phhd = phhd.pk_phhd
                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",
                    paint_date=util.get_default_datetime()))
                session.execute(update_flag)
                session.commit()
                session.close()
        else:
            # 没有查询到新案子，等待一段时间后再查
            logging.info(f"暂未查询到需要涂抹的案子，等待{SLEEP_MINUTES}分钟...")
            sleep(SLEEP_MINUTES * 60)