167 lines
5.6 KiB
Python
167 lines
5.6 KiB
Python
|
# ===============================================
|
|||
|
|
# =============== OCR - 任务管理器 ===============
|
|||
|
|
# ===============================================
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
一种任务管理器为全局单例，不同标签页要执行同一种任务，要访问对应的任务管理器。
|
|||
|
|
任务管理器中有一个引擎API实例，所有任务均使用该API。
|
|||
|
|
标签页可以向任务管理器提交一组任务队列，其中包含了每一项任务的信息，及总体的参数和回调。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
from umi_log import logger
|
|||
|
|
from .mission import Mission
|
|||
|
|
from ..ocr.tbpu import getParser, IgnoreArea
|
|||
|
|
from ..ocr.api import getApiOcr, getLocalOptions
|
|||
|
|
from ..utils.utils import argdIntConvert
|
|||
|
|
|
|||
|
|
# 合法文件后缀
|
|||
|
|
ImageSuf = [
|
|||
|
|
".jpg",
|
|||
|
|
".jpe",
|
|||
|
|
".jpeg",
|
|||
|
|
".jfif",
|
|||
|
|
".png",
|
|||
|
|
".webp",
|
|||
|
|
".bmp",
|
|||
|
|
".tif",
|
|||
|
|
".tiff",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class __MissionOcrClass(Mission):
|
|||
|
|
def __init__(self):
|
|||
|
|
super().__init__()
|
|||
|
|
self._apiKey = "" # 当前api类型
|
|||
|
|
self._api = None # 当前引擎api对象
|
|||
|
|
|
|||
|
|
# ========================= 【重载】 =========================
|
|||
|
|
|
|||
|
|
# msnInfo: { 回调函数"onXX", 参数"argd":{"tbpu.xx", "ocr.xx"} }
|
|||
|
|
# msnList: [ { "path", "bytes", "base64" } ]
|
|||
|
|
def addMissionList(self, msnInfo, msnList): # 添加任务列表
|
|||
|
|
# 实例化 tbpu 文本后处理模块
|
|||
|
|
msnInfo["tbpu"] = []
|
|||
|
|
argd = msnInfo["argd"]
|
|||
|
|
# 忽略区域
|
|||
|
|
if "tbpu.ignoreArea" in argd:
|
|||
|
|
iArea = argd["tbpu.ignoreArea"]
|
|||
|
|
if isinstance(iArea, list) and len(iArea) > 0:
|
|||
|
|
msnInfo["tbpu"].append(IgnoreArea(iArea))
|
|||
|
|
# 获取排版解析器对象
|
|||
|
|
if "tbpu.parser" in argd:
|
|||
|
|
msnInfo["tbpu"].append(getParser(argd["tbpu.parser"]))
|
|||
|
|
# 检查任务合法性
|
|||
|
|
for i in range(len(msnList) - 1, -1, -1):
|
|||
|
|
if "path" in msnList[i]:
|
|||
|
|
p = msnList[i]["path"]
|
|||
|
|
if os.path.splitext(p)[-1].lower() not in ImageSuf:
|
|||
|
|
logger.warning(f"添加OCR任务时，第{i}项的路径path不是图片：{p}")
|
|||
|
|
del msnList[i]
|
|||
|
|
elif "bytes" not in msnList[i] and "base64" not in msnList[i]:
|
|||
|
|
logger.warning(f"添加OCR任务时，第{i}项不含 path、bytes、base64")
|
|||
|
|
del msnList[i]
|
|||
|
|
return super().addMissionList(msnInfo, msnList)
|
|||
|
|
|
|||
|
|
def msnPreTask(self, msnInfo): # 用于更新api和参数
|
|||
|
|
# 检查API对象
|
|||
|
|
if not self._api:
|
|||
|
|
return "[Error] MissionOCR: API object is None."
|
|||
|
|
# 检查参数更新
|
|||
|
|
startInfo = self._dictShortKey(msnInfo["argd"])
|
|||
|
|
# 恢复int类型
|
|||
|
|
argdIntConvert(startInfo)
|
|||
|
|
msg = self._api.start(startInfo)
|
|||
|
|
if msg.startswith("[Error]"):
|
|||
|
|
logger.error(f"OCR引擎启动失败： {msg}")
|
|||
|
|
return msg # 更新失败，结束该队列
|
|||
|
|
else:
|
|||
|
|
return "" # 更新成功 TODO: continue
|
|||
|
|
|
|||
|
|
def msnTask(self, msnInfo, msn): # 执行msn
|
|||
|
|
if "path" in msn:
|
|||
|
|
res = self._api.runPath(msn["path"])
|
|||
|
|
res["path"] = msn["path"] # 结果字典中补充参数
|
|||
|
|
elif "bytes" in msn:
|
|||
|
|
res = self._api.runBytes(msn["bytes"])
|
|||
|
|
elif "base64" in msn:
|
|||
|
|
res = self._api.runBase64(msn["base64"])
|
|||
|
|
else:
|
|||
|
|
res = {
|
|||
|
|
"code": 901,
|
|||
|
|
"data": f"[Error] Unknown task type.\n【异常】未知的任务类型。\n{str(msn)[:100]}",
|
|||
|
|
}
|
|||
|
|
# 任务成功时的后处理
|
|||
|
|
if res["code"] == 100:
|
|||
|
|
# 计算平均置信度
|
|||
|
|
score, num = 0, 0
|
|||
|
|
for r in res["data"]:
|
|||
|
|
score += r["score"]
|
|||
|
|
num += 1
|
|||
|
|
if num > 0:
|
|||
|
|
score /= num
|
|||
|
|
res["score"] = score
|
|||
|
|
# 执行 tbpu
|
|||
|
|
if msnInfo["tbpu"]:
|
|||
|
|
for tbpu in msnInfo["tbpu"]:
|
|||
|
|
res["data"] = tbpu.run(res["data"])
|
|||
|
|
# 如果忽略区域等处理将所有文本删除，则结束tbpu
|
|||
|
|
if not res["data"]:
|
|||
|
|
res["code"] = 101
|
|||
|
|
res["data"] = ""
|
|||
|
|
break
|
|||
|
|
return res
|
|||
|
|
|
|||
|
|
# ========================= 【qml接口】 =========================
|
|||
|
|
|
|||
|
|
def getStatus(self): # 返回当前状态
|
|||
|
|
return {
|
|||
|
|
"apiKey": self._apiKey,
|
|||
|
|
"missionListsLength": self.getMissionListsLength(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def setApi(self, apiKey, info): # 设置api
|
|||
|
|
# 成功返回 [Success] ，失败返回 [Error] 开头的字符串
|
|||
|
|
self._apiKey = apiKey
|
|||
|
|
info = self._dictShortKey(info)
|
|||
|
|
# 如果api对象已启动，则先停止
|
|||
|
|
if self._api:
|
|||
|
|
self._api.stop()
|
|||
|
|
# 获取新api对象
|
|||
|
|
res = getApiOcr(apiKey, info)
|
|||
|
|
# 失败
|
|||
|
|
if isinstance(res, str):
|
|||
|
|
self._apiKey = ""
|
|||
|
|
self._api = None
|
|||
|
|
return res
|
|||
|
|
# 成功
|
|||
|
|
else:
|
|||
|
|
self._api = res
|
|||
|
|
return "[Success]"
|
|||
|
|
|
|||
|
|
# 将字典中配置项的长key转为短key
|
|||
|
|
# 如： ocr.win32_PaddleOCR-json.path → path
|
|||
|
|
def _dictShortKey(self, d):
|
|||
|
|
newD = {}
|
|||
|
|
key1 = "ocr."
|
|||
|
|
key2 = key1 + self._apiKey + "."
|
|||
|
|
for k in d:
|
|||
|
|
if k.startswith(key2):
|
|||
|
|
newD[k[len(key2) :]] = d[k]
|
|||
|
|
elif k.startswith(key1):
|
|||
|
|
newD[k[len(key1) :]] = d[k]
|
|||
|
|
return newD
|
|||
|
|
|
|||
|
|
# ========================= 【qml接口】 =========================
|
|||
|
|
|
|||
|
|
def getLocalOptions(self):
|
|||
|
|
if self._apiKey:
|
|||
|
|
return getLocalOptions(self._apiKey)
|
|||
|
|
else:
|
|||
|
|
return {}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局 OCR任务管理器
|
|||
|
|
MissionOCR = __MissionOcrClass()
|