ysif9/compare-document-processing
0
1import difflib2import tempfile3import time4from io import BytesIO5from pathlib import Path6 7import streamlit as st8from docling.datamodel.base_models import DocumentStream, InputFormat9from docling.document_converter import DocumentConverter, PdfFormatOption, ImageFormatOption10from docling.datamodel.pipeline_options import PdfPipelineOptions, EasyOcrOptions, TesseractOcrOptions11from marker.converters.pdf import PdfConverter12from marker.models import create_model_dict13from marker.output import text_from_rendered14from st_diff_viewer import diff_viewer15 16import fitz17 18@st.cache_resource19def load_marker_models() -> dict:20 """Load Marker models"""21 return create_model_dict()22 23@st.cache_data(show_spinner=False)24def extract_with_marker(pdf_bytes: bytes):25 """Extract text from PDF using Marker"""26 27 try:28 # Save bytes to temporary file since Marker needs a file path29 with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:30 tmp_file.write(pdf_bytes)31 tmp_file_path = tmp_file.name32 33 # Initialize Marker converter34 converter = PdfConverter(35 artifact_dict=load_marker_models(),36 )37 38 start_time = time.time()39 rendered = converter(tmp_file_path)40 text, _, images = text_from_rendered(rendered)41 end_time = time.time()42 43 # Clean up temp file44 Path(tmp_file_path).unlink()45 46 processing_time = end_time - start_time47 48 return text, processing_time, None49 50 except Exception as e:51 return None, None, str(e)52 53 54def pdf_to_images(pdf_bytes: bytes, dpi: int = 200) -> list[bytes]:55 """Convert PDF pages to PIL Images using PyMuPDF"""56 images = []57 pdf_doc = fitz.open(stream=pdf_bytes, filetype="pdf")58 59 zoom = float(dpi) / 72.060 mat = fitz.Matrix(zoom, zoom)61 62 try:63 for page in pdf_doc:64 pix = page.get_pixmap(matrix=mat)65 66 img_data = pix.tobytes("png")67 # img = Image.open(BytesIO(img_data))68 images.append(img_data)69 70 finally:71 pdf_doc.close()72 73 return images74 75@st.cache_data(show_spinner=False)76def extract_with_docling(pdf_bytes: bytes, filename: str, ocr_engine: str = "EasyOCR", full_ocr_mode: bool = False):77 """Extract text from PDF using Docling with configurable OCR options78 79 Args:80 pdf_bytes: PDF file content as bytes81 filename: Name of the PDF file82 ocr_engine: OCR engine to use ("EasyOCR" or "Tesseract")83 full_ocr_mode: If True, converts pages to images and applies full OCR84 """85 86 try:87 if full_ocr_mode:88 # Convert PDF pages to images first89 images = pdf_to_images(pdf_bytes, dpi=300)90 91 pipeline_options = PdfPipelineOptions()92 pipeline_options.do_ocr = True93 if ocr_engine == "Tesseract":94 pipeline_options.ocr_options = TesseractOcrOptions(force_full_page_ocr=True)95 else:96 pipeline_options.ocr_options = EasyOcrOptions(force_full_page_ocr=True)97 98 # Initialize converter for images99 converter = DocumentConverter(100 format_options={101 InputFormat.IMAGE: ImageFormatOption(102 pipeline_options=pipeline_options103 )104 }105 )106 107 all_markdown = []108 total_processing_time = 0.0109 for i, img in enumerate(images):110 # img_buffer = BytesIO()111 # img.save(img_buffer, format='PNG')112 img_bytes = BytesIO(img)113 114 # Create DocumentStream for the image115 img_stream = DocumentStream(116 name=f"{filename}_page_{i+1}.png",117 stream=img_bytes118 )119 120 # Convert image with OCR121 start_time = time.time()122 result = converter.convert(img_stream)123 end_time = time.time()124 processing_time = end_time - start_time125 total_processing_time += processing_time126 page_markdown = result.document.export_to_markdown()127 128 if page_markdown.strip():129 all_markdown.append(f"# Page {i+1}\n\n{page_markdown}")130 131 # Combine all pages132 markdown_text = "\n\n---\n\n".join(all_markdown)133 return markdown_text, total_processing_time, None134 135 else:136 # Standard PDF processing137 buf = BytesIO(pdf_bytes)138 source = DocumentStream(name=filename, stream=buf)139 140 # Configure pipeline options141 pipeline_options = PdfPipelineOptions()142 143 # Configure OCR engine144 if ocr_engine == "Tesseract":145 pipeline_options.ocr_options = TesseractOcrOptions()146 else:147 pipeline_options.ocr_options = EasyOcrOptions()148 149 # Initialize Docling converter with custom options150 converter = DocumentConverter(151 format_options={152 InputFormat.PDF: PdfFormatOption(153 pipeline_options=pipeline_options154 )155 }156 )157 158 start_time = time.time()159 result = converter.convert(source)160 end_time = time.time()161 markdown_text = result.document.export_to_markdown()162 processing_time = end_time - start_time163 return markdown_text, processing_time, None164 except Exception as e:165 return None, None, str(e)166 167 168def calculate_similarity(text1: str, text2: str) -> float:169 """Calculate similarity ratio between two texts"""170 return difflib.SequenceMatcher(None, text1, text2).ratio()171 172 173def main() -> None:174 """175 Main function for the application, providing an interface for comparing PDF-to-Markdown176 extraction performance between the Marker library and the Docling library. The function177 is executed in a Streamlit environment and utilizes its widgets and layout.178 179 This function handles file uploads, extraction using the two libraries, and displays180 various processing metrics, outputs, and comparisons to the user in an accessible format.181 182 :raises ValueError: If invalid or unsupported inputs are provided during processing.183 """184 st.set_page_config(185 page_title="PDF Extraction Comparison: Marker vs Docling",186 page_icon="📄",187 layout="wide"188 )189 190 st.title("📄 PDF Extraction Comparison: Marker vs Docling")191 st.markdown("Compare PDF-to-Markdown extraction performance between **Marker**, **Docling Standard** (PDF text extraction), and **Docling Full OCR** (page-to-image + OCR processing)")192 193 # File upload194 st.header("📤 Upload PDF Document")195 uploaded_file = st.file_uploader(196 "Choose a PDF file",197 type="pdf",198 help="Upload a PDF document to compare extraction performance"199 )200 201 # OCR Configuration Section202 st.header("⚙️ OCR Configuration")203 204 ocr_engine = st.selectbox(205 "OCR Engine",206 options=["EasyOCR", "Tesseract"],207 index=0,208 help="Choose the OCR engine for text extraction. EasyOCR is generally faster, while Tesseract may be more accurate for certain document types."209 )210 211 st.info("📋 **Processing modes**: The app will run both Docling Standard (PDF text extraction) and Docling Full OCR (page-to-image + OCR) modes for comparison.")212 213 if uploaded_file is not None:214 st.success(f"File uploaded: {uploaded_file.name}")215 pdf_bytes = uploaded_file.read()216 217 # Process with all three methods218 st.header("🔄 Processing...")219 220 # Create columns for parallel processing display221 col1, col2, col3 = st.columns(3)222 223 with col1:224 st.subheader("🏷️ Marker Processing")225 marker_placeholder = st.empty()226 227 with col2:228 st.subheader("📋 Docling Standard")229 docling_standard_placeholder = st.empty()230 231 with col3:232 st.subheader("🔍 Docling Full OCR")233 docling_ocr_placeholder = st.empty()234 235 # Process with Marker236 with marker_placeholder.container():237 with st.spinner("Processing with Marker..."):238 marker_text, marker_time, marker_error = extract_with_marker(pdf_bytes)239 240 # Process with Docling Standard Mode241 with docling_standard_placeholder.container():242 with st.spinner(f"Processing with Docling Standard ({ocr_engine} OCR)..."):243 docling_standard_text, docling_standard_time, docling_standard_error = extract_with_docling(244 pdf_bytes,245 uploaded_file.name,246 ocr_engine=ocr_engine,247 full_ocr_mode=False248 )249 250 # Process with Docling Full OCR Mode251 with docling_ocr_placeholder.container():252 with st.spinner(f"Processing with Docling Full OCR ({ocr_engine} OCR)..."):253 docling_ocr_text, docling_ocr_time, docling_ocr_error = extract_with_docling(254 pdf_bytes,255 uploaded_file.name,256 ocr_engine=ocr_engine,257 full_ocr_mode=True258 )259 260 # Display results261 st.header("📊 Results")262 263 # Performance metrics264 if marker_time is not None and docling_standard_time is not None and docling_ocr_time is not None:265 metrics_col1, metrics_col2, metrics_col3 = st.columns(3)266 267 with metrics_col1:268 st.metric(269 "Marker Processing Time",270 f"{marker_time:.2f}s"271 )272 273 with metrics_col2:274 st.metric(275 "Docling Standard Time",276 f"{docling_standard_time:.2f}s"277 )278 279 with metrics_col3:280 st.metric(281 "Docling Full OCR Time",282 f"{docling_ocr_time:.2f}s"283 )284 285 # Text comparison286 if marker_text is not None and docling_standard_text is not None and docling_ocr_text is not None:287 # Calculate similarities between all methods288 similarity_marker_standard = calculate_similarity(marker_text, docling_standard_text)289 similarity_marker_ocr = calculate_similarity(marker_text, docling_ocr_text)290 similarity_standard_ocr = calculate_similarity(docling_standard_text, docling_ocr_text)291 292 # Display similarity metrics293 st.subheader("📝 Text Similarity Comparison")294 sim_col1, sim_col2, sim_col3 = st.columns(3)295 296 with sim_col1:297 st.metric("Marker ↔ Docling Standard", f"{similarity_marker_standard:.1%}")298 with sim_col2:299 st.metric("Marker ↔ Docling Full OCR", f"{similarity_marker_ocr:.1%}")300 with sim_col3:301 st.metric("Docling Standard ↔ Full OCR", f"{similarity_standard_ocr:.1%}")302 303 # Length comparison304 len_col1, len_col2, len_col3 = st.columns(3)305 with len_col1:306 st.info(f"Marker output: {len(marker_text)} characters")307 with len_col2:308 st.info(f"Docling Standard: {len(docling_standard_text)} characters")309 with len_col3:310 st.info(f"Docling Full OCR: {len(docling_ocr_text)} characters")311 312 # Three-way comparison tabs313 st.subheader("📄 Markdown Output Comparison")314 315 tab1, tab2, tab3, tab4 = st.tabs(["Marker Output", "Docling Standard", "Docling Full OCR", "Diff View"])316 317 with tab1:318 st.markdown("### Marker Output")319 st.text_area(320 "Marker Markdown",321 marker_text,322 height=800,323 key="marker_output"324 )325 326 with tab2:327 st.markdown("### Docling Standard Output")328 st.text_area(329 "Docling Standard Markdown",330 docling_standard_text,331 height=800,332 key="docling_standard_output"333 )334 335 with tab3:336 st.markdown("### Docling Full OCR Output")337 st.text_area(338 "Docling Full OCR Markdown",339 docling_ocr_text,340 height=800,341 key="docling_ocr_output"342 )343 344 with tab4:345 st.markdown("### Text Differences")346 347 # Allow user to choose which comparison to view348 diff_option = st.selectbox(349 "Choose comparison:",350 ["Marker vs Docling Standard", "Marker vs Docling Full OCR", "Docling Standard vs Full OCR"]351 )352 353 try:354 if diff_option == "Marker vs Docling Standard":355 diff_viewer(356 old_text=marker_text,357 new_text=docling_standard_text,358 left_title="Marker",359 right_title="Docling Standard",360 )361 elif diff_option == "Marker vs Docling Full OCR":362 diff_viewer(363 old_text=marker_text,364 new_text=docling_ocr_text,365 left_title="Marker",366 right_title="Docling Full OCR",367 )368 else: # Docling Standard vs Full OCR369 diff_viewer(370 old_text=docling_standard_text,371 new_text=docling_ocr_text,372 left_title="Docling Standard",373 right_title="Docling Full OCR",374 )375 except ImportError as e:376 st.error(f"streamlit-diff-viewer not available: {e}")377 378 # Error handling379 if marker_error:380 st.error(f"Marker Error: {marker_error}")381 382 if docling_standard_error:383 st.error(f"Docling Standard Error: {docling_standard_error}")384 385 if docling_ocr_error:386 st.error(f"Docling Full OCR Error: {docling_ocr_error}")387 388 else:389 st.info("👆 Please upload a PDF file to begin comparison")390 391 392if __name__ == "__main__":393 main()394