50 lines
2.3 KiB
Python
50 lines
2.3 KiB
Python
|
import torch
|
||
|
|
from pathlib import Path
|
||
|
|
from docling.datamodel.base_models import InputFormat
|
||
|
|
from docling.datamodel.pipeline_options import PdfPipelineOptions
|
||
|
|
from docling.document_converter import DocumentConverter, PdfFormatOption
|
||
|
|
from docling_core.types.doc import ImageRefMode, PictureItem
|
||
|
|
|
||
|
|
# 1. Hardware Monkeypatch for Apple Silicon M3 CPU routing
|
||
|
|
torch.backends.mps.is_available = lambda: False
|
||
|
|
torch.backends.mps.is_built = lambda: False
|
||
|
|
torch.set_default_device("cpu")
|
||
|
|
|
||
|
|
def convert_with_images(pdf_name: str, output_md_name: str, image_folder_name: str = "extracted_images"):
|
||
|
|
source = f"./{pdf_name}"
|
||
|
|
output_dir = Path(image_folder_name)
|
||
|
|
output_dir.mkdir(exist_ok=True)
|
||
|
|
|
||
|
|
# 2. Configure pipeline to capture images
|
||
|
|
pipeline_options = PdfPipelineOptions()
|
||
|
|
pipeline_options.images_scale = 2.0 # Increase for higher resolution (2.0 = roughly 144 DPI)
|
||
|
|
pipeline_options.generate_picture_images = True # Tells the engine to extract figures/diagrams
|
||
|
|
|
||
|
|
# 3. Initialize converter with our custom options
|
||
|
|
converter = DocumentConverter(
|
||
|
|
format_options={
|
||
|
|
InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
print(f"🚀 Processing {pdf_name} and extracting figures...")
|
||
|
|
result = converter.convert(source)
|
||
|
|
|
||
|
|
# 4. Loop through the document items and save the pictures out to your folder
|
||
|
|
picture_counter = 0
|
||
|
|
for element, _level in result.document.iterate_items():
|
||
|
|
if isinstance(element, PictureItem):
|
||
|
|
picture_counter += 1
|
||
|
|
image_path = output_dir / f"figure_{picture_counter}.png"
|
||
|
|
|
||
|
|
# Save the raw image file cleanly using Pillow (built into docling)
|
||
|
|
element.get_image(result.document).save(image_path, "PNG")
|
||
|
|
|
||
|
|
# 5. Save the final Markdown file
|
||
|
|
# ImageRefMode.REFERENCED automatically adds clean image links like  into your MD file!
|
||
|
|
result.document.save_as_markdown(Path(output_md_name), image_mode=ImageRefMode.REFERENCED)
|
||
|
|
|
||
|
|
print(f"✅ Success! Saved {output_md_name} and extracted {picture_counter} figures to ./{image_folder_name}/")
|
||
|
|
|
||
|
|
# Run it for your textbook
|
||
|
|
convert_with_images("1-AS_NZS_3000_2018_Wiring_Rules_Standards_Australia.pdf", "1-AS_NZS_3000_2018_Wiring_Rules_Standards_Australia.md")
|