487 lines
No EOL
18 KiB
Python
487 lines
No EOL
18 KiB
Python
#!/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 ensure_results_folder():
|
|
"""Create the results folder if it doesn't exist."""
|
|
results_dir = Path("results")
|
|
if not results_dir.exists():
|
|
results_dir.mkdir()
|
|
print(f"Created results directory: {results_dir}")
|
|
return results_dir
|
|
|
|
|
|
def parse_arguments():
|
|
"""Parse command line arguments."""
|
|
results_dir = ensure_results_folder()
|
|
|
|
parser = argparse.ArgumentParser(description='Convert an image or PDF to docling format')
|
|
parser.add_argument('--image', '-i', type=str, required=True,
|
|
help='Path to local image file, PDF file, or URL')
|
|
parser.add_argument('--prompt', '-p', type=str, default="Convert this page to docling.",
|
|
help='Prompt for the model')
|
|
parser.add_argument('--output', '-o', type=str, default=str(results_dir / "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 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
|
|
|
|
# Ensure results folder exists
|
|
results_dir = ensure_results_folder()
|
|
|
|
# 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 = results_dir / 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 in results folder
|
|
doctags_path = results_dir / f"{output_path.stem}.doctags.txt"
|
|
with open(doctags_path, 'w', encoding='utf-8') as f:
|
|
f.write(output)
|
|
print(f"Raw DocTags saved to: {doctags_path}")
|
|
|
|
# Save the tag analysis if in DocTags-only mode
|
|
if args.doctags_only:
|
|
tags_path = results_dir / f"{output_path.stem}.tags.md"
|
|
with open(tags_path, 'w', encoding='utf-8') as f:
|
|
f.write(tags_analysis)
|
|
print(f"DocTags analysis saved to: {tags_path}")
|
|
|
|
# Save a copy of the processed image for reference if in debug mode
|
|
if args.debug:
|
|
img_debug_path = results_dir / f"{output_path.stem}.debug.png"
|
|
pil_image.save(img_debug_path)
|
|
print(f"Saved debug image to: {img_debug_path}")
|
|
|
|
return output_path
|
|
|
|
|
|
def main():
|
|
# Ensure results folder exists
|
|
ensure_results_folder()
|
|
|
|
# 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() |