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

Push code produced this Week end

Browse source
This commit is contained in:
pjmalandrino 2025-05-13 15:20:48 +02:00
commit 9c5e13f342
4 changed files with 1430 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
README.md Normal file
View file

@ -0,0 +1,6 @@
python analyzer.py --image document.pdf --page 13
python visualizer.py --doctags output.doctags.txt --pdf document.pdf --page 13 --adjust --show
python analyzer.py --image document.pdf --page 13 && python visualizer.py --doctags output.doctags.txt --pdf document.pdf --page 13 --adjust --show

650
analyzer.py Normal file
View file

@ -0,0 +1,650 @@
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "docling-core",
# "mlx-vlm",
# "pillow",
# "requests",
# "argparse",
# "pdf2image",
# "pymupdf", # Optional for better PDF handling
# ]
# ///
import argparse
import os
import tempfile
import re
import base64
from io import BytesIO
from pathlib import Path
from urllib.parse import urlparse
import requests
from PIL import Image, UnidentifiedImageError
from pdf2image import convert_from_path, convert_from_bytes
from docling_core.types.doc import ImageRefMode
from docling_core.types.doc.document import DocTagsDocument, DoclingDocument
def parse_arguments():
"""Parse command line arguments."""
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",
help='Output file path')
parser.add_argument('--show', '-s', action='store_true',
help='Show output in browser')
parser.add_argument('--page', type=int, default=1,
help='Page number to process for PDF files (starts at 1)')
parser.add_argument('--dpi', type=int, default=200,
help='DPI for PDF rendering (higher values produce larger images)')
parser.add_argument('--debug', '-d', action='store_true',
help='Enable debug mode with extra output')
parser.add_argument('--doctags-only', action='store_true',
help='Generate only raw DocTags output without processing')
parser.add_argument('--all-pages', '-a', action='store_true',
help='Process all pages in a PDF without asking')
parser.add_argument('--start-page', type=int, default=1,
help='Start processing PDF from this page number')
parser.add_argument('--end-page', type=int, default=None,
help='Stop processing PDF at this page number')
parser.add_argument('--max-pages', type=int, default=None,
help='Maximum number of pages to process')
return parser.parse_args()
def load_image(image_path, page_num=1, dpi=200):
"""Load image from URL, local image file, or PDF."""
if urlparse(image_path).scheme in ['http', 'https']: # it is a URL
try:
response = requests.get(image_path, stream=True, timeout=10)
response.raise_for_status()
content = response.content
# Check if it's a PDF
if image_path.lower().endswith('.pdf') or response.headers.get('Content-Type') == 'application/pdf':
print(f"Converting PDF from URL (page {page_num})...")
pdf_images = convert_from_bytes(content, 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 first (and only) page
else:
return Image.open(BytesIO(content))
except requests.exceptions.RequestException as e:
raise Exception(f"Error loading image from URL: {e}")
else: # it is a local file
image_path = Path(image_path)
if not image_path.exists():
raise FileNotFoundError(f"File not found: {image_path}")
# Check if it's a PDF
if image_path.suffix.lower() == '.pdf':
print(f"Converting PDF to image (page {page_num}, DPI: {dpi})...")
try:
pdf_images = convert_from_path(
image_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}")
else:
try:
return Image.open(image_path)
except UnidentifiedImageError:
raise Exception(f"Cannot identify image file: {image_path}. Make sure it's a valid image format or PDF.")
def cleanup_doctags(doctags_text):
"""Clean up the DocTags structure."""
print("Cleaning up DocTags structure...")
# Simplified approach to extract valuable information
# Extract headers
headers = re.findall(r'<section_header_level_1>.*?>(.*?)</section_header_level_1>', doctags_text)
# Extract text paragraphs
paragraphs = re.findall(r'<text>.*?>(.*?)</text>', doctags_text)
# Extract list items
list_items = re.findall(r'<list_item>.*?>(.*?)</list_item>', doctags_text)
# Extract footer
footer = re.search(r'<page_footer>.*?>(.*?)</page_footer>', doctags_text)
footer_text = footer.group(1) if footer else ""
# Create a clean doctags structure
clean_doctags = "<doctag>\n"
# Add headers
for header in headers:
clean_doctags += f"<section_header_level_1>{header}</section_header_level_1>\n"
# Add text
for paragraph in paragraphs:
clean_doctags += f"<text>{paragraph}</text>\n"
# Add list if any items found
if list_items:
clean_doctags += "<unordered_list>\n"
for item in list_items:
clean_doctags += f"<list_item>{item}</list_item>\n"
clean_doctags += "</unordered_list>\n"
# Add footer if present
if footer_text:
clean_doctags += f"<page_footer>{footer_text}</page_footer>\n"
clean_doctags += "</doctag>"
return clean_doctags
def extract_all_tags(doctags_text):
"""Extract all unique DocTags from the text."""
print("Extracting all DocTags...")
# Use regex to find all tags
tag_pattern = r'</?(\w+)(?:\s[^>]*)?>'
all_tags = re.findall(tag_pattern, doctags_text)
# Remove duplicates and sort
unique_tags = sorted(set(all_tags))
# Create a list of tags with their frequencies
tag_counts = {}
for tag in all_tags:
tag_counts[tag] = tag_counts.get(tag, 0) + 1
# Format the output
tags_output = "# DocTags Found\n\n"
tags_output += "| Tag | Count |\n"
tags_output += "|-----|-------|\n"
for tag in unique_tags:
tags_output += f"| {tag} | {tag_counts[tag]} |\n"
# Add examples section
tags_output += "\n\n# DocTags Examples\n\n"
# Find example usages for each tag
for tag in unique_tags:
# Find an opening tag with content
open_pattern = f'<{tag}(?:\\s[^>]*)?>.*?</{tag}>'
examples = re.findall(open_pattern, doctags_text, re.DOTALL)
if examples:
# Limit to first example
example = examples[0]
# Truncate if too long
if len(example) > 200:
example = example[:197] + "..."
tags_output += f"## {tag}\n\n```xml\n{example}\n```\n\n"
return tags_output
def debug_doctags(doctags_text, debug_mode=False):
"""Debug function to analyze doctag structure."""
if not debug_mode:
return doctags_text
print("\nAnalyzing DocTags structure:")
print(f"Total length: {len(doctags_text)} characters")
# Check for valid opening and closing tags
opening_tags = []
i = 0
while i < len(doctags_text):
open_tag_start = doctags_text.find('<', i)
if open_tag_start == -1:
break
open_tag_end = doctags_text.find('>', open_tag_start)
if open_tag_end == -1:
print(f"WARNING: Unclosed tag starting at position {open_tag_start}")
break
tag_content = doctags_text[open_tag_start+1:open_tag_end]
if tag_content.startswith('/'):
# This is a closing tag
if not opening_tags:
print(f"WARNING: Closing tag {tag_content} without matching opening tag")
else:
last_open = opening_tags.pop()
if last_open != tag_content[1:]:
print(f"WARNING: Mismatched tags: opening <{last_open}> vs closing <{tag_content}>")
elif not tag_content.startswith('!') and not ' ' in tag_content:
# This is an opening tag (not a comment or self-closing)
opening_tags.append(tag_content)
i = open_tag_end + 1
if opening_tags:
print(f"WARNING: Unclosed tags: {', '.join(opening_tags)}")
# Count all tags
tag_counts = {}
tag_pattern = r'<(\w+)(?:\s|>)'
import re
for tag in re.findall(tag_pattern, doctags_text):
tag_counts[tag] = tag_counts.get(tag, 0) + 1
print("Tag counts:")
for tag, count in tag_counts.items():
print(f" <{tag}>: {count}")
# Special check for doctag
if doctags_text.count('<doctag>') != 1:
print(f"WARNING: Expected 1 <doctag> tag, found {doctags_text.count('<doctag>')}")
if doctags_text.count('</doctag>') != 1:
print(f"WARNING: Expected 1 </doctag> tag, found {doctags_text.count('</doctag>')}")
return doctags_text
def create_html_document(image_path, doctags_text, pil_image, args):
"""Create HTML directly from DocTags and image."""
print("Creating HTML document directly...")
# Get filename for document title
doc_name = Path(image_path).stem
# Extract content from doctags
title_match = re.search(r'<section_header_level_1>.*?>(.*?)</section_header_level_1>', doctags_text)
title = title_match.group(1) if title_match else doc_name
# Clean title by removing any location tags
title = re.sub(r'<loc_\d+>', '', title)
# Extract text content
text_matches = re.findall(r'<text>.*?>(.*?)</text>', doctags_text)
paragraphs = []
for text in text_matches:
# Clean text by removing any location tags
cleaned_text = re.sub(r'<loc_\d+>', '', text)
paragraphs.append(cleaned_text)
# Extract list items
list_items = []
item_matches = re.findall(r'<list_item>.*?>(.*?)</list_item>', doctags_text)
for item in item_matches:
# Clean item by removing any location tags
cleaned_item = re.sub(r'<loc_\d+>', '', item)
list_items.append(cleaned_item)
# Extract footer
footer_match = re.search(r'<page_footer>.*?>(.*?)</page_footer>', doctags_text)
footer_text = ""
if footer_match:
footer_text = re.sub(r'<loc_\d+>', '', footer_match.group(1))
# Create image data URI
max_dimension = 1200
img_for_embedding = pil_image
if max(pil_image.size) > max_dimension:
ratio = min(max_dimension / pil_image.size[0], max_dimension / pil_image.size[1])
new_size = (int(pil_image.size[0] * ratio), int(pil_image.size[1] * ratio))
img_for_embedding = pil_image.resize(new_size, Image.LANCZOS)
buffered = BytesIO()
img_for_embedding.save(buffered, format="PNG", optimize=True, quality=90)
img_str = base64.b64encode(buffered.getvalue()).decode()
img_data_uri = f"data:image/png;base64,{img_str}"
# Build HTML paragraphs
paragraphs_html = "".join([f"<p>{text}</p>\n" for text in paragraphs]) if paragraphs else ""
# Build list HTML
list_html = ""
if list_items:
list_html = "<ul>\n"
for item in list_items:
list_html += f"<li>{item}</li>\n"
list_html += "</ul>\n"
# Create footer HTML
footer_html = f"<div class='page-footer'>{footer_text}</div>\n" if footer_text else ""
# Create HTML document
html_content = f'''<!DOCTYPE html>
<html>
<head>
<meta charset="UTF-8">
<title>{title}</title>
<style>
body {{ font-family: Arial, sans-serif; margin: 0 auto; max-width: 800px; padding: 20px; line-height: 1.6; }}
img {{ max-width: 100%; height: auto; }}
h1 {{ color: #333; border-bottom: 1px solid #eee; padding-bottom: 10px; }}
figure {{ margin: 20px 0; text-align: center; }}
figcaption {{ color: #666; font-style: italic; margin-top: 5px; }}
ul {{ margin-left: 20px; padding-left: 20px; }}
li {{ margin-bottom: 10px; }}
.page-footer {{ color: #666; margin-top: 20px; border-top: 1px solid #eee; padding-top: 5px; text-align: center; }}
</style>
</head>
<body>
<div class='page'>
<h1>{title}</h1>
<figure>
<img src="{img_data_uri}" alt="Document Image"/>
<figcaption>Page {args.page}</figcaption>
</figure>
{paragraphs_html}
{list_html}
{footer_html}
</div>
</body>
</html>'''
return html_content
def create_markdown_document(image_path, doctags_text, args):
"""Create Markdown document from DocTags."""
print("Creating Markdown document...")
# Get filename for document title
doc_name = Path(image_path).stem
# Extract content from doctags
title_match = re.search(r'<section_header_level_1>.*?>(.*?)</section_header_level_1>', doctags_text)
title = title_match.group(1) if title_match else doc_name
# Clean title by removing any location tags
title = re.sub(r'<loc_\d+>', '', title)
# Extract text content
text_matches = re.findall(r'<text>.*?>(.*?)</text>', doctags_text)
paragraphs = []
for text in text_matches:
# Clean text by removing any location tags
cleaned_text = re.sub(r'<loc_\d+>', '', text)
paragraphs.append(cleaned_text)
# Extract list items
list_items = []
item_matches = re.findall(r'<list_item>.*?>(.*?)</list_item>', doctags_text)
for item in item_matches:
# Clean item by removing any location tags
cleaned_item = re.sub(r'<loc_\d+>', '', item)
list_items.append(cleaned_item)
# Extract footer
footer_match = re.search(r'<page_footer>.*?>(.*?)</page_footer>', doctags_text)
footer_text = ""
if footer_match:
footer_text = re.sub(r'<loc_\d+>', '', footer_match.group(1))
# Build markdown content
markdown = f"# {title}\n\n"
markdown += f"*Page {args.page}*\n\n"
# Add paragraphs
for p in paragraphs:
markdown += f"{p}\n\n"
# Add list
if list_items:
for item in list_items:
markdown += f"* {item}\n"
markdown += "\n"
# Add footer
if footer_text:
markdown += f"---\n{footer_text}\n"
return markdown
def process_page(args, model, processor, config, image_path, pil_image, page_num=1):
"""Process a single page from a PDF or image file."""
from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import stream_generate
# Prepare input
prompt = args.prompt
output_base = Path(args.output)
output_path = output_base
# If processing PDF and output is a path without explicit numbering, add page numbers
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}"
print(f"Processing page {page_num}, output will be saved to {output_path}")
# Create a temporary file for the image
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_img_file:
temp_img_path = temp_img_file.name
pil_image.save(temp_img_path, format='PNG')
print(f"Saved temporary image to: {temp_img_path}")
try:
# Apply chat template
formatted_prompt = apply_chat_template(processor, config, prompt, num_images=1)
# Generate output
print(f"Generating DocTags for page {page_num}: \n\n")
output = ""
for token in stream_generate(
model, processor, formatted_prompt, [temp_img_path], max_tokens=4096, verbose=False
):
output += token.text
print(token.text, end="")
if "</doctag>" in token.text:
break
print("\n\n")
# Debug and clean the doctags content
output = debug_doctags(output, args.debug)
# Extract all tags if in DocTags-only mode
if args.doctags_only:
tags_analysis = extract_all_tags(output)
# Clean the output for document creation
cleaned_output = cleanup_doctags(output)
finally:
# Clean up the temporary file
if os.path.exists(temp_img_path):
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')
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')
with open(tags_path, 'w', encoding='utf-8') as f:
f.write(tags_analysis)
print(f"DocTags analysis saved to: {tags_path}")
# If not in DocTags-only mode, create documents
if not args.doctags_only:
# Create markdown document
md_content = create_markdown_document(image_path, output, args)
md_path = output_path.with_suffix('.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')
with open(html_path, 'w', encoding='utf-8') as f:
f.write(html_content)
print(f"HTML saved to: {html_path}")
# Calculate final file size
html_size = os.path.getsize(html_path)
print(f"Final HTML file size: {html_size} bytes")
# Open in browser if requested (only for the first page or single pages)
if args.show and html_size > 100 and page_num <= 1:
import webbrowser
webbrowser.open(f"file:///{str(html_path.resolve())}")
# 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')
pil_image.save(img_debug_path)
print(f"Saved debug image to: {img_debug_path}")
return output_path
def main():
# Parse arguments
args = parse_arguments()
# Settings
DEBUG_MODE = args.debug
DOCTAGS_ONLY = args.doctags_only
image_path = args.image
# Load the model
print("Loading model...")
try:
from mlx_vlm import load, generate
from mlx_vlm.utils import load_config
model_path = "ds4sd/SmolDocling-256M-preview-mlx-bf16"
model, processor = load(model_path)
config = load_config(model_path)
except Exception as e:
print(f"Error loading model: {e}")
import traceback
traceback.print_exc()
return
# Check if the input is a PDF
pdf_path = Path(image_path)
is_pdf = False
pdf_page_count = 0
if pdf_path.suffix.lower() == '.pdf':
is_pdf = True
# Try to get page count
try:
try:
import fitz # PyMuPDF
pdf_document = fitz.open(pdf_path)
pdf_page_count = len(pdf_document)
print(f"PDF detected with {pdf_page_count} pages")
pdf_document.close()
except ImportError:
print("PyMuPDF not installed. Using pdf2image to estimate page count...")
from pdf2image import pdfinfo_from_path
pdf_info = pdfinfo_from_path(pdf_path)
pdf_page_count = pdf_info["Pages"]
print(f"PDF detected with {pdf_page_count} pages")
except Exception as e:
print(f"Could not determine PDF page count: {e}")
print("Will process the specified page only.")
pdf_page_count = args.page
# Determine which pages to process
process_all_pages = args.all_pages
start_page = args.start_page
end_page = args.end_page if args.end_page else pdf_page_count
max_pages = args.max_pages
# Validate page ranges
if start_page < 1:
start_page = 1
if end_page and end_page > pdf_page_count:
end_page = pdf_page_count
# Ask user if they want to process all pages or just one (if not specified by arguments)
if is_pdf and pdf_page_count > 1 and not process_all_pages and args.start_page == 1 and not args.end_page:
if args.page > 1:
print(f"You specified to process page {args.page}.")
process_all_pages = False
start_page = args.page
end_page = args.page
else:
user_input = input("Do you want to process all pages? [y/N]: ")
process_all_pages = user_input.lower() in ['y', 'yes']
if not process_all_pages:
# Ask for specific page or range
page_input = input(f"Enter page number(s) to process (e.g., 3 or 1-5) [1]: ")
if page_input.strip():
if '-' in page_input:
try:
start_str, end_str = page_input.split('-')
start_page = int(start_str.strip())
end_page = int(end_str.strip())
except ValueError:
print("Invalid range format. Using default page 1.")
start_page = end_page = 1
else:
try:
start_page = end_page = int(page_input.strip())
except ValueError:
print("Invalid page number. Using default page 1.")
start_page = end_page = 1
# Apply max_pages limit if specified
if max_pages and end_page - start_page + 1 > max_pages:
end_page = start_page + max_pages - 1
# Process pages
if is_pdf and (process_all_pages or start_page != end_page):
page_range = range(start_page, end_page + 1)
print(f"Processing pages {start_page} to {end_page} ({len(page_range)} pages)...")
processed_pages = []
for page_num in page_range:
print(f"\n{'='*50}\nProcessing PDF page {page_num}/{end_page}\n{'='*50}\n")
# Update the page argument
args.page = page_num
# Load the specific page
try:
pil_image = load_image(image_path, page_num=page_num, dpi=args.dpi)
print(f"Page {page_num} loaded: {pil_image.size}")
# Process the page
output_path = process_page(args, model, processor, config, image_path, pil_image, page_num)
processed_pages.append(output_path)
except Exception as e:
print(f"Error processing page {page_num}: {e}")
import traceback
traceback.print_exc()
continue
print(f"\nProcessed {len(processed_pages)} pages from PDF.")
print(f"Output files: {', '.join([str(p) for p in processed_pages])}")
else:
# Process just one page (either it's not a PDF or user only wants one page)
try:
# Load image resource
print(f"Loading {'PDF page' if is_pdf else 'image'} from: {image_path}")
pil_image = load_image(image_path, page_num=args.page, dpi=args.dpi)
print(f"Image loaded: {pil_image.size}")
# Process the single page
process_page(args, model, processor, config, image_path, pil_image)
except Exception as e:
print(f"Error processing {image_path}: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()

BIN
document.pdf Normal file
View file

Binary file not shown.

774
visualizer.py Normal file
View file

@ -0,0 +1,774 @@
#!/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'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>'
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 <doctag> tags
doctag_pattern = r'<doctag>(.*?)</doctag>'
doctag_match = re.search(doctag_pattern, doctags_content, re.DOTALL)
if not doctag_match:
raise ValueError("No <doctag> 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_name})>'
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}(.*?)</{tag_name}>'
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"""<!DOCTYPE html>
<html>
<head>
<meta charset="UTF-8">
<title>DocTags Zone Visualization - Page {page_num}</title>
<style>
body {{ font-family: Arial, sans-serif; margin: 0; padding: 20px; background-color: #f5f5f5; }}
.container {{ position: relative; display: inline-block; margin: 0 auto; background-color: white; box-shadow: 0 2px 10px rgba(0,0,0,0.2); width: {img_width}px; height: {img_height}px; }}
.page-image {{ display: block; max-width: 100%; height: auto; }}
.zone-overlay {{
position: absolute;
border: 2px solid;
pointer-events: none;
background-color: rgba(255, 255, 255, 0.1);
transition: background-color 0.3s;
}}
.zone-overlay:hover {{
background-color: rgba(255, 255, 255, 0.3);
}}
.zone-label {{
position: absolute;
top: 0;
left: 0;
font-size: 12px;
padding: 2px 5px;
background-color: rgba(255, 255, 255, 0.8);
border-radius: 0 0 3px 0;
max-width: 100%;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}}
.controls {{
margin-bottom: 10px;
padding: 10px;
background-color: white;
border-radius: 3px;
box-shadow: 0 1px 3px rgba(0,0,0,0.1);
}}
.legend {{
display: flex;
flex-wrap: wrap;
gap: 10px;
margin-top: 10px;
}}
.legend-item {{
display: flex;
align-items: center;
font-size: 14px;
}}
.legend-color {{
width: 16px;
height: 16px;
margin-right: 5px;
border: 1px solid rgba(0,0,0,0.2);
}}
.zone-list {{
margin-top: 20px;
background-color: white;
border-radius: 3px;
padding: 10px;
box-shadow: 0 1px 3px rgba(0,0,0,0.1);
}}
.zone-item {{
margin-bottom: 10px;
padding-bottom: 10px;
border-bottom: 1px solid #eee;
}}
.zone-type {{
font-weight: bold;
display: inline-block;
padding: 2px 6px;
border-radius: 3px;
color: white;
margin-bottom: 5px;
}}
.zone-content {{
margin-left: 10px;
font-size: 14px;
color: #333;
}}
.pagination {{
margin-top: 10px;
display: flex;
gap: 5px;
align-items: center;
}}
.pagination a {{
display: inline-block;
padding: 5px 10px;
background-color: #f0f0f0;
text-decoration: none;
color: #333;
border-radius: 3px;
}}
.pagination a:hover {{
background-color: #e0e0e0;
}}
.pagination .current {{
background-color: #2196F3;
color: white;
}}
.pagination-info {{
margin-left: 10px;
color: #666;
}}
</style>
</head>
<body>
<div class="controls">
<h2>DocTags Zone Visualization - {Path(pdf_path).stem}</h2>
<div style="display: flex; justify-content: space-between; align-items: center;">
<div>
<label><input type="checkbox" id="toggle-labels" checked> Show Labels</label>
<label><input type="checkbox" id="toggle-zones" checked> Show Zones</label>
<label><input type="checkbox" id="toggle-transparency" checked> Transparent Zones</label>
<div class="legend">"""
# 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"""
<div class="legend-item">
<div class="legend-color" style="background-color: {color};"></div>
{zone_type}
</div>"""
html += """
</div>
</div>
<div class="pagination">"""
# Add page navigation
if total_pages > 1:
for p in range(1, total_pages + 1):
if p == page_num:
html += f"""
<span class="current">{p}</span>"""
else:
# Create the filename for this page
page_output = Path(output_path).with_stem(f"{Path(output_path).stem.split('_page_')[0]}_page_{p}")
html += f"""
<a href="{page_output.name}">{p}</a>"""
html += f"""
<span class="pagination-info">Page {page_num} of {total_pages}</span>"""
html += """
</div>
</div>
</div>
<div class="container">
<img src="data:image/png;base64,""" + img_base64 + """" class="page-image">
"""
# 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"""
<div class="zone-overlay" style="left: {zone['x1']}px; top: {zone['y1']}px; width: {width}px; height: {height}px; border-color: {color};">
<div class="zone-label" style="border-bottom: 2px solid {color}; border-right: 2px solid {color};">
{zone_type}
</div>
</div>"""
html += """
</div>
<div class="zone-list">
<h3>Detected Zones:</h3>
"""
# Add zone details
for zone in zones:
zone_type = zone['type']
color = zone_colors.get(zone_type, zone_colors['default'])
html += f"""
<div class="zone-item">
<div class="zone-type" style="background-color: {color};">{zone_type}</div>
<div class="zone-coords">Position: ({zone['x1']}, {zone['y1']}) - ({zone['x2']}, {zone['y2']})</div>
"""
if zone['content']:
html += f"""
<div class="zone-content">{zone['content']}</div>
"""
html += """
</div>
"""
html += """
</div>
<script>
// Toggle labels visibility
document.getElementById('toggle-labels').addEventListener('change', function() {
var labels = document.querySelectorAll('.zone-label');
labels.forEach(function(label) {
label.style.display = this.checked ? 'block' : 'none';
}, this);
});
// Toggle zones visibility
document.getElementById('toggle-zones').addEventListener('change', function() {
var zones = document.querySelectorAll('.zone-overlay');
zones.forEach(function(zone) {
zone.style.display = this.checked ? 'block' : 'none';
}, this);
});
// Toggle zone transparency
document.getElementById('toggle-transparency').addEventListener('change', function() {
var zones = document.querySelectorAll('.zone-overlay');
zones.forEach(function(zone) {
zone.style.backgroundColor = this.checked ? 'rgba(255, 255, 255, 0.1)' : 'rgba(255, 255, 255, 0.6)';
}, this);
});
</script>
</body>
</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
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=False):
"""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}")
# Apply scaling to coordinates if required
if scale != 1.0 or scale_x is not None or scale_y is not None:
x_scale = scale_x if scale_x is not None else scale
y_scale = scale_y if scale_y is not None else scale
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 scaling: X={x_scale}, Y={y_scale}")
# Auto-adjust scaling if requested
if adjust:
# Find the maximum coordinates in the zones
max_x = max([zone['x2'] for zone in zones]) if zones else 0
max_y = max([zone['y2'] for zone in zones]) if zones else 0
# If max coordinates exceed image dimensions or are too small
if max_x > 0 and max_y > 0:
width, height = image.size
# Calculate appropriate scaling factors
if max_x > width * 1.1 or max_x < width * 0.5:
x_scale = width / max_x
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 > height * 1.1 or max_y < height * 0.5:
y_scale = height / max_y
print(f"Auto-adjusted Y scale to {y_scale:.3f} (image height: {height}, max zone y: {max_y})")
else:
y_scale = 1.0
# 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 '<doctag>' in line:
in_doctags = True
if in_doctags:
doctags_lines.append(line)
if '</doctag>' 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())