FORMAT_CONVERSION_MAXCOMPUTE/atc/atc_interface.py

import base64
import json
import os
import sys
import time
import traceback
from glob import glob
import numpy as np
import cv2
import tensorflow as tf
from PIL import Image
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
MAX_COMPUTE = False

if not MAX_COMPUTE:
    # tensorflow 内存设置
    try:
        gpus = tf.config.list_physical_devices('GPU')
        if len(gpus) > 0:
            tf.config.experimental.set_virtual_device_configuration(
                gpus[0],
                [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=512)])
    except:
        traceback.print_exc()
        # pass
        # gpus = tf.config.list_physical_devices('GPU')
        # for gpu in gpus:  # 如果使用多块GPU时
        #     tf.config.experimental.set_memory_growth(gpu, True)

        os.environ['CUDA_CACHE_MAXSIZE'] = str(2147483648)
        os.environ['CUDA_CACHE_DISABLE'] = str(0)
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.05)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from format_convert import _global
from format_convert.utils import log, bytes2np, request_post, get_md5_from_bytes
from flask import Flask, request
from atc.predict import ModelRelationExtraction


def atc(data, model):
    log("into atc_interface atc")
    try:
        # print("data", data)
        _, classification = model.evaluate(data)
        return {"classification": classification}
    except TimeoutError:
        return {"classification": [-5]}
    except:
        traceback.print_exc()
        return {"classification": [-1]}


# 接口配置
app = Flask(__name__)


@app.route('/atc', methods=['POST'])
def _atc():
    _global._init()
    _global.update({"port": globals().get("port")})
    start_time = time.time()

    log("into atc_interface _atc")
    try:
        if not request.form:
            log("atc no data!")
            return json.dumps({"classification": str([-9])})
        data = request.form.get("data")
        log("atc_interface get data time" + str(time.time()-start_time))

        _md5 = request.form.get("md5")
        _global.update({"md5": _md5})

        atc_model = globals().get("global_atc_model")
        if atc_model is None:
            print("=========== init atc model ===========")
            atc_model = AtcModels().get_model()
            globals().update({"global_atc_model": atc_model})

        classification = atc(data, atc_model).get("classification")
        return json.dumps({"classification": classification})
    except TimeoutError:
        return json.dumps({"classification": str([-5])})
    except:
        traceback.print_exc()
        return json.dumps({"classification": str([-1])})
    finally:
        log("atc interface finish time " + str(time.time()-start_time))


class AtcModels:
    def __init__(self):
        self.model = ModelRelationExtraction()

    def get_model(self):
        return self.model


def test_atc_model(from_remote=False):
    atc_model = AtcModels().get_model()
    text = '''招标文件项目编号：SDGP370302202102000110项目名称:淄川经济开发区中心小学校园智能化采购项目采购人:山东淄川经
    济开发区管理委员会采购代理机构：淄博正益招标有限公司发出日期：2021年8月目录第一章投标邀请7一、项目基本情况7二、申请人的资格要
    求8三、获取招标文件8四、提交投标文件截止时间、开标时间和地点8五、公告期限9六、其他补充事宜9第二章投标人须知11一、总则161.采
    购人、采购代理机构及投标人162.资金来源183.投标费用184.适用法律18二、招标文件185.招标文件构成186.招标文件的澄清与修改207.投
    标截止时间的顺延20三、投标文件的编制208.编制要求209.投标范围及投标文件中标准和计量单位的使用2110.投标文件构成2211.投标报价241
    2.电子版投标文件2513.投标保证金2614.投标有效期2615.投标文件的签署及规定26四、投标文件的递交2616.投标文件的递交2617.递交
    投标文件的截止时间2718.投标文件的接收、修改与撤回27五、开标及评标2719.开标2720.资格审查2821.组建评标委员会2922.投标文件符
    合性审查与澄清3023.投标偏离3224.投标无效3225.比较和评价3326.废标3527.保密要求36六、确定中标3628.中标候选人的确定原则及标
    准3629.确定中标候选人和中标人3630.采购任务取消3631.中标通知书3632.签订合同3633.履约保证金3734.政府采购融资担保3735.预付
    款3736.廉洁自律规定3737.人员回避3738.质疑与接收3739.项目其他相关费用3940.合同公示3941.验收4042.履约验收公示4043.招标文
    件解释权40第三章货物需求41一、项目概述41
    '''
    _md5 = get_md5_from_bytes(text.encode('utf-8'))[0]

    _global._init()
    _global.update({"port": 18060, "md5": _md5})

    if from_remote:
        file_json = {"data": text, "md5": _md5}
        # _url = "http://192.168.2.102:17000/ocr"
        _url = "http://127.0.0.1:18060/atc"
        print(json.loads(request_post(_url, file_json)))
    else:
        result = atc(text, atc_model)
        # print(result)
        if type(result.get("angle")) == list:
            print(result)
        else:
            classification = result.get("classification")
            print("classification", classification)


if __name__ == "__main__":
    test_atc_model()