diff --git a/.gitignore b/.gitignore index 33755f5..19fd7f4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ .idea document_output +results diff --git a/README.md b/README.md index d96d108..e37f2b6 100644 --- a/README.md +++ b/README.md @@ -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 \ No newline at end of file +# 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. \ No newline at end of file diff --git a/analyzer.py b/analyzer.py index e454aa7..ea8ccdd 100644 --- a/analyzer.py +++ b/analyzer.py @@ -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() diff --git a/visualizer.py b/visualizer.py index d987dd2..f4f1ff2 100644 --- a/visualizer.py +++ b/visualizer.py @@ -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'' +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')