172 lines
7.6 KiB
Python
172 lines
7.6 KiB
Python
# 双层可搜索 searchable pdf
|
||
# https://github.com/pymupdf/PyMuPDF/discussions/2299
|
||
|
||
import os
|
||
import fitz # PyMuPDF
|
||
|
||
from umi_log import logger
|
||
from .output import Output
|
||
|
||
|
||
class OutputPdfLayered(Output):
|
||
def __init__(self, argd):
|
||
self.dir = argd["outputDir"] # 输出路径（文件夹）
|
||
self.originPath = argd["originPath"] # 原始文件路径
|
||
self.fileName = argd["outputFileName"] # 文件名
|
||
self.password = argd["password"] # 密码
|
||
self.outputPath = f"{self.dir}/{self.fileName}.layered.pdf" # 输出路径
|
||
self.pdf = None
|
||
self.existentPages = [] # 已处理的页数
|
||
self.isInsertFont = False # 是否有字体嵌入
|
||
self.opacity = 0 # 文本透明度为0
|
||
try:
|
||
self.font = fitz.Font("cjk") # 字体
|
||
except Exception as e:
|
||
raise Exception(f"Failed to load cjk font. {e}\n无法加载cjk字体。")
|
||
try:
|
||
self.pdf = self._getPDF(self.originPath) # 加载pymupdf对象
|
||
except Exception as e:
|
||
raise Exception(
|
||
f"Failed to load doc file. {e}\n无法加载文档。\n{self.originPath}"
|
||
)
|
||
|
||
# 获取pdf文档对象，或将其它类型的文档转为PDF对象
|
||
def _getPDF(self, path):
|
||
# https://github.com/pymupdf/PyMuPDF-Utilities/blob/master/examples/convert-document/convert.py
|
||
doc = fitz.open(path)
|
||
# 如果已加密，则尝试解密
|
||
if doc.is_encrypted and not doc.authenticate(self.password):
|
||
raise Exception(
|
||
f'The document is encrypted, and the password "{self.password}" is incorrect.\n文档已加密，输入密码不正确。'
|
||
)
|
||
if doc.is_pdf:
|
||
return doc
|
||
b = doc.convert_to_pdf() # 转换为PDF格式的二进制数据
|
||
pdf = fitz.open("pdf", b) # 创建PDF文档对象
|
||
try:
|
||
pdf.set_toc(doc.get_toc()) # 复制原始文档的目录
|
||
except Exception:
|
||
logger.warning("pdf.set_toc error", exc_info=True, stack_info=True)
|
||
# 复制原始文档的元数据（如作者、标题等）
|
||
meta = doc.metadata
|
||
if not meta["producer"]:
|
||
meta["producer"] = "Umi-OCR & PyMuPDF v" + fitz.VersionBind
|
||
if not meta["creator"]:
|
||
meta["creator"] = "Umi-OCR & PyMuPDF PDF converter"
|
||
pdf.set_metadata(meta)
|
||
# 复制原始文档的链接
|
||
for pinput in doc:
|
||
links = pinput.get_links()
|
||
pout = pdf[pinput.number]
|
||
for link in links:
|
||
if link["kind"] == fitz.LINK_NAMED: # 不处理 named links
|
||
continue
|
||
pout.insert_link(link) # 写入新文档
|
||
doc.close() # 释放原文档
|
||
return pdf
|
||
|
||
# 计算填满宽和高的一行字体大小
|
||
def _calculateFontSize(self, text, w, h):
|
||
if h > w: # 竖排转为横排计算
|
||
w, h = h, w
|
||
fontsize = round(h) # 字体大小初值，假设为行高
|
||
minSize = 5 # 大小下限
|
||
getLen = lambda text, s: self.font.text_length(text, fontsize=s)
|
||
while getLen(text, fontsize) > w and fontsize >= minSize:
|
||
fontsize -= 1 # 尝试减小字体，直到行宽刚好小于界限
|
||
while getLen(text, fontsize) < w:
|
||
fontsize += 1 # 尝试增大字体，直到行宽刚好超过界限
|
||
while getLen(text, fontsize) > w and fontsize >= minSize:
|
||
fontsize -= 0.1 # 再次减小字体，将精度提升到 0.1
|
||
return fontsize
|
||
|
||
def print(self, res): # 输出图片结果
|
||
if not self.pdf:
|
||
logger.error("self.pdf 未初始化。")
|
||
return
|
||
pno = res["page"] - 1 # 当前页数
|
||
self.existentPages.append(pno) # 记录已处理的页面
|
||
if not res["code"] == 100:
|
||
return # 忽略空白
|
||
|
||
page = self.pdf[pno] # 当前页对象
|
||
page.clean_contents() # 内容流清理、语法更正，减少错误
|
||
protation = page.rotation # 获取页面旋转角度
|
||
isInsertFont = False # 当前是否进行过字体注入
|
||
# 插入文本，用shape.insert_text（可编辑）或page.insert_text（不可编辑）
|
||
for tb in res["data"]:
|
||
if self.opacity == 0 and "from" in tb and tb["from"] == "text":
|
||
continue # 双层（透明文字）模式下，跳过直接提取的文本，只写入OCR文本
|
||
if not isInsertFont: # 页面插入字体
|
||
self.isInsertFont = isInsertFont = True
|
||
page.insert_font(fontname="cjk", fontbuffer=self.font.buffer)
|
||
text = tb["text"]
|
||
box = tb["box"]
|
||
x0, y0 = box[0]
|
||
x2, y2 = box[2]
|
||
w = x2 - x0
|
||
h = y2 - y0
|
||
fontsize = self._calculateFontSize(text, w, h)
|
||
# 插入点的 旋转后的坐标
|
||
point = fitz.Point(x0, y2) * page.derotation_matrix
|
||
page.insert_text(
|
||
point,
|
||
text,
|
||
fontsize,
|
||
fontname="cjk",
|
||
rotate=protation, # 文本角度设定
|
||
stroke_opacity=self.opacity, # 描边透明度
|
||
fill_opacity=self.opacity, # 填充（字体）透明度
|
||
)
|
||
|
||
def onEnd(self): # 结束时保存。
|
||
if not self.pdf:
|
||
return
|
||
# 删除未处理的页数
|
||
for i in range(len(self.pdf) - 1, -1, -1):
|
||
if i not in self.existentPages:
|
||
self.pdf.delete_page(i)
|
||
logger.info(f"保存{len(self.pdf)}页PDF：{self.outputPath}")
|
||
if self.isInsertFont: # 有任意页面嵌入字体，则构建字体子集
|
||
try: # 对于部分PDF，如用txt直接打印的，构建字体子集会失败。
|
||
self.pdf.subset_fonts() # 构建字体子集，减小文件大小。需要 fontTools 库
|
||
except Exception: # TODO: 失败原因？可能文件中实际并没有字体？
|
||
logger.error("构建字体子集失败。", exc_info=True, stack_info=True)
|
||
# 保存：压缩并进行3级垃圾回收。等同 ez_save
|
||
self.save(self.pdf, self.outputPath, deflate=True, garbage=3)
|
||
else:
|
||
# 无嵌入字体，则直接保存，不压缩
|
||
self.save(self.pdf, self.outputPath)
|
||
|
||
def save(self, pdf, path, **options): # 保存并关闭 pdf 对象
|
||
try:
|
||
# 尝试保存到指定路径
|
||
pdf.save(path, **options)
|
||
except Exception:
|
||
# 保存失败，尝试保存到 ".temp" 路径
|
||
tempPath = self.outputPath + ".temp"
|
||
logger.warning(f"保存PDF失败。 path: {path}", exc_info=True)
|
||
try:
|
||
pdf.save(tempPath, **options)
|
||
pdf.close()
|
||
except Exception as e1:
|
||
logger.error(
|
||
f"保存PDF到临时路径失败。 tempPath: {tempPath}", exc_info=True
|
||
)
|
||
raise Exception(f"[Error] Unable to save PDF to [{tempPath}]: {e1}")
|
||
# 已保存到 .temp 并 close 原对象，尝试替换文件
|
||
try:
|
||
if os.path.exists(path):
|
||
os.remove(path)
|
||
os.rename(tempPath, path)
|
||
except Exception as e2:
|
||
logger.warning(
|
||
f"保存PDF文件替换失败。保存到临时文件: {tempPath}", exc_info=True
|
||
)
|
||
|
||
raise Exception(
|
||
f"[Warning] Unable to save PDF: [{path}]. Exception: {e2}. Saved to temporary path: [{tempPath}]."
|
||
)
|
||
else: # 正常结束
|
||
pdf.close()
|