diff --git a/.gitignore b/.gitignore index 37aebdf6..848d0818 100644 --- a/.gitignore +++ b/.gitignore @@ -173,4 +173,8 @@ cython_debug/ .pytest_cache .DS_Store +# Test File +test.py +configs/*.py + diff --git a/README.md b/README.md index 5fb0bcba..78ec0691 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,7 @@ OpenAI GPT API 的调用,并将在后续持续扩充对各类模型及模型 A 🚩 本项目未涉及微调、训练过程,但可利用微调或训练对本项目效果进行优化。 🌐 [AutoDL 镜像](https://www.codewithgpu.com/i/chatchat-space/Langchain-Chatchat/Langchain-Chatchat) 中 `v14` + 版本所使用代码已更新至本项目 `v0.2.10` 版本。 🐳 [Docker 镜像](registry.cn-beijing.aliyuncs.com/chatchat/chatchat:0.2.7) 已经更新到 ```0.2.7``` 版本。 diff --git a/document_loaders/mypdfloader.py b/document_loaders/mypdfloader.py index faaf63dd..71dfd137 100644 --- a/document_loaders/mypdfloader.py +++ b/document_loaders/mypdfloader.py @@ -1,5 +1,8 @@ from typing import List from langchain.document_loaders.unstructured import UnstructuredFileLoader +import cv2 +from PIL import Image +import numpy as np from configs import PDF_OCR_THRESHOLD from document_loaders.ocr import get_ocr import tqdm @@ -7,6 +10,30 @@ import tqdm class RapidOCRPDFLoader(UnstructuredFileLoader): def _get_elements(self) -> List: + def rotate_img(img, angle): + ''' + img --image + angle --rotation angle + return--rotated img + ''' + + h, w = img.shape[:2] + rotate_center = (w/2, h/2) + #获取旋转矩阵 + # 参数1为旋转中心点; + # 参数2为旋转角度,正值-逆时针旋转;负值-顺时针旋转 + # 参数3为各向同性的比例因子,1.0原图,2.0变成原来的2倍,0.5变成原来的0.5倍 + M = cv2.getRotationMatrix2D(rotate_center, angle, 1.0) + #计算图像新边界 + new_w = int(h * np.abs(M[0, 1]) + w * np.abs(M[0, 0])) + new_h = int(h * np.abs(M[0, 0]) + w * np.abs(M[0, 1])) + #调整旋转矩阵以考虑平移 + M[0, 2] += (new_w - w) / 2 + M[1, 2] += (new_h - h) / 2 + + rotated_img = cv2.warpAffine(img, M, (new_w, new_h)) + return rotated_img + def pdf2text(filepath): import fitz # pyMuPDF里面的fitz包,不要与pip install fitz混淆 import numpy as np @@ -30,7 +57,16 @@ class RapidOCRPDFLoader(UnstructuredFileLoader): 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) + samples = pix.samples + if int(page.rotation)!=0: #如果Page有旋转角度,则旋转图片 + img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, -1) + tmp_img = Image.fromarray(img_array); + ori_img = cv2.cvtColor(np.array(tmp_img),cv2.COLOR_RGB2BGR) + rot_img = rotate_img(img=ori_img, angle=360-page.rotation) + img_array = cv2.cvtColor(rot_img, cv2.COLOR_RGB2BGR) + else: + 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] @@ -46,6 +82,6 @@ class RapidOCRPDFLoader(UnstructuredFileLoader): if __name__ == "__main__": - loader = RapidOCRPDFLoader(file_path="../tests/samples/ocr_test.pdf") + loader = RapidOCRPDFLoader(file_path="/Users/tonysong/Desktop/test.pdf") docs = loader.load() print(docs) diff --git a/img/qr_code_86.jpg b/img/qr_code_86.jpg new file mode 100644 index 00000000..f353a04d Binary files /dev/null and b/img/qr_code_86.jpg differ diff --git a/img/qr_code_87.jpg b/img/qr_code_87.jpg new file mode 100644 index 00000000..ce8fc77a Binary files /dev/null and b/img/qr_code_87.jpg differ diff --git a/requirements.txt b/requirements.txt index c53eb563..7912ea1d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -64,4 +64,4 @@ streamlit-modal==0.1.0 streamlit-aggrid==0.3.4.post3 httpx==0.26.0 watchdog==3.0.0 -pyjwt==2.8.0 \ No newline at end of file +pyjwt==2.8.0 diff --git a/requirements_lite.txt b/requirements_lite.txt index 14566d15..aae3a284 100644 --- a/requirements_lite.txt +++ b/requirements_lite.txt @@ -32,4 +32,4 @@ watchdog~=3.0.0 # pymilvus>=2.3.4 # psycopg2==2.9.9 # pgvector>=0.2.4 -# chromadb==0.4.13 \ No newline at end of file +# chromadb==0.4.13