FORMAT_CONVERSION_MAXCOMPUTE/otr/otr_interface.py

import base64
import json
import multiprocessing as mp
import os
# os.environ['TF_XLA_FLAGS'] = '--tf_xla_cpu_global_jit'
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
import time
import traceback
from multiprocessing.context import Process
import multiprocessing
import logging
# from table_line import *
import cv2
import numpy as np
import tensorflow as tf
from flask import Flask, request

from format_convert.utils import request_post, judge_error_code, get_intranet_ip, log, get_md5_from_bytes
from otr.table_line import get_best_predict_size, table_line, get_points, get_split_line, get_points_row, \
    get_points_col, \
    delete_close_points, fix_outline, get_bbox, get_outline_point, table_net, delete_contain_bbox, points_to_line, \
    fix_inner, merge_line, fix_corner, add_continue_bbox, delete_outline
from format_convert import _global


# 接口配置
app = Flask(__name__)


@app.route('/otr', methods=['POST'])
def _otr():
    try:
        if not request.form:
            log("otr no data!")
            return json.dumps({"list_line": str([-9])})

        otr_model = globals().get("global_otr_model")
        if otr_model is None:
            otr_model = OtrModels().get_model()
            globals().update({"global_otr_model": otr_model})

        data = request.form.get("data")
        is_from_pdf = request.form.get("is_from_pdf")
        img_data = base64.b64decode(data)
        _md5 = get_md5_from_bytes(img_data)[0]
        _global.update({"md5": _md5})
        if is_from_pdf:
            list_lines = line_detect(img_data, otr_model, prob=0.2)
        else:
            list_lines = line_detect(img_data, otr_model, prob=0.5)
        return json.dumps(list_lines)
    except TimeoutError:
        return json.dumps({"list_line": str([-5])})
    except:
        traceback.print_exc()
        return json.dumps({"list_line": str([-1])})


def otr(data, otr_model, is_from_pdf):
    try:
        img_data = base64.b64decode(data)
        # points_and_lines = pool.apply(table_detect, (img_data,))
        if is_from_pdf:
            list_lines = line_detect(img_data, otr_model, prob=0.2)
        else:
            list_lines = line_detect(img_data, otr_model, prob=0.5)
        return list_lines
    except TimeoutError:
        raise TimeoutError


flag = 0
# model_path = "models/table-line.h5"
def table_detect2(img_data, otr_model):
    log("into otr_interface table_detect")
    start_time = time.time()
    try:
        start_time1 = time.time()
        # 二进制数据流转np.ndarray [np.uint8: 8位像素]
        img = cv2.imdecode(np.frombuffer(img_data, np.uint8), cv2.IMREAD_COLOR)
        # log("into otr_interface table_detect 1")
        # cv2.imwrite("111111.jpg", img)

        # 将bgr转为rbg
        image_np = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        # log("into otr_interface table_detect 2")

        # 选择与图片最接近分辨率，以防失真
        # best_h, best_w = get_best_predict_size(img)
        print("image_np.shape", image_np.shape)
        best_h, best_w, _ = image_np.shape
        log("otr preprocess time: " + str(round(float(time.time()-start_time1), 4)) + "s")

        # 调用模型
        # rows, cols = table_line(image_np, otr_model)
        start_time1 = time.time()
        rows, cols, image_np = table_line(image_np, otr_model, size=(best_w, best_h), hprob=0.5, vprob=0.5)
        log("otr model predict time: " + str(round(float(time.time()-start_time1), 4)) + "s")

        start_time1 = time.time()
        if not rows or not cols:
            print("points", 0, "split_lines", 0, "bboxes", 0)
            return {"points": str([]), "split_lines": str([]),
                    "bboxes": str([]), "outline_points": str([]),
                    "lines": str([])}

        # 查看是否正确输出rows,cols
        # for line in rows+cols:
        #     cv2.line(img, (int(line[0]), int(line[1])), (int(line[2]), int(line[3])),
        #              (255, 0, 0), 2)
        #     cv2.imshow("rows-cols1", img)
        #     cv2.waitKey(0)

        # 处理结果
        # 合并错开线
        rows = merge_line(rows, axis=0)
        cols = merge_line(cols, axis=1)

        # 计算交点、分割线
        points = get_points(rows, cols, (image_np.shape[0], image_np.shape[1]))
        # log("into otr_interface table_detect 5")
        if not points:
            print("points", 0, "split_lines", 0, "bboxes", 0)
            return {"points": str([]), "split_lines": str([]),
                    "bboxes": str([]), "outline_points": str([]),
                    "lines": str([])}

        # 清掉外围的没用的线
        rows, cols = delete_outline(rows, cols, points)

        split_lines, split_y = get_split_line(points, cols, image_np)
        # log("into otr_interface table_detect 6")

        # 计算交点所在行列，剔除相近交点
        row_point_list = get_points_row(points, split_y, 5)
        col_point_list = get_points_col(points, split_y, 5)
        # log("into otr_interface table_detect 7")
        points = delete_close_points(points, row_point_list, col_point_list)
        # log("into otr_interface table_detect 8")

        # 查看是否正确输出点
        # for p in points:
        #     cv2.circle(img, (p[0], p[1]), 3, (0, 0, 255))
        #     cv2.imshow("points", img)
        #     cv2.waitKey(0)

        # 查看是否正确输出rows,cols
        # for line in rows+cols:
        #     cv2.line(img, (int(line[0]), int(line[1])), (int(line[2]), int(line[3])),
        #              (0, 255, 0), 2)
        #     cv2.imshow("rows-cols0", img)
        #     cv2.waitKey(0)

        # 修复边框
        new_rows, new_cols, long_rows, long_cols = fix_outline(image_np, rows, cols, points,
                                                               split_y)
        # print(new_cols, new_rows)
        if new_rows or new_cols:
            # 连接至补线的延长线
            if long_rows:
                rows = long_rows
            if long_cols:
                cols = long_cols
            # 新的补线
            if new_rows:
                rows += new_rows
            if new_cols:
                cols += new_cols

            # 修复边框后重新计算交点、分割线
            points = get_points(rows, cols, (image_np.shape[0], image_np.shape[1]))

            # log("into otr_interface table_detect 10")
            split_lines, split_y = get_split_line(points, cols, image_np)

            # 计算交点所在行列，剔除相近交点
            row_point_list = get_points_row(points, split_y, 0)
            col_point_list = get_points_col(points, split_y, 0)
            # log("into otr_interface table_detect 11")
            points = delete_close_points(points, row_point_list, col_point_list)
            # row_point_list = get_points_row(points, split_y)
            # col_point_list = get_points_col(points, split_y)
            # log("into otr_interface table_detect 12")

        # 查看是否正确输出rows,cols
        # for line in rows+cols:
        #     cv2.line(img, (int(line[0]), int(line[1])), (int(line[2]), int(line[3])),
        #              (255, 0, 0), 2)
        #     cv2.imshow("rows-cols1", img)
        #     cv2.waitKey(0)

        # 修复表格4个角
        rows, cols = fix_corner(rows, cols, split_y)
        points = get_points(rows, cols, (image_np.shape[0], image_np.shape[1]))
        # row_point_list = get_points_row(points, split_y, 5)
        # col_point_list = get_points_col(points, split_y, 5)

        # print("row_point_list", row_point_list)
        # print("col_point_list", col_point_list)

        # 修复内部缺线
        points = fix_inner(rows, cols, points, split_y)
        if not points:
            print("points", 0, "split_lines", 0, "bboxes", 0)
            return {"points": str([]), "split_lines": str([]),
                    "bboxes": str([]), "outline_points": str([]),
                    "lines": str([])}
        row_point_list = get_points_row(points, split_y, 5)
        col_point_list = get_points_col(points, split_y, 5)

        # 查看是否正确输出点
        # for p in points:
        #     cv2.circle(img, (p[0], p[1]), 1, (0, 255, 0), 3)
        #     cv2.imshow("points fix", img)
        #     cv2.waitKey(0)

        # 查看是否正确输出rows,cols
        # for line in rows+cols:
        #     cv2.line(img, (int(line[0]), int(line[1])), (int(line[2]), int(line[3])),
        #              (255, 0, 0), 2)
        #     cv2.imshow("rows-cols2", img)
        #     cv2.waitKey(0)

        # 根据分行分列重新得到rows、cols,避免线延长导致后续bbox生成失败
        # rows = points_to_line(row_point_list, axis=0)
        # cols = points_to_line(col_point_list, axis=1)
        # points = get_points(rows, cols, (image_np.shape[0], image_np.shape[1]))
        # row_point_list = get_points_row(points, split_y, 0)
        # col_point_list = get_points_col(points, split_y, 0)

        # 获取bbox 单元格
        bboxes = get_bbox(image_np, row_point_list, col_point_list, split_y, rows, cols)
        # log("into otr_interface table_detect 13")

        # 删除包含bbox
        if bboxes:
            bboxes = delete_contain_bbox(bboxes)

        # 查看是否能输出正确框
        # for box in bboxes:
        #     cv2.rectangle(img, box[0], box[1], (0, 0, 255), 3)
        #     cv2.imshow("bbox", img)
        #     cv2.waitKey(0)

        # 补充连续框
        # if bboxes:
        #     bboxes = add_continue_bbox(bboxes)
        #
        #     # 删除包含bbox
        #     bboxes = delete_contain_bbox(bboxes)

        # 查看是否能输出正确框
        # cv2.namedWindow('bbox', 0)
        # for box in bboxes:
        #     cv2.rectangle(img, box[0], box[1], (0, 255, 0), 3)
        #     cv2.imshow("bbox", img)
        #     cv2.waitKey(0)

        # 查看是否正确输出点
        # cv2.namedWindow('points', 0)
        # for p in points:
        #     cv2.circle(img, (p[0], p[1]), 3, (0, 0, 255))
        #     cv2.imshow("points", img)
        #     cv2.waitKey(0)

        # 查看是否正确输出区域分割线
        # cv2.namedWindow('split_lines', 0)
        # for line in split_lines:
        #     cv2.line(img, line[0], line[1], (0, 0, 255), 2)
        #     cv2.imshow("split_lines", img)
        #     cv2.waitKey(0)

        # 获取每个表格的左上右下两个点
        outline_points = get_outline_point(points, split_y)
        # log("into otr_interface table_detect 14")

        if bboxes:
            print("bboxes number", len(bboxes))
            # print("bboxes", bboxes)
        else:
            print("bboxes number", "None")
        log("otr postprocess time: " + str(round(float(time.time()-start_time1), 4)) + "s")
        log("use time: " + str(time.time()-start_time))
        return {"points": str(points), "split_lines": str(split_lines),
                "bboxes": str(bboxes), "outline_points": str(outline_points),
                "lines": str(rows+cols)}

    except TimeoutError:
        raise TimeoutError
    except Exception as e:
        log("otr_interface cannot detected table!")
        print("otr_interface cannot detected table!", traceback.print_exc())
        print("points", 0, "split_lines", 0, "bboxes", 0)
        log("otr postprocess time: " + str(round(float(time.time()-start_time1), 4)) + "s")
        return {"points": str([]), "split_lines": str([]), "bboxes": str([]),
                "outline_points": str([]), "lines": str([])}


def line_detect(img_data, otr_model, prob=0.2):
    log("into otr_interface table_detect")
    start_time = time.time()
    try:
        start_time1 = time.time()
        # 二进制数据流转np.ndarray [np.uint8: 8位像素]
        img = cv2.imdecode(np.frombuffer(img_data, np.uint8), cv2.IMREAD_COLOR)
        # log("into otr_interface table_detect 1")
        # cv2.imwrite("111111.jpg", img)

        # 将bgr转为rbg
        image_np = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        # log("into otr_interface table_detect 2")

        # 选择与图片最接近分辨率，以防失真
        # best_h, best_w = get_best_predict_size(img)
        log("image_np.shape" + str(image_np.shape))
        best_h, best_w, _ = image_np.shape
        log("otr preprocess time: " + str(round(float(time.time()-start_time1), 4)) + "s")

        # 调用模型
        # rows, cols = table_line(image_np, otr_model)
        start_time1 = time.time()
        list_line = table_line(image_np, otr_model, size=(best_w, best_h), prob=prob)
        log("otr finish " + str(round(float(time.time()-start_time1), 4)) + "s")
        return {"list_line": str(list_line)}
    except TimeoutError:
        raise TimeoutError
    except Exception as e:
        log("otr_interface cannot detected table!")
        print("otr_interface cannot detected table!", traceback.print_exc())
        log("otr postprocess time: " + str(round(float(time.time()-start_time1), 4)) + "s")
        return {"list_line": str([])}


class OtrModels:
    def __init__(self):
        # python文件所在目录
        _dir = os.path.abspath(os.path.dirname(os.path.abspath(__file__)))
        model_path = _dir + "/models/table-line.h5"
        self.otr_model = table_net((None, None, 3), 2)
        self.otr_model.load_weights(model_path)

    def get_model(self):
        return self.otr_model


def test_otr_model():
    file_path = "C:/Users/Administrator/Desktop/error2.png"
    with open(file_path, "rb") as f:
        file_bytes = f.read()
    file_base64 = base64.b64encode(file_bytes)
    file_json = {"data": file_base64, "is_from_pdf": False}

    _url = "http://192.168.2.103:18000/otr"
    r = json.loads(request_post(_url, file_json))
    print(r)


# otr_model = table_net((None, None, 3), 2)
# otr_model.load_weights(model_path)
if __name__ == '__main__':
    if len(sys.argv) == 2:
        port = int(sys.argv[1])
    elif len(sys.argv) == 3:
        port = int(sys.argv[1])
        using_gpu_index = int(sys.argv[2])
    else:
        port = 18000
        using_gpu_index = 0
    _global._init()
    _global.update({"port": str(port)})

    # 日志格式设置
    # ip = get_intranet_ip()
    # logging.basicConfig(level=logging.INFO,
    #                     format='%(asctime)s - %(name)s - %(levelname)s - '
    #                            + ip + ' - ' + str(port) + ' - %(message)s')

    # 限制tensorflow显存
    memory_limit_scale = 0.3
    os.environ['CUDA_VISIBLE_DEVICES'] = str(using_gpu_index)
    os.environ['CUDA_CACHE_MAXSIZE'] = str(2147483648)
    os.environ['CUDA_CACHE_DISABLE'] = str(0)
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=memory_limit_scale)
    sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

    app.run(host='0.0.0.0', port=port, processes=1, threaded=False, debug=False)
    log("OTR running "+str(port))

    # test_otr_model()

    # print(json.dumps([-2]))

    # otr_model = OtrModels().get_model()
    # otr("11", otr_model)