Merge pull request #2816 from songpb/master

修复PDF旋转的BUG (Issues #2792)

.gitignore

@@ -173,4 +173,8 @@ cython_debug/
 .pytest_cache
 .DS_Store
 
+# Test File
+test.py
+configs/*.py
+
 

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``` 版本。

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)
+                        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)