Spaces:
Running
Running
| 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) | |