reformat

src/plugins/pdf2txt.py

@@ -13,26 +13,7 @@ def pdf2txt(pdf_path, return_text=False):
 
     table_engine = PPStructure(recovery=True, lang='ch')
 
-    imgs = []
-    img_dir = os.path.join(output_dir, 'imgs')
-    if not os.path.exists(img_dir):
-        os.makedirs(img_dir)
-        pdfDoc = fitz.open(pdf_path)
-        totalPage = pdfDoc.page_count
-        for pg in tqdm(range(totalPage), desc='to imgs', ncols=100):
-            page = pdfDoc[pg]
-            rotate = int(0)
-            zoom_x = 2
-            zoom_y = 2
-            mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
-            pix = page.get_pixmap(matrix=mat, alpha=False)
-            img_filename = os.path.join(img_dir, f'images_{pg+1}.png')
-            pix.save(img_filename)  # os.sep
-            imgs.append(img_filename)
-    else:
-        img_names = sorted(os.listdir(img_dir))
-        imgs = [os.path.join(img_dir, img_name) for img_name in img_names]
-
+    imgs = convert_imgs(pdf_path, output_dir)
     respath = os.path.join(output_dir, 'res.txt')
 
     text = []
@@ -59,7 +40,7 @@ def pdf2txt(pdf_path, return_text=False):
                     continue  # 如果不是字典或者缺少 'text' 键，跳过当前循环
             text.append('\n')
 
-    whole_text = '\n'.join(text)
+    whole_text = ''.join(text)
     with open(respath, 'w', encoding='utf-8') as f:
         f.write(whole_text)
 
@@ -68,6 +49,29 @@ def pdf2txt(pdf_path, return_text=False):
 
     return respath
 
+def convert_imgs(pdf_path, output_dir):
+    imgs = []
+    img_dir = os.path.join(output_dir, 'imgs')
+    if not os.path.exists(img_dir):
+        os.makedirs(img_dir)
+        pdfDoc = fitz.open(pdf_path)
+        totalPage = pdfDoc.page_count
+        for pg in tqdm(range(totalPage), desc='to imgs', ncols=100):
+            page = pdfDoc[pg]
+            rotate = int(0)
+            zoom_x = 2
+            zoom_y = 2
+            mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
+            pix = page.get_pixmap(matrix=mat, alpha=False)
+            img_filename = os.path.join(img_dir, f'images_{pg+1}.png')
+            pix.save(img_filename)  # os.sep
+            imgs.append(img_filename)
+    else:
+        img_names = sorted(os.listdir(img_dir))
+        imgs = [os.path.join(img_dir, img_name) for img_name in img_names]
+
+    return imgs
+
 if __name__ == "__main__":
-    pdf_path = r'data/file/焙烤食品工艺学.pdf'
+    pdf_path = r'saves/data/uploads/2e04d5_保健食品.pdf'
     print(pdf2txt(pdf_path))