diff --git a/visualizer.py b/visualizer.py index f4f1ff2..4826e91 100644 --- a/visualizer.py +++ b/visualizer.py @@ -1,20 +1,17 @@ #!/usr/bin/env python3 """ DocTags Zone Visualizer - Simple script to visualize zones identified in DocTags format. +PNG-only version: Creates debug images with rectangles around zones. Usage: - python doctags_visualizer.py --doctags output.doctags.txt --pdf document.pdf --page 8 --output visualization.html + python visualizer_png_only.py --doctags output.doctags.txt --pdf document.pdf --page 8 """ import argparse import os import re import sys -import base64 -from io import BytesIO -import tempfile from pathlib import Path -import xml.etree.ElementTree as ET from PIL import Image, ImageDraw import pdf2image @@ -33,19 +30,17 @@ def parse_arguments(): """Parse command line arguments.""" results_dir = ensure_results_folder() - parser = argparse.ArgumentParser(description='Visualize zones identified in DocTags format') + parser = argparse.ArgumentParser(description='Visualize zones identified in DocTags format as PNG images') parser.add_argument('--doctags', '-d', type=str, required=True, help='Path to DocTags file') parser.add_argument('--pdf', '-p', type=str, required=True, 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=str(results_dir / "visualization.html"), - help='Output HTML file path') + parser.add_argument('--output', '-o', type=str, default=None, + help='Output PNG file path (default: results/visualization_page_X.png)') parser.add_argument('--dpi', type=int, default=200, help='DPI for PDF rendering') - parser.add_argument('--show', '-s', action='store_true', - help='Open HTML in default browser') parser.add_argument('--page-count', action='store_true', help='Just count pages in the PDF and exit') parser.add_argument('--scale', type=float, default=1.0, @@ -174,298 +169,6 @@ def parse_doctags(doctags_path): return zones -def create_visualization(image, zones, pdf_path, page_num, total_pages, output_path): - """Create HTML visualization with the image and overlay zones.""" - # Convert image to base64 for embedding - buffered = BytesIO() - image.save(buffered, format="PNG") - img_base64 = base64.b64encode(buffered.getvalue()).decode() - - # Image dimensions - img_width, img_height = image.size - - # Define colors for different zone types - zone_colors = { - 'section_header_level_1': '#FF5722', # Orange - 'text': '#2196F3', # Blue - 'picture': '#4CAF50', # Green - 'table': '#9C27B0', # Purple - 'page_header': '#FFC107', # Amber - 'page_footer': '#795548', # Brown - 'default': '#607D8B' # Blue Grey - } - - # Create HTML with CSS for zones - html = f""" - - - - DocTags Zone Visualization - Page {page_num} - - - -
-

DocTags Zone Visualization - {Path(pdf_path).stem}

-
-
- - - - -
""" - - # Add legend items for each zone type - zone_types = set(zone['type'] for zone in zones) - for zone_type in zone_types: - color = zone_colors.get(zone_type, zone_colors['default']) - html += f""" -
-
- {zone_type} -
""" - - html += """ -
-
- - -
-
- -
- -""" - - # Add zone overlays - for i, zone in enumerate(zones): - zone_type = zone['type'] - color = zone_colors.get(zone_type, zone_colors['default']) - - width = zone['x2'] - zone['x1'] - height = zone['y2'] - zone['y1'] - - html += f""" -
-
- {zone_type} -
-
""" - - html += """ -
- -
-

Detected Zones:

-""" - - # Add zone details - for zone in zones: - zone_type = zone['type'] - color = zone_colors.get(zone_type, zone_colors['default']) - - html += f""" -
-
{zone_type}
-
Position: ({zone['x1']}, {zone['y1']}) - ({zone['x2']}, {zone['y2']})
-""" - if zone['content']: - html += f""" -
{zone['content']}
-""" - - html += """ -
-""" - - html += """ -
- - - - -""" - - # Save HTML file - with open(output_path, 'w', encoding='utf-8') as f: - f.write(html) - - print(f"Visualization saved to: {output_path}") - - return output_path - -def normalize_coordinates(zones, image_width, image_height, grid_size=500): - """ - Normalize coordinates from the DocTags grid (0-500) to actual image dimensions. - - Args: - zones: List of zone dictionaries with x1, y1, x2, y2 coordinates - image_width: Width of the PDF page image in pixels - image_height: Height of the PDF page image in pixels - grid_size: The grid size used in DocTags (default 500) - - Returns: - The same zones list with updated coordinates - """ - # Create a copy of the zones to avoid modifying the original - normalized_zones = [] - - for zone in zones: - # Clone the zone - new_zone = zone.copy() - - # Convert from grid coordinates to actual page dimensions - new_zone['x1'] = int(zone['x1'] * image_width / grid_size) - new_zone['y1'] = int(zone['y1'] * image_height / grid_size) - new_zone['x2'] = int(zone['x2'] * image_width / grid_size) - new_zone['y2'] = int(zone['y2'] * image_height / grid_size) - - normalized_zones.append(new_zone) - - return normalized_zones - def create_debug_image(image, zones, page_num, output_path): """Create a debug image with rectangles around zones.""" # Create a copy of the input image @@ -527,15 +230,45 @@ def create_debug_image(image, zones, page_num, output_path): return debug_img -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): +def normalize_coordinates(zones, image_width, image_height, grid_size=500): + """ + Normalize coordinates from the DocTags grid (0-500) to actual image dimensions. + + Args: + zones: List of zone dictionaries with x1, y1, x2, y2 coordinates + image_width: Width of the PDF page image in pixels + image_height: Height of the PDF page image in pixels + grid_size: The grid size used in DocTags (default 500) + + Returns: + The same zones list with updated coordinates + """ + # Create a copy of the zones to avoid modifying the original + normalized_zones = [] + + for zone in zones: + # Clone the zone + new_zone = zone.copy() + + # Convert from grid coordinates to actual page dimensions + new_zone['x1'] = int(zone['x1'] * image_width / grid_size) + new_zone['y1'] = int(zone['y1'] * image_height / grid_size) + new_zone['x2'] = int(zone['x2'] * image_width / grid_size) + new_zone['y2'] = int(zone['y2'] * image_height / grid_size) + + normalized_zones.append(new_zone) + + return normalized_zones + +def process_page(pdf_path, page_num, doctags_path, output_path, dpi=200, 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 = results_dir / output_name - debug_output = output_path.with_suffix('.debug.png') + # Generate output path if not provided + if output_path is None: + output_name = f"visualization_page_{page_num}.png" + output_path = results_dir / output_name # Load the page image try: @@ -605,23 +338,12 @@ def process_page(pdf_path, page_num, doctags_path, output_base, dpi=200, show=Fa print(f"Error parsing DocTags: {e}") return False - # Get total page count - total_pages = count_pdf_pages(pdf_path) - # Create debug image with zones - create_debug_image(image, zones, page_num, debug_output) - - # Create HTML visualization - output_html = create_visualization(image, zones, pdf_path, page_num, total_pages, output_path) - - # Open in browser if requested - if show: - import webbrowser - webbrowser.open(f"file:///{os.path.abspath(output_html)}") + create_debug_image(image, zones, page_num, output_path) return True -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): +def process_all_pages(pdf_path, doctags_path, output_base, dpi=200, 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() @@ -634,26 +356,18 @@ def process_all_pages(pdf_path, doctags_path, output_base, dpi=200, show_last=Fa print(f"Processing all {total_pages} pages of the PDF...") - last_html = None - # Process each page for page_num in range(1, total_pages + 1): print(f"\nProcessing page {page_num} of {total_pages}...") # Generate output paths for this page - output_name = f"{Path(output_base).stem}_page_{page_num}{Path(output_base).suffix}" - output_path = results_dir / output_name + if output_base is None: + output_path = results_dir / f"visualization_page_{page_num}.png" + else: + output_path = Path(output_base).with_stem(f"{Path(output_base).stem}_page_{page_num}") # Process the page - success = process_page(pdf_path, page_num, doctags_path, output_path, dpi, False, scale, scale_x, scale_y, adjust) - - if success: - last_html = output_path - - # Open the last successful page in browser if requested - if show_last and last_html: - import webbrowser - webbrowser.open(f"file:///{os.path.abspath(last_html)}") + process_page(pdf_path, page_num, doctags_path, output_path, dpi, scale, scale_x, scale_y, adjust) return True @@ -686,20 +400,22 @@ def main(): args.doctags, args.output, args.dpi, - args.show, args.scale, args.scale_x, args.scale_y, args.adjust ) else: + # Determine output path + results_dir = ensure_results_folder() + output_path = args.output if args.output else results_dir / f"visualization_page_{args.page}.png" + process_page( args.pdf, args.page, args.doctags, - args.output, + output_path, args.dpi, - args.show, args.scale, args.scale_x, args.scale_y, @@ -707,123 +423,4 @@ def main(): ) if __name__ == "__main__": - main() - -# -------------------------------------------------------- -# Helper script to generate DocTags from a PDF -# -------------------------------------------------------- - -def generate_doctags_from_pdf(input_pdf, page_num, output_doctags, prompt="Convert this page to docling."): - """ - Helper function to generate DocTags from a PDF page using the analyzer.py script. - This is a wrapper for the original script functionality. - - Args: - input_pdf: Path to the PDF file - page_num: Page number to process (starting from 1) - output_doctags: Path to save the generated DocTags - prompt: Prompt for the model (default: "Convert this page to docling.") - - 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 - - try: - print(f"Generating DocTags for page {page_num} of {input_pdf}...") - - # Construct command to run analyzer.py - cmd = [ - sys.executable, # Use the current Python interpreter - "analyzer.py", # Path to the script - "--image", str(input_pdf), - "--page", str(page_num), - "--prompt", prompt - ] - - # Run the command - result = subprocess.run( - cmd, - capture_output=True, - text=True, - check=False - ) - - # Check if the command was successful - if result.returncode != 0: - print(f"Error running analyzer.py: {result.stderr}") - return False - - # Extract DocTags from output - doctags_content = None - output_lines = result.stdout.split('\n') - in_doctags = False - doctags_lines = [] - - for line in output_lines: - if '' in line: - in_doctags = True - - if in_doctags: - doctags_lines.append(line) - - if '' in line: - in_doctags = False - - if not doctags_lines: - print("No DocTags found in the output") - return False - - # Save DocTags to file - with open(output_doctags, 'w', encoding='utf-8') as f: - f.write('\n'.join(doctags_lines)) - - print(f"DocTags saved to: {output_doctags}") - return True - - except Exception as e: - print(f"Error generating DocTags: {e}") - return False - -# 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=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') - - args = parser.parse_args() - - # Generate DocTags - success = generate_doctags_from_pdf( - args.pdf, - args.page, - args.output, - args.prompt - ) - - return 0 if success else 1 - -# Run the generator if called directly with --generate flag -if __name__ == "__main__" and '--generate' in sys.argv: - sys.argv.remove('--generate') - sys.exit(generate_doctags_cli()) \ No newline at end of file + main() \ No newline at end of file