#!/usr/bin/env python3 """ DocTags Zone Visualizer - Simple script to visualize zones identified in DocTags format. Usage: python doctags_visualizer.py --doctags output.doctags.txt --pdf document.pdf --page 8 --output visualization.html """ import argparse import os import re import sys # Add missing import 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 # Regular expression to extract location data LOC_PATTERN = r'' def parse_arguments(): """Parse command line arguments.""" parser = argparse.ArgumentParser(description='Visualize zones identified in DocTags format') 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="visualization.html", help='Output HTML file path') 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, help='Scaling factor for zone coordinates (default: 1.0)') parser.add_argument('--scale-x', type=float, default=None, help='X-axis scaling factor (overrides --scale)') parser.add_argument('--scale-y', type=float, default=None, help='Y-axis scaling factor (overrides --scale)') parser.add_argument('--adjust', action='store_true', help='Try to automatically adjust scaling') return parser.parse_args() def count_pdf_pages(pdf_path): """Count the number of pages in a PDF file.""" if not os.path.exists(pdf_path): print(f"Error: PDF file not found: {pdf_path}") return 0 try: from pdf2image.pdf2image import pdfinfo_from_path info = pdfinfo_from_path(pdf_path, userpw=None, poppler_path=None) return info["Pages"] except Exception as e: print(f"Warning: pdfinfo failed: {e}") # Fallback method if pdfinfo fails try: images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=1, last_page=1) # Try to load the last page - increment until we get an error page_count = 1 while True: try: images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=page_count+1, last_page=page_count+1) if not images: break page_count += 1 except: break return page_count except Exception as e2: print(f"Error counting PDF pages: {e2}") return 0 def load_image_from_pdf(pdf_path, page_num=1, dpi=200): """Load a specific page from PDF as an image.""" if not os.path.exists(pdf_path): raise FileNotFoundError(f"PDF file not found: {pdf_path}") print(f"Converting PDF page {page_num} to image (DPI: {dpi})...") try: pdf_images = pdf2image.convert_from_path( pdf_path, dpi=dpi, first_page=page_num, last_page=page_num ) if not pdf_images: raise Exception(f"Could not extract page {page_num} from PDF") return pdf_images[0] # Return the requested page except Exception as e: raise Exception(f"Error converting PDF to image: {e}") def parse_doctags(doctags_path): """Parse DocTags file and extract zones with their coordinates.""" if not os.path.exists(doctags_path): raise FileNotFoundError(f"DocTags file not found: {doctags_path}") with open(doctags_path, 'r', encoding='utf-8') as f: doctags_content = f.read() # Extract content between tags doctag_pattern = r'(.*?)' doctag_match = re.search(doctag_pattern, doctags_content, re.DOTALL) if not doctag_match: raise ValueError("No tags found in the file") doctag_content = doctag_match.group(1) # Find all tags with location information zones = [] # Find all tag starts tag_starts = re.finditer(r'<(\w+)>', doctag_content) for tag_match in tag_starts: tag_name = tag_match.group(1) # Skip location tags themselves if tag_name.startswith('loc_'): continue # Find the end of the tag tag_start_pos = tag_match.start() tag_end_pattern = f'' tag_end_match = re.search(tag_end_pattern, doctag_content[tag_start_pos:]) if not tag_end_match: continue # Skip if no closing tag # Extract the tag content tag_content = doctag_content[tag_start_pos:tag_start_pos + tag_end_match.end()] # Look for location pattern loc_match = re.search(LOC_PATTERN, tag_content) if loc_match: # Extract coordinates x1, y1, x2, y2 = map(int, loc_match.groups()) # Extract text content if available text_content = "" # Look for content between the location info and the closing tag content_pattern = f'{LOC_PATTERN}(.*?)' content_match = re.search(content_pattern, tag_content, re.DOTALL) if content_match: text_content = content_match.group(5).strip() zones.append({ 'type': tag_name, 'x1': x1, 'y1': y1, 'x2': x2, 'y2': y2, 'content': text_content }) 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 create_debug_image(image, zones, page_num, output_path): """Create a debug image with rectangles around zones.""" # Create a copy of the input image debug_img = image.copy() draw = ImageDraw.Draw(debug_img) # Define colors for different zone types zone_colors = { 'section_header_level_1': (255, 87, 34), # Orange 'text': (33, 150, 243), # Blue 'picture': (76, 175, 80), # Green 'table': (156, 39, 176), # Purple 'page_header': (255, 193, 7), # Amber 'page_footer': (121, 85, 72), # Brown 'default': (96, 125, 139) # Blue Grey } # Draw rectangles for each zone for zone in zones: zone_type = zone['type'] color = zone_colors.get(zone_type, zone_colors['default']) draw.rectangle( [(zone['x1'], zone['y1']), (zone['x2'], zone['y2'])], outline=color, width=2 ) # Add zone type label label_width = len(zone_type) * 7 + 6 label_x = min(zone['x1'], image.width - label_width) # Keep label on image draw.rectangle( [(label_x, zone['y1']), (label_x + label_width, zone['y1'] + 20)], fill=(255, 255, 255, 180), outline=color ) draw.text( (label_x + 3, zone['y1'] + 3), zone_type, fill=color ) # Draw page number on the debug image draw.rectangle( [(10, 10), (100, 40)], fill=(0, 0, 0, 180), outline=(255, 255, 255) ) draw.text( (15, 15), f"Page {page_num}", fill=(255, 255, 255) ) # Save the debug image debug_img.save(output_path) print(f"Debug image saved to: {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.""" # 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 debug_output = output_path.with_suffix('.debug.png') # Load the page image try: image = load_image_from_pdf(pdf_path, page_num, dpi) print(f"Page {page_num} loaded: {image.size}") except Exception as e: print(f"Error loading page {page_num}: {e}") return False # Parse DocTags try: zones = parse_doctags(doctags_path) print(f"Found {len(zones)} zones in DocTags") # Debug output to understand scaling issues if zones: max_x = max([zone['x2'] for zone in zones]) max_y = max([zone['y2'] for zone in zones]) min_x = min([zone['x1'] for zone in zones]) min_y = min([zone['y1'] for zone in zones]) img_width, img_height = image.size print(f"Image dimensions: {img_width}x{img_height}") print(f"Zone boundaries: X({min_x}-{max_x}), Y({min_y}-{max_y})") print(f"Suggested scale factors: X={img_width/max_x:.3f}, Y={img_height/max_y:.3f}") # Always apply automatic scaling (making adjust=True by default) # This solves the common scaling issue between DocTags coordinates and image dimensions if adjust: width, height = image.size # Calculate appropriate scaling factors with better heuristics # Use smaller scaling to avoid cutting off content if max_x > 0: x_scale = min(width / max_x, 1.0) if max_x > width else max(width / max_x, 0.5) print(f"Auto-adjusted X scale to {x_scale:.3f} (image width: {width}, max zone x: {max_x})") else: x_scale = 1.0 if max_y > 0: y_scale = min(height / max_y, 1.0) if max_y > height else max(height / max_y, 0.5) print(f"Auto-adjusted Y scale to {y_scale:.3f} (image height: {height}, max zone y: {max_y})") else: y_scale = 1.0 # Apply more aggressive adjustment if image and zones are very different in scale if max_x > width * 5 or max_x < width / 5: x_scale = width / max_x print(f"Major X scale adjustment to {x_scale:.3f}") if max_y > height * 5 or max_y < height / 5: y_scale = height / max_y print(f"Major Y scale adjustment to {y_scale:.3f}") # Apply the scaling to all zones if x_scale != 1.0 or y_scale != 1.0: for zone in zones: zone['x1'] = int(zone['x1'] * x_scale) zone['y1'] = int(zone['y1'] * y_scale) zone['x2'] = int(zone['x2'] * x_scale) zone['y2'] = int(zone['y2'] * y_scale) print(f"Applied auto-scaling: X={x_scale}, Y={y_scale}") except Exception as e: 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)}") 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): """Process all pages of the PDF and create visualizations.""" # Get total page count total_pages = count_pdf_pages(pdf_path) if total_pages == 0: print("Error: Could not determine the number of pages in the PDF.") return False 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 = Path(output_base).parent / output_name # 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)}") return True def main(): # Parse arguments args = parse_arguments() # If just counting pages if args.page_count: page_count = count_pdf_pages(args.pdf) print(f"The PDF has {page_count} pages.") return # Check if files exist if not os.path.exists(args.pdf): print(f"Error: PDF file not found: {args.pdf}") return if not os.path.exists(args.doctags): print(f"Error: DocTags file not found: {args.doctags}") return # Process page(s) if args.page == 0: # Special case: process all pages process_all_pages( args.pdf, args.doctags, args.output, args.dpi, args.show, args.scale, args.scale_x, args.scale_y, args.adjust ) else: process_page( args.pdf, args.page, args.doctags, args.output, args.dpi, args.show, args.scale, args.scale_x, args.scale_y, args.adjust ) 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 """ 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.""" 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', 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())