diff --git a/src/plugins/pdf2txt.py b/src/plugins/pdf2txt.py index e823d021..5d2a72de 100644 --- a/src/plugins/pdf2txt.py +++ b/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))