373 lines
16 KiB
Python
373 lines
16 KiB
Python
|
# ===============================================
|
|||
|
|
# =============== 文档 - 任务管理器 ===============
|
|||
|
|
# ===============================================
|
|||
|
|
|
|||
|
|
# API所有页数page 均为1开始
|
|||
|
|
|
|||
|
|
import fitz # PyMuPDF
|
|||
|
|
import time
|
|||
|
|
import math
|
|||
|
|
from PIL import Image
|
|||
|
|
from io import BytesIO
|
|||
|
|
|
|||
|
|
from umi_log import logger
|
|||
|
|
from .mission import Mission
|
|||
|
|
from .mission_ocr import MissionOCR
|
|||
|
|
from ..ocr.tbpu import getParser
|
|||
|
|
from ..ocr.tbpu import IgnoreArea
|
|||
|
|
from ..ocr.tbpu.parser_tools.paragraph_parse import word_separator # 上下句间隔符
|
|||
|
|
|
|||
|
|
MinSize = 1080 # 最小渲染分辨率
|
|||
|
|
|
|||
|
|
# 合法文件后缀
|
|||
|
|
DocSuf = [
|
|||
|
|
".pdf",
|
|||
|
|
".xps",
|
|||
|
|
".epub",
|
|||
|
|
".mobi",
|
|||
|
|
".fb2",
|
|||
|
|
".cbz",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FitzOpen:
|
|||
|
|
def __init__(self, path):
|
|||
|
|
self._path = path
|
|||
|
|
self._doc = None
|
|||
|
|
|
|||
|
|
def __enter__(self):
|
|||
|
|
self._doc = fitz.open(self._path)
|
|||
|
|
return self._doc
|
|||
|
|
|
|||
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|||
|
|
self._doc.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# https://pymupdf.readthedocs.io/en/latest/matrix.html#matrix
|
|||
|
|
# 从变换矩阵中提取角度值，返回0~359整数
|
|||
|
|
def transform_to_rotation(matrix):
|
|||
|
|
# [1, 0, 0, 1, 0, 0] -> [cos(deg), sin(deg), -sin(deg), cos(deg), 0, 0].
|
|||
|
|
a, b, c, d, _, _ = matrix
|
|||
|
|
# 处理缩放和反射
|
|||
|
|
scale = math.sqrt(a**2 + b**2)
|
|||
|
|
if scale < 1e-6:
|
|||
|
|
return 0
|
|||
|
|
# 归一化以消除缩放影响
|
|||
|
|
cos_theta = a / scale
|
|||
|
|
sin_theta = b / scale
|
|||
|
|
# 检查反射
|
|||
|
|
determinant = a * d - b * c
|
|||
|
|
if determinant < 0:
|
|||
|
|
# 反射情况，调整角度计算
|
|||
|
|
cos_theta = -cos_theta
|
|||
|
|
theta_rad = math.atan2(sin_theta, cos_theta)
|
|||
|
|
theta_deg = math.degrees(theta_rad)
|
|||
|
|
rounded_deg = round(theta_deg) % 360
|
|||
|
|
return rounded_deg
|
|||
|
|
|
|||
|
|
|
|||
|
|
class _MissionDocClass(Mission):
|
|||
|
|
def __init__(self):
|
|||
|
|
super().__init__()
|
|||
|
|
self._schedulingMode = "1234" # 调度方式：顺序
|
|||
|
|
self._minInterval = 0.1 # msnTask最短调用间隔
|
|||
|
|
self._lastCallTime = 0 # 上一次调用时间
|
|||
|
|
|
|||
|
|
# 添加一个文档任务
|
|||
|
|
# msnInfo: { 回调函数"onXX", 参数"argd":{"tbpu.xx", "ocr.xx"} }
|
|||
|
|
# msnPath: 单个文档路径
|
|||
|
|
# pageRange: 页数范围。可选： None 全部页 , [1,3] 页面范围（含开头结束）。
|
|||
|
|
# pageList: 指定多个页数。可选： [] 使用pageRange设置 , [1,2,3] 指定页数
|
|||
|
|
# password: 密码（非必填）
|
|||
|
|
def addMission(self, msnInfo, msnPath, pageRange=None, pageList=[], password=""):
|
|||
|
|
# =============== 加载文档，获取文档操作对象 ===============
|
|||
|
|
try:
|
|||
|
|
doc = fitz.open(msnPath)
|
|||
|
|
except Exception as e:
|
|||
|
|
return f"[Error] fitz.open error: {msnPath} {e}"
|
|||
|
|
if doc.is_encrypted and not doc.authenticate(password):
|
|||
|
|
if password:
|
|||
|
|
msg = f"[Error] Incorrect password. 文档已加密，密码错误。 [{password}]"
|
|||
|
|
else:
|
|||
|
|
msg = "[Error] Doc encrypted. 文档已加密，请提供密码。"
|
|||
|
|
return msg
|
|||
|
|
msnInfo["doc"] = doc
|
|||
|
|
msnInfo["path"] = msnPath
|
|||
|
|
# =============== 拦截 onEnd ===============
|
|||
|
|
msnInfo["sourceOnEnd"] = msnInfo["onEnd"] if "onEnd" in msnInfo else None
|
|||
|
|
msnInfo["onEnd"] = self._preOnEnd
|
|||
|
|
# =============== pageRange 页面范围 ===============
|
|||
|
|
page_count = doc.page_count
|
|||
|
|
if len(pageList) == 0:
|
|||
|
|
if isinstance(pageRange, (tuple, list)) and len(pageRange) == 2:
|
|||
|
|
a, b = pageRange[0], pageRange[1]
|
|||
|
|
if a < 0:
|
|||
|
|
a += page_count + 1
|
|||
|
|
if b < 0:
|
|||
|
|
b += page_count + 1
|
|||
|
|
if a < 1:
|
|||
|
|
return f"[Error] pageRange {pageRange} 范围起始不能小于1"
|
|||
|
|
if b > page_count:
|
|||
|
|
return f"[Error] pageRange {pageRange} 范围结束不能大于页数 {doc.page_count}"
|
|||
|
|
if a > b:
|
|||
|
|
return f"[Error] pageRange {pageRange} 范围错误"
|
|||
|
|
pageList = list(range(a - 1, b))
|
|||
|
|
else:
|
|||
|
|
pageList = list(range(0, page_count))
|
|||
|
|
# 检查页数列表合法性
|
|||
|
|
if len(pageList) == 0:
|
|||
|
|
return "[Error] 页数列表为空"
|
|||
|
|
for p in pageList:
|
|||
|
|
if not isinstance(p, int):
|
|||
|
|
return "[Error] 页数列表内容非整数"
|
|||
|
|
if not 0 <= p < page_count:
|
|||
|
|
return f"[Error] 页数列表超出 1~{page_count} 范围"
|
|||
|
|
msnInfo["pageList"] = pageList
|
|||
|
|
# =============== tbpu文本块后处理 msnInfo["tbpu"] ===============
|
|||
|
|
argd = msnInfo["argd"] # 参数
|
|||
|
|
msnInfo["tbpu"] = []
|
|||
|
|
msnInfo["ignoreArea"] = {}
|
|||
|
|
# 忽略区域
|
|||
|
|
if "tbpu.ignoreArea" in argd:
|
|||
|
|
iArea = argd["tbpu.ignoreArea"]
|
|||
|
|
if isinstance(iArea, list) and len(iArea) > 0:
|
|||
|
|
msnInfo["ignoreArea"]["obj"] = IgnoreArea(iArea)
|
|||
|
|
# 范围，负数转为倒数第x页
|
|||
|
|
igStart = argd.get("tbpu.ignoreRangeStart", 1)
|
|||
|
|
igEnd = argd.get("tbpu.ignoreRangeEnd", page_count)
|
|||
|
|
if igStart < 0:
|
|||
|
|
igStart += page_count + 1
|
|||
|
|
if igEnd < 0:
|
|||
|
|
igEnd += page_count + 1
|
|||
|
|
msnInfo["ignoreArea"]["start"] = igStart - 1 # -1是将起始1页转为起始0页
|
|||
|
|
msnInfo["ignoreArea"]["end"] = igEnd - 1
|
|||
|
|
logger.debug(f"忽略区域范围： {igStart} ~ {igEnd} 。")
|
|||
|
|
# 获取排版解析器对象
|
|||
|
|
if "tbpu.parser" in argd:
|
|||
|
|
msnInfo["tbpu"].append(getParser(argd["tbpu.parser"]))
|
|||
|
|
return self.addMissionList(msnInfo, pageList)
|
|||
|
|
|
|||
|
|
def msnTask(self, msnInfo, pno): # 执行msn。pno为当前页数
|
|||
|
|
doc = msnInfo["doc"] # 文档对象
|
|||
|
|
page = doc[pno] # 页面对象
|
|||
|
|
argd = msnInfo["argd"] # 参数
|
|||
|
|
extractionMode = argd["doc.extractionMode"] # OCR内容模式
|
|||
|
|
""" mixed - 混合OCR/拷贝文本
|
|||
|
|
fullPage - 整页强制OCR
|
|||
|
|
imageOnly - 仅OCR图片
|
|||
|
|
textOnly - 仅拷贝原有文本 """
|
|||
|
|
errMsg = "" # 本次任务流程的异常信息
|
|||
|
|
|
|||
|
|
# =============== 提取图片和原文本 ===============
|
|||
|
|
imgs = [] # 待OCR的图片列表
|
|||
|
|
tbs = [] # text box 文本块列表
|
|||
|
|
page_rotation = page.rotation # 获取页面的旋转角度
|
|||
|
|
if extractionMode == "fullPage": # 模式：整页强制OCR
|
|||
|
|
# 检查页面边长，如果低于阈值，则增加放大系数，以提高渲染清晰度
|
|||
|
|
rect = page.rect
|
|||
|
|
w, h = abs(rect[2] - rect[0]), abs(rect[3] - rect[1])
|
|||
|
|
m = min(w, h)
|
|||
|
|
if m < MinSize:
|
|||
|
|
zoom = MinSize / max(m, 1)
|
|||
|
|
matrix = fitz.Matrix(zoom, zoom)
|
|||
|
|
else:
|
|||
|
|
zoom = 1
|
|||
|
|
matrix = fitz.Identity
|
|||
|
|
p = page.get_pixmap(matrix=matrix)
|
|||
|
|
bytes = p.tobytes("png")
|
|||
|
|
scale = 1 / zoom
|
|||
|
|
imgs.append(
|
|||
|
|
{"bytes": bytes, "xy": (0, 0), "scale_w": scale, "scale_h": scale}
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
# 获取元素 https://pymupdf.readthedocs.io/en/latest/_images/img-textpage.png
|
|||
|
|
# https://pymupdf.readthedocs.io/en/latest/textpage.html#structure-of-dictionary-outputs
|
|||
|
|
# 确保越界图像能被采集 https://github.com/pymupdf/PyMuPDF/issues/3171
|
|||
|
|
p = page.get_text("dict", clip=fitz.INFINITE_RECT())
|
|||
|
|
for t in p["blocks"]: # 遍历区块（段落）
|
|||
|
|
# ========== 获取图片 ==========
|
|||
|
|
if t["type"] == 1 and (
|
|||
|
|
extractionMode == "imageOnly" or extractionMode == "mixed"
|
|||
|
|
):
|
|||
|
|
# 提取图片相对旋转角，加上页面旋转角，得到图片绝对旋转角
|
|||
|
|
transform = t["transform"]
|
|||
|
|
img_rotation = transform_to_rotation(transform)
|
|||
|
|
abs_rotation = round(page_rotation+img_rotation) % 360
|
|||
|
|
img_bytes = t["image"] # 图片字节
|
|||
|
|
bbox = t["bbox"] # 图片包围盒
|
|||
|
|
# 图片视觉大小、原始大小、缩放比例
|
|||
|
|
w1, h1 = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
|||
|
|
w2, h2 = t["width"], t["height"]
|
|||
|
|
# 特殊情况：图片宽高为0
|
|||
|
|
if w2 <= 0 or h2 <= 0:
|
|||
|
|
continue
|
|||
|
|
# 单独计算宽高的缩放比例
|
|||
|
|
scale_w = w1 / w2
|
|||
|
|
scale_h = h1 / h2
|
|||
|
|
# 如果图片有绝对旋转，则逆向旋转图片字节
|
|||
|
|
if page_rotation != 0 or img_rotation != 0:
|
|||
|
|
logger.debug(f"P{pno}-{len(imgs)} 旋转：页面{page_rotation}°，图片{img_rotation}°，绝对{abs_rotation}°")
|
|||
|
|
if abs_rotation != 0:
|
|||
|
|
try:
|
|||
|
|
with Image.open(BytesIO(img_bytes)) as pimg:
|
|||
|
|
# 记录原图格式
|
|||
|
|
format = pimg.format
|
|||
|
|
if not format:
|
|||
|
|
format = "PNG"
|
|||
|
|
# PDF的旋转是顺时针，需要逆时针旋转图片
|
|||
|
|
pimg = pimg.rotate(-abs_rotation, expand=True)
|
|||
|
|
# 将旋转后的图片转回bytes
|
|||
|
|
buffered = BytesIO()
|
|||
|
|
pimg.save(buffered, format=format)
|
|||
|
|
img_bytes = buffered.getvalue()
|
|||
|
|
except Exception:
|
|||
|
|
logger.error(
|
|||
|
|
"旋转文档图片异常。", exc_info=True, stack_info=True
|
|||
|
|
)
|
|||
|
|
# 记录图片
|
|||
|
|
imgs.append(
|
|||
|
|
{
|
|||
|
|
"bytes": img_bytes,
|
|||
|
|
"xy": (bbox[0], bbox[1]),
|
|||
|
|
"scale_w": scale_w,
|
|||
|
|
"scale_h": scale_h,
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
# ========== 获取文本块 ==========
|
|||
|
|
elif t["type"] == 0 and (
|
|||
|
|
extractionMode == "textOnly" or extractionMode == "mixed"
|
|||
|
|
):
|
|||
|
|
l = len(t["lines"]) - 1
|
|||
|
|
for index, line in enumerate(t["lines"]): # 遍历每一行
|
|||
|
|
# 拼接该行所有子文本块的内容
|
|||
|
|
text = ""
|
|||
|
|
for span in line["spans"]:
|
|||
|
|
text += span["text"]
|
|||
|
|
# 提取其他信息，组装为OCR文本块格式
|
|||
|
|
if text:
|
|||
|
|
# 获取该行的的包围盒
|
|||
|
|
b = line["bbox"]
|
|||
|
|
if page_rotation == 0: # 页面没有旋转，直接提取
|
|||
|
|
box = [
|
|||
|
|
[b[0], b[1]],
|
|||
|
|
[b[2], b[1]],
|
|||
|
|
[b[2], b[3]],
|
|||
|
|
[b[0], b[3]],
|
|||
|
|
]
|
|||
|
|
else: # 页面有旋转，默认文本行无相对旋转，则反向消除文本的绝对旋转
|
|||
|
|
# https://pymupdf.readthedocs.io/en/latest/page.html#Page.derotation_matrix
|
|||
|
|
rotation_matrix = page.rotation_matrix
|
|||
|
|
b01 = fitz.Point(b[0], b[1]) * rotation_matrix
|
|||
|
|
b23 = fitz.Point(b[2], b[3]) * rotation_matrix
|
|||
|
|
x0 = min(b01.x, b23.x)
|
|||
|
|
x1 = max(b01.x, b23.x)
|
|||
|
|
y0 = min(b01.y, b23.y)
|
|||
|
|
y1 = max(b01.y, b23.y)
|
|||
|
|
box = [
|
|||
|
|
[x0, y0],
|
|||
|
|
[x1, y0],
|
|||
|
|
[x1, y1],
|
|||
|
|
[x0, y1],
|
|||
|
|
]
|
|||
|
|
# 组装文本块
|
|||
|
|
tb = {
|
|||
|
|
"box": box,
|
|||
|
|
"text": text,
|
|||
|
|
"score": 1,
|
|||
|
|
"end": "\n" if index == l else "", # 结尾符
|
|||
|
|
"from": "text", # 来源：直接提取文本
|
|||
|
|
}
|
|||
|
|
tbs.append(tb)
|
|||
|
|
# 补充结尾符
|
|||
|
|
for i1 in range(len(tbs) - 1):
|
|||
|
|
if tbs[i1]["end"]: # 跳过已有结尾符的
|
|||
|
|
continue
|
|||
|
|
i2 = i1 + 1
|
|||
|
|
sep = word_separator(tbs[i1]["text"][-1], tbs[i2]["text"][0])
|
|||
|
|
tbs[i1]["end"] = sep
|
|||
|
|
|
|||
|
|
# =============== 调用OCR，将 imgs 的内容提取出来放入 tbs ===============
|
|||
|
|
if imgs:
|
|||
|
|
# 提取 "ocr." 开头的参数，组装OCR参数字典
|
|||
|
|
ocrArgd = {}
|
|||
|
|
for k in argd:
|
|||
|
|
if k.startswith("ocr."):
|
|||
|
|
ocrArgd[k] = argd[k]
|
|||
|
|
# 调用OCR，堵塞等待任务完成
|
|||
|
|
ocrList = MissionOCR.addMissionWait(ocrArgd, imgs)
|
|||
|
|
# 整理OCR结果
|
|||
|
|
for o in ocrList:
|
|||
|
|
res = o["result"]
|
|||
|
|
if res["code"] == 100:
|
|||
|
|
x, y = o["xy"]
|
|||
|
|
scale_w = o["scale_w"]
|
|||
|
|
scale_h = o["scale_h"]
|
|||
|
|
for r in res["data"]:
|
|||
|
|
# 将所有文本块的坐标，从图片相对坐标系，转为页面绝对坐标系
|
|||
|
|
for bi in range(4):
|
|||
|
|
r["box"][bi][0] = r["box"][bi][0] * scale_w + x
|
|||
|
|
r["box"][bi][1] = r["box"][bi][1] * scale_h + y
|
|||
|
|
r["from"] = "ocr" # 来源：OCR
|
|||
|
|
tbs.append(r)
|
|||
|
|
elif res["code"] != 101:
|
|||
|
|
errMsg += f'[Error] OCR code:{res["code"]} msg:{res["data"]}\n'
|
|||
|
|
|
|||
|
|
# =============== tbpu文本块后处理 ===============
|
|||
|
|
# 忽略区域
|
|||
|
|
if msnInfo["ignoreArea"] and tbs:
|
|||
|
|
# 检查范围
|
|||
|
|
igStart = msnInfo["ignoreArea"]["start"]
|
|||
|
|
igEnd = msnInfo["ignoreArea"]["end"]
|
|||
|
|
if pno >= igStart and pno <= igEnd:
|
|||
|
|
tbs = msnInfo["ignoreArea"]["obj"].run(tbs)
|
|||
|
|
# 其他tbpu
|
|||
|
|
if msnInfo["tbpu"] and tbs:
|
|||
|
|
for tbpu in msnInfo["tbpu"]:
|
|||
|
|
tbs = tbpu.run(tbs)
|
|||
|
|
|
|||
|
|
# =============== 组装结果字典 resDict ===============
|
|||
|
|
if errMsg:
|
|||
|
|
logger.error(f"文档识别异常。P{pno}, errMsg: {errMsg}")
|
|||
|
|
errMsg = f"[Error] Doc P{pno}\n" + errMsg
|
|||
|
|
|
|||
|
|
if tbs: # 有文本
|
|||
|
|
resDict = {"code": 100, "data": tbs}
|
|||
|
|
elif errMsg: # 无文本，有异常
|
|||
|
|
resDict = {"code": 102, "data": errMsg}
|
|||
|
|
else: # 无文本，无异常
|
|||
|
|
resDict = {"code": 101, "data": ""}
|
|||
|
|
|
|||
|
|
# ===== 仅提取文本时任务速度过快，频繁回调会导致UI卡死，因此故意引入延迟 =====
|
|||
|
|
currentTime = time.time()
|
|||
|
|
elapsedTime = currentTime - self._lastCallTime
|
|||
|
|
# 如果与上一次调用的时间差小于最短间隔，则睡至满足最短间隔
|
|||
|
|
if elapsedTime < self._minInterval:
|
|||
|
|
t = self._minInterval - elapsedTime
|
|||
|
|
time.sleep(t)
|
|||
|
|
self._lastCallTime = currentTime
|
|||
|
|
return resDict
|
|||
|
|
|
|||
|
|
# 获取一个文档的信息，如页数
|
|||
|
|
def getDocInfo(self, path):
|
|||
|
|
try:
|
|||
|
|
with FitzOpen(path) as doc:
|
|||
|
|
info = {
|
|||
|
|
"path": path,
|
|||
|
|
"page_count": doc.page_count,
|
|||
|
|
"is_encrypted": doc.is_encrypted,
|
|||
|
|
}
|
|||
|
|
return info
|
|||
|
|
except Exception as e:
|
|||
|
|
return {"path": path, "error": e}
|
|||
|
|
|
|||
|
|
# 结束前的处理
|
|||
|
|
def _preOnEnd(self, msnInfo, msg):
|
|||
|
|
# 先关闭文档对象，再触发原本的 onEnd ，防止新文档保存到原路径时的冲突
|
|||
|
|
msnInfo["doc"].close()
|
|||
|
|
if msnInfo["sourceOnEnd"]:
|
|||
|
|
msnInfo["sourceOnEnd"](msnInfo, msg)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局 DOC 任务管理器
|
|||
|
|
MissionDOC = _MissionDocClass()
|