769 lines
No EOL
26 KiB
Python
769 lines
No EOL
26 KiB
Python
#!/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
|
|
|
|
# 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 '<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()) |