[Python] How can I extract images with their nearby captions or annotations using PyMuPDF...

I'm working on a script using PyMuPDF (fitz) to extract both text and images from PDF documents. The goal is to also retrieve any nearby captions or annotations that are close to the images—especially those directly below or above the image, as often seen in lecture slides or academic papers.

This is part of a larger workflow where the extracted content (text, hyperlinks, images and captions) will be converted into a Jupyter Book. The intention is for an AI agent to use this structured data to automatically generate high-quality lecture notes in MyST Markdown format, complete with images and proper references.

import fitz
import os

# Define the folder containing PDF files
pdf_folder = "pdf_files" # Change this to the folder containing your PDFs
output_folder = "output" # Folder to save extracted text and images
image_dir = os.path.join(output_folder, "images")

# Create output directories if they don't exist
if not os.path.exists(output_folder):
os.makedirs(output_folder)
if not os.path.exists(image_dir):
os.makedirs(image_dir)

# Iterate through all files in the folder
for pdf_file in os.listdir(pdf_folder):
if pdf_file.endswith(".pdf"): # Process only PDF files
pdf_path = os.path.join(pdf_folder, pdf_file)
output_txt = os.path.join(output_folder, f"{os.path.splitext(pdf_file)[0]}.txt")

# Open the PDF file
doc = fitz.open(pdf_path)

# Initialize a list to hold text content
text_content = []

# Iterate through each page in the PDF
for page_num in range(len(doc)):
page = doc[page_num]

# Extract text from the page
text = page.get_text()
text_content.append(text)

# Extract hyperlinks from the page
links = page.get_links()
for link in links:
if "uri" in link:
text_content.append(f"Link: {link['uri']}")

# Extract images from the page
images = page.get_images(full=True)
for img_index, img in enumerate(images):
xref = img[0]
base_image = doc.extract_image(xref)
image_bytes = base_image["image"]
image_filename = os.path.join(image_dir, f"{os.path.splitext(pdf_file)[0]}_page_{page_num + 1}_img_{img_index + 1}.png")

# Save the image to the output directory
with open(image_filename, "wb") as img_file:
img_file.write(image_bytes)

# Add placeholder in text
text_content.append(f"[[image:{image_filename}|Image from page {page_num + 1}]]")

# Add page break
text_content.append("\n--- Page Break ---\n")

# Write the text content to the output file
with open(output_txt, "w", encoding="utf-8") as txt_file:
for line in text_content:
txt_file.write(line + "\n")

# Close the PDF document
doc.close()

print(f"Extraction complete for '{pdf_file}'. Text and image references saved to '{output_txt}'. Images saved to '{image_dir}/'.")

Continue reading...