DaVinciCode/doctra-document-parser
0
1"""2Doctra - Document Parser for Hugging Face Spaces3 4This is a Hugging Face Spaces deployment of the Doctra document parsing library.5It provides a comprehensive web interface for PDF parsing, table/chart extraction,6image restoration, and enhanced document processing.7"""8 9import os10import shutil11import tempfile12import re13import html as _html14import base6415import json16from pathlib import Path17from typing import Optional, Tuple, List, Dict, Any18 19import gradio as gr20import pandas as pd21 22# Mock google.genai to avoid import errors23import sys24from unittest.mock import MagicMock25 26# Create a mock google.genai module27mock_google_genai = MagicMock()28sys.modules['google.genai'] = mock_google_genai29sys.modules['google.genai.types'] = MagicMock()30 31# Now import Doctra components32try:33 from doctra.parsers.structured_pdf_parser import StructuredPDFParser34 from doctra.parsers.table_chart_extractor import ChartTablePDFParser35 from doctra.parsers.enhanced_pdf_parser import EnhancedPDFParser36 from doctra.ui.docres_wrapper import DocResUIWrapper37 from doctra.utils.pdf_io import render_pdf_to_images38except ImportError as e:39 print(f"Warning: Some Doctra components may not be available: {e}")40 # Create mock classes if imports fail41 StructuredPDFParser = None42 ChartTablePDFParser = None43 EnhancedPDFParser = None44 DocResUIWrapper = None45 render_pdf_to_images = None46 47 48# UI Theme and Styling Constants49THEME = gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")50 51CUSTOM_CSS = """52/* Full-width layout */53.gradio-container {max-width: 100% !important; padding-left: 24px; padding-right: 24px}54.container {max-width: 100% !important}55.app {max-width: 100% !important}56 57/* Header and helpers */58.header {margin-bottom: 8px}59.subtitle {color: var(--body-text-color-subdued)}60.card {border:1px solid var(--border-color); border-radius:12px; padding:8px}61.status-ok {color: var(--color-success)}62 63/* Scrollable gallery styling */64.scrollable-gallery {65 max-height: 600px !important;66 overflow-y: auto !important;67 border: 1px solid var(--border-color) !important;68 border-radius: 8px !important;69 padding: 8px !important;70}71 72/* Page content styling */73.page-content img {74 max-width: 100% !important;75 height: auto !important;76 display: block !important;77 margin: 10px auto !important;78 border: 1px solid #ddd !important;79 border-radius: 8px !important;80 box-shadow: 0 2px 4px rgba(0,0,0,0.1) !important;81}82 83.page-content {84 max-height: none !important;85 overflow: visible !important;86}87 88/* Table styling */89.page-content table.doc-table { 90 width: 100% !important; 91 border-collapse: collapse !important; 92 margin: 12px 0 !important; 93}94.page-content table.doc-table th,95.page-content table.doc-table td { 96 border: 1px solid #e5e7eb !important; 97 padding: 8px 10px !important; 98 text-align: left !important; 99}100.page-content table.doc-table thead th { 101 background: #f9fafb !important; 102 font-weight: 600 !important; 103}104.page-content table.doc-table tbody tr:nth-child(even) td { 105 background: #fafafa !important; 106}107 108/* Clickable image buttons */109.image-button {110 background: #0066cc !important;111 color: white !important;112 border: none !important;113 padding: 5px 10px !important;114 border-radius: 4px !important;115 cursor: pointer !important;116 margin: 2px !important;117 font-size: 14px !important;118}119 120.image-button:hover {121 background: #0052a3 !important;122}123"""124 125 126def gather_outputs(127 out_dir: Path, 128 allowed_kinds: Optional[List[str]] = None, 129 zip_filename: Optional[str] = None, 130 is_structured_parsing: bool = False131) -> Tuple[List[tuple[str, str]], List[str], str]:132 """133 Gather output files and create a ZIP archive for download.134 """135 gallery_items: List[tuple[str, str]] = []136 file_paths: List[str] = []137 138 if out_dir.exists():139 if is_structured_parsing:140 # For structured parsing, include all files141 for file_path in sorted(out_dir.rglob("*")):142 if file_path.is_file():143 file_paths.append(str(file_path))144 else:145 # For full parsing, include specific main files146 main_files = [147 "result.html",148 "result.md", 149 "tables.html",150 "tables.xlsx"151 ]152 153 for main_file in main_files:154 file_path = out_dir / main_file155 if file_path.exists():156 file_paths.append(str(file_path))157 158 # Include images based on allowed kinds159 if allowed_kinds:160 for kind in allowed_kinds:161 p = out_dir / kind162 if p.exists():163 for img in sorted(p.glob("*.png")):164 file_paths.append(str(img))165 166 images_dir = out_dir / "images" / kind167 if images_dir.exists():168 for img in sorted(images_dir.glob("*.jpg")):169 file_paths.append(str(img))170 else:171 # Include all images if no specific kinds specified172 for p in (out_dir / "charts").glob("*.png"):173 file_paths.append(str(p))174 for p in (out_dir / "tables").glob("*.png"):175 file_paths.append(str(p))176 for p in (out_dir / "images").rglob("*.jpg"):177 file_paths.append(str(p))178 179 # Include Excel files based on allowed kinds180 if allowed_kinds:181 if "charts" in allowed_kinds and "tables" in allowed_kinds:182 excel_files = ["parsed_tables_charts.xlsx"]183 elif "charts" in allowed_kinds:184 excel_files = ["parsed_charts.xlsx"]185 elif "tables" in allowed_kinds:186 excel_files = ["parsed_tables.xlsx"]187 else:188 excel_files = []189 190 for excel_file in excel_files:191 excel_path = out_dir / excel_file192 if excel_path.exists():193 file_paths.append(str(excel_path))194 195 # Build gallery items for image display196 kinds = allowed_kinds if allowed_kinds else ["tables", "charts", "figures"]197 for sub in kinds:198 p = out_dir / sub199 if p.exists():200 for img in sorted(p.glob("*.png")):201 gallery_items.append((str(img), f"{sub}: {img.name}"))202 203 images_dir = out_dir / "images" / sub204 if images_dir.exists():205 for img in sorted(images_dir.glob("*.jpg")):206 gallery_items.append((str(img), f"{sub}: {img.name}"))207 208 # Create ZIP archive209 tmp_zip_dir = Path(tempfile.mkdtemp(prefix="doctra_zip_"))210 211 if zip_filename:212 safe_filename = re.sub(r'[<>:"/\\|?*]', '_', zip_filename)213 zip_base = tmp_zip_dir / safe_filename214 else:215 zip_base = tmp_zip_dir / "doctra_outputs"216 217 filtered_dir = tmp_zip_dir / "filtered_outputs"218 shutil.copytree(out_dir, filtered_dir, ignore=shutil.ignore_patterns('~$*', '*.tmp', '*.temp'))219 220 zip_path = shutil.make_archive(str(zip_base), 'zip', root_dir=str(filtered_dir))221 222 return gallery_items, file_paths, zip_path223 224 225def validate_vlm_config(use_vlm: bool, vlm_api_key: str, vlm_provider: str = "gemini") -> Optional[str]:226 """227 Validate VLM configuration parameters.228 """229 if use_vlm and vlm_provider not in ["ollama"] and not vlm_api_key:230 return "❌ Error: VLM API key is required when using VLM (except for Ollama)"231 232 if use_vlm and vlm_api_key and vlm_provider not in ["ollama"]:233 # Basic API key validation234 if len(vlm_api_key.strip()) < 10:235 return "❌ Error: VLM API key appears to be too short or invalid"236 if vlm_api_key.strip().startswith('sk-') and len(vlm_api_key.strip()) < 20:237 return "❌ Error: OpenAI API key appears to be invalid (too short)"238 239 return None240 241 242def create_page_html_content(page_content: List[str], base_dir: Optional[Path] = None) -> str:243 """244 Convert page content lines to HTML with inline images and proper formatting.245 """246 processed_content = []247 paragraph_buffer = []248 249 def flush_paragraph():250 """Flush accumulated paragraph content to HTML"""251 nonlocal paragraph_buffer252 if paragraph_buffer:253 joined = '<br/>'.join(_html.escape(l) for l in paragraph_buffer)254 processed_content.append(f'<p>{joined}</p>')255 paragraph_buffer = []256 257 def is_markdown_table_header(s: str) -> bool:258 return '|' in s and ('---' in s or '—' in s)259 260 def render_markdown_table(lines: List[str]) -> str:261 rows = [l.strip().strip('|').split('|') for l in lines]262 rows = [[_html.escape(c.strip()) for c in r] for r in rows]263 if len(rows) < 2:264 return ""265 266 header = rows[0]267 body = rows[2:] if len(rows) > 2 else []268 thead = '<thead><tr>' + ''.join(f'<th>{c}</th>' for c in header) + '</tr></thead>'269 tbody = '<tbody>' + ''.join('<tr>' + ''.join(f'<td>{c}</td>' for c in r) + '</tr>' for r in body) + '</tbody>'270 return f'<table class="doc-table">{thead}{tbody}</table>'271 272 i = 0273 n = len(page_content)274 275 while i < n:276 raw_line = page_content[i]277 line = raw_line.rstrip('\r\n')278 stripped = line.strip()279 280 # Handle image references281 if stripped.startswith(':282 flush_paragraph()283 match = re.match(r'!\[([^\]]+)\]\(([^)]+)\)', stripped)284 if match and base_dir is not None:285 caption = match.group(1)286 rel_path = match.group(2).replace('\\\\', '/').replace('\\', '/').lstrip('/')287 abs_path = (base_dir / rel_path).resolve()288 try:289 with open(abs_path, 'rb') as f:290 b64 = base64.b64encode(f.read()).decode('ascii')291 processed_content.append(f'<figure><img src="data:image/jpeg;base64,{b64}" alt="{_html.escape(caption)}"/><figcaption>{_html.escape(caption)}</figcaption></figure>')292 except Exception as e:293 print(f"❌ Failed to embed image {rel_path}: {e}")294 processed_content.append(f'<div>{_html.escape(caption)} (image not found)</div>')295 else:296 processed_content.append(f'<div>{_html.escape(stripped)}</div>')297 i += 1298 continue299 300 # Handle markdown tables301 if (stripped.startswith('|') or stripped.count('|') >= 2) and i + 1 < n and is_markdown_table_header(page_content[i + 1]):302 flush_paragraph()303 table_block = [stripped]304 i += 1305 table_block.append(page_content[i].strip())306 i += 1307 while i < n:308 nxt = page_content[i].rstrip('\r\n')309 if nxt.strip() == '' or (not nxt.strip().startswith('|') and nxt.count('|') < 2):310 break311 table_block.append(nxt.strip())312 i += 1313 html_table = render_markdown_table(table_block)314 if html_table:315 processed_content.append(html_table)316 else:317 for tl in table_block:318 paragraph_buffer.append(tl)319 continue320 321 # Handle headers and content322 if stripped.startswith('## '):323 flush_paragraph()324 processed_content.append(f'<h3>{_html.escape(stripped[3:])}</h3>')325 elif stripped.startswith('# '):326 flush_paragraph()327 processed_content.append(f'<h2>{_html.escape(stripped[2:])}</h2>')328 elif stripped == '':329 flush_paragraph()330 processed_content.append('<br/>')331 else:332 paragraph_buffer.append(raw_line)333 i += 1334 335 flush_paragraph()336 return "\n".join(processed_content)337 338 339def run_full_parse(340 pdf_file: str,341 use_vlm: bool,342 vlm_provider: str,343 vlm_api_key: str,344 layout_model_name: str,345 dpi: int,346 min_score: float,347 ocr_lang: str,348 ocr_psm: int,349 ocr_oem: int,350 ocr_extra_config: str,351 box_separator: str,352) -> Tuple[str, Optional[str], List[tuple[str, str]], List[str], str]:353 """Run full PDF parsing with structured output."""354 if not pdf_file:355 return ("No file provided.", None, [], [], "")356 357 # Check if Doctra components are available358 if StructuredPDFParser is None:359 return ("❌ Error: Doctra library not properly installed. Please check the requirements.", None, [], [], "")360 361 # Validate VLM configuration362 vlm_error = validate_vlm_config(use_vlm, vlm_api_key, vlm_provider)363 if vlm_error:364 return (vlm_error, None, [], [], "")365 366 original_filename = Path(pdf_file).stem367 368 # Create temporary directory for processing369 tmp_dir = Path(tempfile.mkdtemp(prefix="doctra_"))370 input_pdf = tmp_dir / f"{original_filename}.pdf"371 shutil.copy2(pdf_file, input_pdf)372 373 # Initialize parser with configuration374 parser = StructuredPDFParser(375 use_vlm=use_vlm,376 vlm_provider=vlm_provider,377 vlm_api_key=vlm_api_key or None,378 layout_model_name=layout_model_name,379 dpi=int(dpi),380 min_score=float(min_score),381 ocr_lang=ocr_lang,382 ocr_psm=int(ocr_psm),383 ocr_oem=int(ocr_oem),384 ocr_extra_config=ocr_extra_config or "",385 box_separator=box_separator or "\n",386 )387 388 try:389 parser.parse(str(input_pdf))390 except Exception as e:391 import traceback392 traceback.print_exc()393 try:394 error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')395 return (f"❌ VLM processing failed: {error_msg}", None, [], [], "")396 except Exception:397 return (f"❌ VLM processing failed: <Unicode encoding error>", None, [], [], "")398 399 # Find output directory400 outputs_root = Path("outputs")401 out_dir = outputs_root / original_filename / "full_parse"402 if not out_dir.exists():403 candidates = sorted(outputs_root.glob("*/"), key=lambda p: p.stat().st_mtime, reverse=True)404 if candidates:405 out_dir = candidates[0] / "full_parse"406 else:407 out_dir = outputs_root408 409 # Read markdown file if it exists410 md_file = next(out_dir.glob("*.md"), None)411 md_preview = None412 if md_file and md_file.exists():413 try:414 with md_file.open("r", encoding="utf-8", errors="ignore") as f:415 md_preview = f.read()416 except Exception:417 md_preview = None418 419 # Gather output files and create ZIP420 gallery_items, file_paths, zip_path = gather_outputs(421 out_dir, 422 zip_filename=original_filename, 423 is_structured_parsing=False424 )425 426 return (427 f"✅ Parsing completed successfully!\n📁 Output directory: {out_dir}", 428 md_preview, 429 gallery_items, 430 file_paths, 431 zip_path432 )433 434 435def run_extract(436 pdf_file: str,437 target: str,438 use_vlm: bool,439 vlm_provider: str,440 vlm_api_key: str,441 layout_model_name: str,442 dpi: int,443 min_score: float,444) -> Tuple[str, str, List[tuple[str, str]], List[str], str]:445 """Run table/chart extraction from PDF."""446 if not pdf_file:447 return ("No file provided.", "", [], [], "")448 449 # Check if Doctra components are available450 if ChartTablePDFParser is None:451 return ("❌ Error: Doctra library not properly installed. Please check the requirements.", "", [], [], "")452 453 # Validate VLM configuration454 vlm_error = validate_vlm_config(use_vlm, vlm_api_key, vlm_provider)455 if vlm_error:456 return (vlm_error, "", [], [], "")457 458 original_filename = Path(pdf_file).stem459 460 # Create temporary directory for processing461 tmp_dir = Path(tempfile.mkdtemp(prefix="doctra_"))462 input_pdf = tmp_dir / f"{original_filename}.pdf"463 shutil.copy2(pdf_file, input_pdf)464 465 # Initialize parser with configuration466 parser = ChartTablePDFParser(467 extract_charts=(target in ("charts", "both")),468 extract_tables=(target in ("tables", "both")),469 use_vlm=use_vlm,470 vlm_provider=vlm_provider,471 vlm_api_key=vlm_api_key or None,472 layout_model_name=layout_model_name,473 dpi=int(dpi),474 min_score=float(min_score),475 )476 477 # Run extraction478 output_base = Path("outputs")479 parser.parse(str(input_pdf), str(output_base))480 481 # Find output directory482 outputs_root = output_base483 out_dir = outputs_root / original_filename / "structured_parsing"484 if not out_dir.exists():485 if outputs_root.exists():486 candidates = sorted(outputs_root.glob("*/"), key=lambda p: p.stat().st_mtime, reverse=True)487 if candidates:488 out_dir = candidates[0] / "structured_parsing"489 else:490 out_dir = outputs_root491 else:492 outputs_root.mkdir(parents=True, exist_ok=True)493 out_dir = outputs_root494 495 # Determine which kinds to include in outputs based on target selection496 allowed_kinds: Optional[List[str]] = None497 if target in ("tables", "charts"):498 allowed_kinds = [target]499 elif target == "both":500 allowed_kinds = ["tables", "charts"]501 502 # Gather output files and create ZIP503 gallery_items, file_paths, zip_path = gather_outputs(504 out_dir, 505 allowed_kinds, 506 zip_filename=original_filename, 507 is_structured_parsing=True508 )509 510 # Build tables HTML preview from Excel data (when VLM enabled)511 tables_html = ""512 try:513 if use_vlm:514 # Find Excel file based on target515 excel_filename = None516 if target in ("tables", "charts"):517 if target == "tables":518 excel_filename = "parsed_tables.xlsx"519 else: # charts520 excel_filename = "parsed_charts.xlsx"521 elif target == "both":522 excel_filename = "parsed_tables_charts.xlsx"523 524 if excel_filename:525 excel_path = out_dir / excel_filename526 if excel_path.exists():527 # Read Excel file and create HTML tables528 xl_file = pd.ExcelFile(excel_path)529 html_blocks = []530 531 for sheet_name in xl_file.sheet_names:532 df = pd.read_excel(excel_path, sheet_name=sheet_name)533 if not df.empty:534 # Create table with title535 title = f"<h3>{_html.escape(sheet_name)}</h3>"536 537 # Convert DataFrame to HTML table538 table_html = df.to_html(539 classes="doc-table",540 table_id=None,541 escape=True,542 index=False,543 na_rep=""544 )545 546 html_blocks.append(title + table_html)547 548 tables_html = "\n".join(html_blocks)549 except Exception as e:550 try:551 error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')552 print(f"Error building tables HTML: {error_msg}")553 except Exception:554 print(f"Error building tables HTML: <Unicode encoding error>")555 tables_html = ""556 557 return (558 f"✅ Parsing completed successfully!\n📁 Output directory: {out_dir}", 559 tables_html, 560 gallery_items, 561 file_paths, 562 zip_path563 )564 565 566def run_docres_restoration(567 pdf_file: str, 568 task: str, 569 device: str, 570 dpi: int, 571 save_enhanced: bool, 572 save_images: bool573) -> Tuple[str, Optional[str], Optional[str], Optional[dict], List[str]]:574 """Run DocRes image restoration on PDF."""575 if not pdf_file:576 return ("No file provided.", None, None, None, [])577 578 # Check if Doctra components are available579 if DocResUIWrapper is None:580 return ("❌ Error: Doctra library not properly installed. Please check the requirements.", None, None, None, [])581 582 try:583 # Initialize DocRes engine584 device_str = None if device == "auto" else device585 docres = DocResUIWrapper(device=device_str)586 587 # Extract filename588 original_filename = Path(pdf_file).stem589 590 # Create output directory591 output_dir = Path("outputs") / f"{original_filename}_docres"592 output_dir.mkdir(parents=True, exist_ok=True)593 594 # Run DocRes restoration595 enhanced_pdf_path = output_dir / f"{original_filename}_enhanced.pdf"596 docres.restore_pdf(597 pdf_path=pdf_file,598 output_path=str(enhanced_pdf_path),599 task=task,600 dpi=dpi601 )602 603 # Prepare outputs604 file_paths = []605 606 if save_enhanced and enhanced_pdf_path.exists():607 file_paths.append(str(enhanced_pdf_path))608 609 if save_images:610 # Look for enhanced images611 images_dir = output_dir / "enhanced_images"612 if images_dir.exists():613 for img_path in sorted(images_dir.glob("*.jpg")):614 file_paths.append(str(img_path))615 616 # Create metadata617 metadata = {618 "task": task,619 "device": str(docres.device),620 "dpi": dpi,621 "original_file": pdf_file,622 "enhanced_file": str(enhanced_pdf_path) if enhanced_pdf_path.exists() else None,623 "output_directory": str(output_dir)624 }625 626 status_msg = f"✅ DocRes restoration completed successfully!\n📁 Output directory: {output_dir}"627 628 enhanced_pdf_file = str(enhanced_pdf_path) if enhanced_pdf_path.exists() else None629 return (status_msg, pdf_file, enhanced_pdf_file, metadata, file_paths)630 631 except Exception as e:632 error_msg = f"❌ DocRes restoration failed: {str(e)}"633 return (error_msg, None, None, None, [])634 635 636def run_enhanced_parse(637 pdf_file: str,638 use_image_restoration: bool,639 restoration_task: str,640 restoration_device: str,641 restoration_dpi: int,642 use_vlm: bool,643 vlm_provider: str,644 vlm_api_key: str,645 layout_model_name: str,646 dpi: int,647 min_score: float,648 ocr_lang: str,649 ocr_psm: int,650 ocr_oem: int,651 ocr_extra_config: str,652 box_separator: str,653) -> Tuple[str, Optional[str], List[str], str, Optional[str], Optional[str], str]:654 """Run enhanced PDF parsing with DocRes image restoration."""655 if not pdf_file:656 return ("No file provided.", None, [], "", None, None, "")657 658 # Check if Doctra components are available659 if EnhancedPDFParser is None:660 return ("❌ Error: Doctra library not properly installed. Please check the requirements.", None, [], "", None, None, "")661 662 # Validate VLM configuration if VLM is enabled663 if use_vlm:664 vlm_error = validate_vlm_config(use_vlm, vlm_api_key, vlm_provider)665 if vlm_error:666 return (vlm_error, None, [], "", None, None, "")667 668 original_filename = Path(pdf_file).stem669 670 # Create temporary directory for processing671 tmp_dir = Path(tempfile.mkdtemp(prefix="doctra_enhanced_"))672 input_pdf = tmp_dir / f"{original_filename}.pdf"673 shutil.copy2(pdf_file, input_pdf)674 675 try:676 # Initialize enhanced parser with configuration677 parser = EnhancedPDFParser(678 use_image_restoration=use_image_restoration,679 restoration_task=restoration_task,680 restoration_device=restoration_device if restoration_device != "auto" else None,681 restoration_dpi=int(restoration_dpi),682 use_vlm=use_vlm,683 vlm_provider=vlm_provider,684 vlm_api_key=vlm_api_key or None,685 layout_model_name=layout_model_name,686 dpi=int(dpi),687 min_score=float(min_score),688 ocr_lang=ocr_lang,689 ocr_psm=int(ocr_psm),690 ocr_oem=int(ocr_oem),691 ocr_extra_config=ocr_extra_config or "",692 box_separator=box_separator or "\n",693 )694 695 # Parse the PDF with enhancement696 parser.parse(str(input_pdf))697 698 except Exception as e:699 import traceback700 traceback.print_exc()701 try:702 error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')703 return (f"❌ Enhanced parsing failed: {error_msg}", None, [], "", None, None, "")704 except Exception:705 return (f"❌ Enhanced parsing failed: <Unicode encoding error>", None, [], "", None, None, "")706 707 # Find output directory708 outputs_root = Path("outputs")709 out_dir = outputs_root / original_filename / "enhanced_parse"710 if not out_dir.exists():711 candidates = sorted(outputs_root.glob("*/"), key=lambda p: p.stat().st_mtime, reverse=True)712 if candidates:713 out_dir = candidates[0] / "enhanced_parse"714 else:715 out_dir = outputs_root716 717 # If still no enhanced_parse directory, try to find any directory with enhanced files718 if not out_dir.exists():719 for candidate_dir in outputs_root.rglob("*"):720 if candidate_dir.is_dir():721 enhanced_pdfs = list(candidate_dir.glob("*enhanced*.pdf"))722 if enhanced_pdfs:723 out_dir = candidate_dir724 break725 726 # Load first page content initially727 md_preview = None728 try:729 pages_dir = out_dir / "pages"730 first_page_path = pages_dir / "page_001.md"731 if first_page_path.exists():732 with first_page_path.open("r", encoding="utf-8", errors="ignore") as f:733 md_content = f.read()734 735 md_lines = md_content.split('\n')736 md_preview = create_page_html_content(md_lines, out_dir)737 else:738 md_file = next(out_dir.glob("*.md"), None)739 if md_file and md_file.exists():740 with md_file.open("r", encoding="utf-8", errors="ignore") as f:741 md_content = f.read()742 743 md_lines = md_content.split('\n')744 md_preview = create_page_html_content(md_lines, out_dir)745 except Exception as e:746 print(f"❌ Error loading initial content: {e}")747 md_preview = None748 749 # Gather output files and create ZIP750 _, file_paths, zip_path = gather_outputs(751 out_dir, 752 zip_filename=f"{original_filename}_enhanced", 753 is_structured_parsing=False754 )755 756 # Look for enhanced PDF file757 enhanced_pdf_path = None758 if use_image_restoration:759 enhanced_pdf_candidates = list(out_dir.glob("*enhanced*.pdf"))760 if enhanced_pdf_candidates:761 enhanced_pdf_path = str(enhanced_pdf_candidates[0])762 else:763 parent_enhanced = list(out_dir.parent.glob("*enhanced*.pdf"))764 if parent_enhanced:765 enhanced_pdf_path = str(parent_enhanced[0])766 767 return (768 f"✅ Enhanced parsing completed successfully!\n📁 Output directory: {out_dir}", 769 md_preview, 770 file_paths, 771 zip_path,772 pdf_file, # Original PDF path773 enhanced_pdf_path, # Enhanced PDF path774 str(out_dir) # Output directory for page-specific content775 )776 777 778def create_tips_markdown() -> str:779 """Create the tips section markdown for the UI."""780 return """781<div class="card">782 <b>Tips</b>783 <ul>784 <li>On Spaces, set a secret <code>VLM_API_KEY</code> to enable VLM features.</li>785 <li>Use <strong>Enhanced Parser</strong> for documents that need image restoration before parsing (scanned docs, low-quality PDFs).</li>786 <li>Use <strong>DocRes Image Restoration</strong> for standalone image enhancement without parsing.</li>787 <li>DocRes tasks: <code>appearance</code> (default), <code>dewarping</code>, <code>deshadowing</code>, <code>deblurring</code>, <code>binarization</code>, <code>end2end</code>.</li>788 <li>Outputs are saved under <code>outputs/<pdf_stem>/</code>.</li>789 <li><strong>Note:</strong> Google Gemini VLM may not be available due to dependency conflicts. Use OpenAI, Anthropic, or other VLM providers.</li>790 </ul>791</div>792 """793 794 795# Create the main Gradio interface796with gr.Blocks(title="Doctra - Document Parser", theme=THEME, css=CUSTOM_CSS) as demo:797 # Header section798 gr.Markdown(799 """800<div class="header">801 <h2 style="margin:0">Doctra — Document Parser</h2>802 <div class="subtitle">Parse PDFs, extract tables/charts, preview markdown, and download outputs.</div>803</div>804 """805 )806 807 # Full Parse Tab808 with gr.Tab("Full Parse"):809 with gr.Row():810 pdf = gr.File(file_types=[".pdf"], label="PDF")811 use_vlm = gr.Checkbox(label="Use VLM (optional)", value=False)812 vlm_provider = gr.Dropdown(["openai", "anthropic", "openrouter", "ollama"], value="openai", label="VLM Provider")813 vlm_api_key = gr.Textbox(type="password", label="VLM API Key", placeholder="Optional if VLM disabled")814 815 with gr.Accordion("Advanced", open=False):816 with gr.Row():817 layout_model = gr.Textbox(value="PP-DocLayout_plus-L", label="Layout model")818 dpi = gr.Slider(100, 400, value=200, step=10, label="DPI")819 min_score = gr.Slider(0, 1, value=0.0, step=0.05, label="Min layout score")820 with gr.Row():821 ocr_lang = gr.Textbox(value="eng", label="OCR Language")822 ocr_psm = gr.Slider(0, 13, value=4, step=1, label="Tesseract PSM")823 ocr_oem = gr.Slider(0, 3, value=3, step=1, label="Tesseract OEM")824 with gr.Row():825 ocr_config = gr.Textbox(value="", label="Extra OCR config")826 box_sep = gr.Textbox(value="\n", label="Box separator")827 828 run_btn = gr.Button("▶ Run Full Parse", variant="primary")829 status = gr.Textbox(label="Status", elem_classes=["status-ok"])830 831 # Full Parse components832 with gr.Row():833 with gr.Column():834 md_preview = gr.HTML(label="Extracted Content", visible=True, elem_classes=["page-content"])835 with gr.Column():836 page_image = gr.Image(label="Page image", interactive=False)837 files_out = gr.Files(label="Download individual output files")838 zip_out = gr.File(label="Download all outputs (ZIP)")839 840 run_btn.click(841 fn=run_full_parse,842 inputs=[pdf, use_vlm, vlm_provider, vlm_api_key, layout_model, dpi, min_score, ocr_lang, ocr_psm, ocr_oem, ocr_config, box_sep],843 outputs=[status, md_preview, files_out, zip_out],844 )845 846 # Tables & Charts Tab847 with gr.Tab("Extract Tables/Charts"):848 with gr.Row():849 pdf_e = gr.File(file_types=[".pdf"], label="PDF")850 target = gr.Dropdown(["tables", "charts", "both"], value="both", label="Target")851 use_vlm_e = gr.Checkbox(label="Use VLM (optional)", value=False)852 vlm_provider_e = gr.Dropdown(["openai", "anthropic", "openrouter", "ollama"], value="openai", label="VLM Provider")853 vlm_api_key_e = gr.Textbox(type="password", label="VLM API Key", placeholder="Optional if VLM disabled")854 855 with gr.Accordion("Advanced", open=False):856 with gr.Row():857 layout_model_e = gr.Textbox(value="PP-DocLayout_plus-L", label="Layout model")858 dpi_e = gr.Slider(100, 400, value=200, step=10, label="DPI")859 min_score_e = gr.Slider(0, 1, value=0.0, step=0.05, label="Min layout score")860 861 run_btn_e = gr.Button("▶ Run Extraction", variant="primary")862 status_e = gr.Textbox(label="Status")863 864 with gr.Row():865 with gr.Column():866 tables_preview_e = gr.HTML(label="Extracted Data", elem_classes=["page-content"])867 with gr.Column():868 image_e = gr.Image(label="Selected Image", interactive=False)869 870 files_out_e = gr.Files(label="Download individual output files")871 zip_out_e = gr.File(label="Download all outputs (ZIP)")872 873 run_btn_e.click(874 fn=lambda f, t, a, b, c, d, e, g: run_extract(875 f.name if f else "",876 t,877 a,878 b,879 c,880 d,881 e,882 g,883 ),884 inputs=[pdf_e, target, use_vlm_e, vlm_provider_e, vlm_api_key_e, layout_model_e, dpi_e, min_score_e],885 outputs=[status_e, tables_preview_e, files_out_e, zip_out_e],886 )887 888 # DocRes Image Restoration Tab889 with gr.Tab("DocRes Image Restoration"):890 with gr.Row():891 pdf_docres = gr.File(file_types=[".pdf"], label="PDF")892 docres_task_standalone = gr.Dropdown(893 ["appearance", "dewarping", "deshadowing", "deblurring", "binarization", "end2end"], 894 value="appearance", 895 label="Restoration Task"896 )897 docres_device_standalone = gr.Dropdown(898 ["auto", "cuda", "cpu"], 899 value="auto", 900 label="Device"901 )902 903 with gr.Row():904 docres_dpi = gr.Slider(100, 400, value=200, step=10, label="DPI")905 docres_save_enhanced = gr.Checkbox(label="Save Enhanced PDF", value=True)906 docres_save_images = gr.Checkbox(label="Save Enhanced Images", value=True)907 908 run_docres_btn = gr.Button("▶ Run DocRes Restoration", variant="primary")909 docres_status = gr.Textbox(label="Status", elem_classes=["status-ok"])910 911 with gr.Row():912 with gr.Column():913 gr.Markdown("### 📄 Original PDF")914 docres_original_pdf = gr.File(label="Original PDF File", interactive=False, visible=False)915 docres_original_page_image = gr.Image(label="Original PDF Page", interactive=False, height=800)916 with gr.Column():917 gr.Markdown("### ✨ Enhanced PDF")918 docres_enhanced_pdf = gr.File(label="Enhanced PDF File", interactive=False, visible=False)919 docres_enhanced_page_image = gr.Image(label="Enhanced PDF Page", interactive=False, height=800)920 921 docres_files_out = gr.Files(label="Download enhanced files")922 923 run_docres_btn.click(924 fn=run_docres_restoration,925 inputs=[pdf_docres, docres_task_standalone, docres_device_standalone, docres_dpi, docres_save_enhanced, docres_save_images],926 outputs=[docres_status, docres_original_pdf, docres_enhanced_pdf, docres_files_out]927 )928 929 # Enhanced Parser Tab930 with gr.Tab("Enhanced Parser"):931 with gr.Row():932 pdf_enhanced = gr.File(file_types=[".pdf"], label="PDF")933 use_image_restoration = gr.Checkbox(label="Use Image Restoration", value=True)934 restoration_task = gr.Dropdown(935 ["appearance", "dewarping", "deshadowing", "deblurring", "binarization", "end2end"], 936 value="appearance", 937 label="Restoration Task"938 )939 restoration_device = gr.Dropdown(940 ["auto", "cuda", "cpu"], 941 value="auto", 942 label="Restoration Device"943 )944 945 with gr.Row():946 use_vlm_enhanced = gr.Checkbox(label="Use VLM (optional)", value=False)947 vlm_provider_enhanced = gr.Dropdown(["openai", "anthropic", "openrouter", "ollama"], value="openai", label="VLM Provider")948 vlm_api_key_enhanced = gr.Textbox(type="password", label="VLM API Key", placeholder="Optional if VLM disabled")949 950 with gr.Accordion("Advanced Settings", open=False):951 with gr.Row():952 restoration_dpi = gr.Slider(100, 400, value=200, step=10, label="Restoration DPI")953 layout_model_enhanced = gr.Textbox(value="PP-DocLayout_plus-L", label="Layout model")954 dpi_enhanced = gr.Slider(100, 400, value=200, step=10, label="Processing DPI")955 min_score_enhanced = gr.Slider(0, 1, value=0.0, step=0.05, label="Min layout score")956 957 with gr.Row():958 ocr_lang_enhanced = gr.Textbox(value="eng", label="OCR Language")959 ocr_psm_enhanced = gr.Slider(0, 13, value=4, step=1, label="Tesseract PSM")960 ocr_oem_enhanced = gr.Slider(0, 3, value=3, step=1, label="Tesseract OEM")961 962 with gr.Row():963 ocr_config_enhanced = gr.Textbox(value="", label="Extra OCR config")964 box_sep_enhanced = gr.Textbox(value="\n", label="Box separator")965 966 run_enhanced_btn = gr.Button("▶ Run Enhanced Parse", variant="primary")967 enhanced_status = gr.Textbox(label="Status", elem_classes=["status-ok"])968 969 with gr.Row():970 with gr.Column():971 gr.Markdown("### 📄 Original PDF")972 enhanced_original_pdf = gr.File(label="Original PDF File", interactive=False, visible=False)973 enhanced_original_page_image = gr.Image(label="Original PDF Page", interactive=False, height=600)974 with gr.Column():975 gr.Markdown("### ✨ Enhanced PDF")976 enhanced_enhanced_pdf = gr.File(label="Enhanced PDF File", interactive=False, visible=False)977 enhanced_enhanced_page_image = gr.Image(label="Enhanced PDF Page", interactive=False, height=600)978 979 with gr.Row():980 enhanced_md_preview = gr.HTML(label="Extracted Content", visible=True, elem_classes=["page-content"])981 982 enhanced_files_out = gr.Files(label="Download individual output files")983 enhanced_zip_out = gr.File(label="Download all outputs (ZIP)")984 985 run_enhanced_btn.click(986 fn=run_enhanced_parse,987 inputs=[988 pdf_enhanced, use_image_restoration, restoration_task, restoration_device, restoration_dpi,989 use_vlm_enhanced, vlm_provider_enhanced, vlm_api_key_enhanced, layout_model_enhanced,990 dpi_enhanced, min_score_enhanced, ocr_lang_enhanced, ocr_psm_enhanced, ocr_oem_enhanced,991 ocr_config_enhanced, box_sep_enhanced992 ],993 outputs=[994 enhanced_status, enhanced_md_preview, enhanced_files_out, enhanced_zip_out,995 enhanced_original_pdf, enhanced_enhanced_pdf996 ]997 )998 999 # Tips section1000 gr.Markdown(create_tips_markdown())1001 1002 1003if __name__ == "__main__":1004 # Launch the interface1005 demo.launch(1006 server_name="0.0.0.0",1007 server_port=int(os.getenv("PORT", "7860")),1008 share=False1009 )1010