Spaces:
Sleeping
Sleeping
| import fitz | |
| import pytesseract | |
| from PIL import Image | |
| import io | |
| import os | |
| import shutil | |
| import pandas as pd | |
| from docling.document_converter import DocumentConverter, PdfFormatOption | |
| from docling.datamodel.base_models import InputFormat | |
| from docling.datamodel.pipeline_options import PdfPipelineOptions | |
| import pandas as pd | |
| class RedactoParser: | |
| def __init__(self): | |
| pipeline_options = PdfPipelineOptions() | |
| pipeline_options.do_ocr = False | |
| pipeline_options.do_table_structure = True | |
| pipeline_options.generate_page_images = False | |
| self.converter = DocumentConverter( | |
| format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)} | |
| ) | |
| def is_scanned_pdf(self, pdf_path, sample_pages=3): | |
| doc = fitz.open(pdf_path) | |
| total_text_length = 0 | |
| pages_to_check = min(sample_pages, len(doc)) | |
| for i in range(pages_to_check): | |
| total_text_length += len(doc[i].get_text("text").strip()) | |
| doc.close() | |
| return total_text_length < 100 | |
| def make_searchable(self, input_pdf_path, output_pdf_path, dpi=300): | |
| print(f"🔍 OCRing Scanned PDF: {input_pdf_path}") | |
| src_doc = fitz.open(input_pdf_path) | |
| out_doc = fitz.open() | |
| for i, page in enumerate(src_doc): | |
| print(f" -> Processing page {i+1} / {len(src_doc)}...") | |
| pix = page.get_pixmap(dpi=dpi) | |
| img = Image.open(io.BytesIO(pix.tobytes("png"))) | |
| pdf_bytes = pytesseract.image_to_pdf_or_hocr(img, extension='pdf', lang='eng') | |
| page_pdf = fitz.open("pdf", pdf_bytes) | |
| out_doc.insert_pdf(page_pdf) | |
| page_pdf.close() | |
| out_doc.save(output_pdf_path, garbage=4, deflate=True) | |
| out_doc.close() | |
| src_doc.close() | |
| return output_pdf_path | |
| def extract_layout(self, input_path, working_path): | |
| if self.is_scanned_pdf(input_path): | |
| self.make_searchable(input_path, working_path) | |
| else: | |
| shutil.copy(input_path, working_path) | |
| print(f"📄 Extracting structure from: {working_path}") | |
| result = self.converter.convert(working_path) | |
| layout_data = [] | |
| current_h1, current_h2 = "Document Start", "" | |
| for item, level in result.document.iterate_items(): | |
| label = getattr(item, "label", None) | |
| if label in ['page_header', 'page_footer']: continue | |
| if label == 'section_header': | |
| if level == 1: current_h1, current_h2 = item.text, "" | |
| elif level == 2: current_h2 = item.text | |
| elif label in ['text', 'list_item', 'paragraph', 'table']: | |
| bbox = item.prov[0].bbox.as_tuple() if hasattr(item, 'prov') and item.prov else None | |
| page_no = item.prov[0].page_no if hasattr(item, 'prov') and item.prov else 1 | |
| text_content = item.export_to_markdown() if label == 'table' else item.text | |
| layout_data.append({ | |
| "context_h1": current_h1, | |
| "context_h2": current_h2, | |
| "text": text_content, | |
| "label": label, | |
| "page": page_no, | |
| "bbox": bbox | |
| }) | |
| return pd.DataFrame(layout_data) |