danvics/docling-studio
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Create a specific folder for results

Browse source
This commit is contained in:
pjmalandrino 2025-05-21 10:56:02 +02:00
parent 14c532b138
commit 3f91fa3e27
4 changed files with 71 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -1,2 +1,3 @@
.idea
document_output
results

10
README.md
View file

@ -1 +1,9 @@
python analyzer.py --image document.pdf --page 8 && python visualizer.py --doctags output.doctags.txt --pdf document.pdf --page 8 --adjust --show
# DocTags Analyzer and Visualizer
Simple command to process PDF pages with DocTags:
```bash
python analyzer.py --image document.pdf --page 8 && python visualizer.py --doctags results/output.doctags.txt --pdf document.pdf --page 8 --adjust --show
```
All output files will be automatically stored in the `results` folder.

33
analyzer.py
View file

@ -26,14 +26,25 @@ from docling_core.types.doc import ImageRefMode
from docling_core.types.doc.document import DocTagsDocument, DoclingDocument
def ensure_results_folder():
"""Create the results folder if it doesn't exist."""
results_dir = Path("results")
if not results_dir.exists():
results_dir.mkdir()
print(f"Created results directory: {results_dir}")
return results_dir
def parse_arguments():
"""Parse command line arguments."""
results_dir = ensure_results_folder()
parser = argparse.ArgumentParser(description='Convert an image or PDF to docling format')
parser.add_argument('--image', '-i', type=str, required=True,
help='Path to local image file, PDF file, or URL')
parser.add_argument('--prompt', '-p', type=str, default="Convert this page to docling.",
help='Prompt for the model')
parser.add_argument('--output', '-o', type=str, default="./output.html",
parser.add_argument('--output', '-o', type=str, default=str(results_dir / "output.html"),
help='Output file path')
parser.add_argument('--show', '-s', action='store_true',
help='Show output in browser')
@ -410,6 +421,9 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import stream_generate
# Ensure results folder exists
results_dir = ensure_results_folder()
# Prepare input
prompt = args.prompt
output_base = Path(args.output)
@ -419,7 +433,7 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
if Path(image_path).suffix.lower() == '.pdf' and page_num > 1:
# Get base filename without extension
base_name = output_base.stem
output_path = output_base.parent / f"{base_name}_page{page_num}{output_base.suffix}"
output_path = results_dir / f"{base_name}_page{page_num}{output_base.suffix}"
print(f"Processing page {page_num}, output will be saved to {output_path}")
@ -460,15 +474,15 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
os.unlink(temp_img_path)
print(f"Removed temporary image file")
# Save the raw DocTags to a txt file
doctags_path = output_path.with_suffix('.doctags.txt')
# Save the raw DocTags to a txt file in results folder
doctags_path = results_dir / f"{output_path.stem}.doctags.txt"
with open(doctags_path, 'w', encoding='utf-8') as f:
f.write(output)
print(f"Raw DocTags saved to: {doctags_path}")
# Save the tag analysis if in DocTags-only mode
if args.doctags_only:
tags_path = output_path.with_suffix('.tags.md')
tags_path = results_dir / f"{output_path.stem}.tags.md"
with open(tags_path, 'w', encoding='utf-8') as f:
f.write(tags_analysis)
print(f"DocTags analysis saved to: {tags_path}")
@ -477,14 +491,14 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
if not args.doctags_only:
# Create markdown document
md_content = create_markdown_document(image_path, output, args)
md_path = output_path.with_suffix('.md')
md_path = results_dir / f"{output_path.stem}.md"
with open(md_path, 'w', encoding='utf-8') as f:
f.write(md_content)
print(f"Markdown saved to: {md_path}")
# Create HTML document
html_content = create_html_document(image_path, output, pil_image, args)
html_path = output_path.with_suffix('.html')
html_path = results_dir / f"{output_path.stem}.html"
with open(html_path, 'w', encoding='utf-8') as f:
f.write(html_content)
print(f"HTML saved to: {html_path}")
@ -500,7 +514,7 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
# Save a copy of the processed image for reference if in debug mode
if args.debug:
img_debug_path = output_path.with_suffix('.debug.png')
img_debug_path = results_dir / f"{output_path.stem}.debug.png"
pil_image.save(img_debug_path)
print(f"Saved debug image to: {img_debug_path}")
@ -508,6 +522,9 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
def main():
# Ensure results folder exists
ensure_results_folder()
# Parse arguments
args = parse_arguments()

44
visualizer.py
View file

@ -9,7 +9,7 @@ Usage:
import argparse
import os
import re
import sys # Add missing import
import sys
import base64
from io import BytesIO
import tempfile
@ -21,8 +21,18 @@ import pdf2image
# Regular expression to extract location data
LOC_PATTERN = r'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>'
def ensure_results_folder():
"""Create the results folder if it doesn't exist."""
results_dir = Path("results")
if not results_dir.exists():
results_dir.mkdir()
print(f"Created results directory: {results_dir}")
return results_dir
def parse_arguments():
"""Parse command line arguments."""
results_dir = ensure_results_folder()
parser = argparse.ArgumentParser(description='Visualize zones identified in DocTags format')
parser.add_argument('--doctags', '-d', type=str, required=True,
help='Path to DocTags file')
@ -30,7 +40,7 @@ def parse_arguments():
help='Path to original PDF file')
parser.add_argument('--page', type=int, default=8,
help='Page number in PDF (starts at 1, default: 8)')
parser.add_argument('--output', '-o', type=str, default="visualization.html",
parser.add_argument('--output', '-o', type=str, default=str(results_dir / "visualization.html"),
help='Output HTML file path')
parser.add_argument('--dpi', type=int, default=200,
help='DPI for PDF rendering')
@ -517,14 +527,14 @@ def create_debug_image(image, zones, page_num, output_path):
return debug_img
# Improved auto-adjustment function for visualizer.py
# Replace this function in your visualizer.py script
def process_page(pdf_path, page_num, doctags_path, output_base, dpi=200, show=False, scale=1.0, scale_x=None, scale_y=None, adjust=True):
"""Process a single page of the PDF with visualization."""
# Ensure results folder exists
results_dir = ensure_results_folder()
# Generate output paths for this page
output_name = f"{Path(output_base).stem}_page_{page_num}{Path(output_base).suffix}"
output_path = Path(output_base).parent / output_name
output_path = results_dir / output_name
debug_output = output_path.with_suffix('.debug.png')
# Load the page image
@ -613,6 +623,9 @@ def process_page(pdf_path, page_num, doctags_path, output_base, dpi=200, show=Fa
def process_all_pages(pdf_path, doctags_path, output_base, dpi=200, show_last=False, scale=1.0, scale_x=None, scale_y=None, adjust=False):
"""Process all pages of the PDF and create visualizations."""
# Ensure results folder exists
results_dir = ensure_results_folder()
# Get total page count
total_pages = count_pdf_pages(pdf_path)
if total_pages == 0:
@ -629,7 +642,7 @@ def process_all_pages(pdf_path, doctags_path, output_base, dpi=200, show_last=Fa
# Generate output paths for this page
output_name = f"{Path(output_base).stem}_page_{page_num}{Path(output_base).suffix}"
output_path = Path(output_base).parent / output_name
output_path = results_dir / output_name
# Process the page
success = process_page(pdf_path, page_num, doctags_path, output_path, dpi, False, scale, scale_x, scale_y, adjust)
@ -645,6 +658,9 @@ def process_all_pages(pdf_path, doctags_path, output_base, dpi=200, show_last=Fa
return True
def main():
# Ensure results folder exists
ensure_results_folder()
# Parse arguments
args = parse_arguments()
@ -711,6 +727,16 @@ def generate_doctags_from_pdf(input_pdf, page_num, output_doctags, prompt="Conve
Returns:
bool: True if successful, False otherwise
"""
# Ensure results folder exists
results_dir = ensure_results_folder()
# Update output path to be in results folder
if not isinstance(output_doctags, Path):
output_doctags = Path(output_doctags)
if output_doctags.parent != results_dir:
output_doctags = results_dir / output_doctags.name
import subprocess
import sys
@ -773,12 +799,14 @@ def generate_doctags_from_pdf(input_pdf, page_num, output_doctags, prompt="Conve
# Command-line interface for the DocTags generator
def generate_doctags_cli():
"""Command-line interface for generating DocTags."""
results_dir = ensure_results_folder()
parser = argparse.ArgumentParser(description='Generate DocTags from a PDF page')
parser.add_argument('--pdf', type=str, required=True,
help='Path to PDF file')
parser.add_argument('--page', type=int, default=1,
help='Page number (starts at 1)')
parser.add_argument('--output', type=str, default='output.doctags.txt',
parser.add_argument('--output', type=str, default=str(results_dir / 'output.doctags.txt'),
help='Output DocTags file')
parser.add_argument('--prompt', type=str, default='Convert this page to docling.',
help='Prompt for the model')