mirror of
https://github.com/RYDE-WORK/Langchain-Chatchat.git
synced 2026-01-19 13:23:16 +08:00
9c525b7fa5
新功能： - 优化 PDF 文件的 OCR，过滤无意义的小图片 by @liunux4odoo #2525 - 支持 Gemini 在线模型 by @yhfgyyf #2630 - 支持 GLM4 在线模型 by @zRzRzRzRzRzRzR - elasticsearch更新https连接 by @xldistance #2390 - 增强对PPT、DOC知识库文件的OCR识别 by @596192804 #2013 - 更新 Agent 对话功能 by @zRzRzRzRzRzRzR - 每次创建对象时从连接池获取连接，避免每次执行方法时都新建连接 by @Lijia0 #2480 - 实现 ChatOpenAI 判断token有没有超过模型的context上下文长度 by @glide-the - 更新运行数据库报错和项目里程碑 by @zRzRzRzRzRzRzR #2659 - 更新配置文件/文档/依赖 by @imClumsyPanda @zRzRzRzRzRzRzR - 添加日文版 readme by @eltociear #2787 修复： - langchain 更新后，PGVector 向量库连接错误 by @HALIndex #2591 - Minimax's model worker 错误 by @xyhshen - ES库无法向量检索.添加mappings创建向量索引 by MSZheng20 #2688
52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
from typing import List
|
||
from langchain.document_loaders.unstructured import UnstructuredFileLoader
|
||
from configs import PDF_OCR_THRESHOLD
|
||
from document_loaders.ocr import get_ocr
|
||
import tqdm
|
||
|
||
|
||
class RapidOCRPDFLoader(UnstructuredFileLoader):
|
||
def _get_elements(self) -> List:
|
||
def pdf2text(filepath):
|
||
import fitz # pyMuPDF里面的fitz包，不要与pip install fitz混淆
|
||
import numpy as np
|
||
ocr = get_ocr()
|
||
doc = fitz.open(filepath)
|
||
resp = ""
|
||
|
||
b_unit = tqdm.tqdm(total=doc.page_count, desc="RapidOCRPDFLoader context page index: 0")
|
||
for i, page in enumerate(doc):
|
||
b_unit.set_description("RapidOCRPDFLoader context page index: {}".format(i))
|
||
b_unit.refresh()
|
||
text = page.get_text("")
|
||
resp += text + "\n"
|
||
|
||
img_list = page.get_image_info(xrefs=True)
|
||
for img in img_list:
|
||
if xref := img.get("xref"):
|
||
bbox = img["bbox"]
|
||
# 检查图片尺寸是否超过设定的阈值
|
||
if ((bbox[2] - bbox[0]) / (page.rect.width) < PDF_OCR_THRESHOLD[0]
|
||
or (bbox[3] - bbox[1]) / (page.rect.height) < PDF_OCR_THRESHOLD[1]):
|
||
continue
|
||
pix = fitz.Pixmap(doc, xref)
|
||
img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, -1)
|
||
result, _ = ocr(img_array)
|
||
if result:
|
||
ocr_result = [line[1] for line in result]
|
||
resp += "\n".join(ocr_result)
|
||
|
||
# 更新进度
|
||
b_unit.update(1)
|
||
return resp
|
||
|
||
text = pdf2text(self.file_path)
|
||
from unstructured.partition.text import partition_text
|
||
return partition_text(text=text, **self.unstructured_kwargs)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
loader = RapidOCRPDFLoader(file_path="../tests/samples/ocr_test.pdf")
|
||
docs = loader.load()
|
||
print(docs)
|