[Python] Zero-Loss PDF Image Extraction in Python

I’m working on a PDF exam paper → Markdown conversion pipeline and I’m struggling with reliably finding and extracting diagrams and visuals from a large set of PDF documents.

I’ve tried several approaches (e.g. Mistral OCR, PyMuPDF, etc.), but none have produced consistently good results.

My current approach looks like this:

1. 
   Convert each page of the PDF to a JPEG at 300 DPI.
   
   
2. 
   Use an LLM to analyse each page image and decide whether it contains any non-text visuals required to answer a question.
   
   • 
     If no relevant visuals are found, I store an entry like:
     
     {
     "stem": "1f2020",
     "page": 2,
     "file": "/data/gcse_maths_edexcel/done/1f2020_p002.jpg",
     "llm_judgment": {
     "outcome": "NO",
     "details": ""
     }
     }
   • 
     If visuals are required, I store them like:
     
     {
     "stem": "1f2020",
     "page": 3,
     "file": "/data/gcse_maths_edexcel/done/1f2020_p003.jpg",
     "llm_judgment": {
     "outcome": "YES",
     "items": [
     {
     "question": "6",
     "summary": "Octagonal spinner divided into 8 equal sectors labeled A, A, A, B, C, C, C, C with a diagonal pointer crossing the centre.",
     "regenerate_prompt": "From the provided page image, for question 6 specifically, isolate and recreate ONLY this visual: Octagonal spinner divided into 8 equal sectors labeled (clockwise from top): A, A, A, B, C, C, C, C with a diagonal spinner pointer crossing through the center.. Include every detail exactly as shown. Do not include any surrounding page text or unrelated elements. Render as a clean vector (SVG preferred; PNG acceptable) with a transparent background and uniform stroke weights. Center it, crop all whitespace to the visual’s tight bounds, and output only the final image (no captions or extra text)."
     },
     {
     "question": "6a",
     "summary": "Horizontal probability scale from 0 to 1 with ticks at 0, 1/2, and 1.",
     "regenerate_prompt": "From the provided page image, for question 6a specifically, isolate and recreate ONLY this linear scale: Horizontal probability scale from 0 to 1 with ticks at 0, 1/2 (fraction), and 1.. Include exactly the same baseline, tick marks, and numeric/label positions as shown. Do not include any surrounding page text or unrelated elements. Render as a clean vector (SVG preferred; PNG acceptable) with a transparent background and uniform stroke weights. Center it, crop all whitespace to the scale’s tight bounds, and output only the final image (no captions or extra text)."
     }
     ]
     }
     }
3. 
   For every required visual, I pass the reference page image plus a detailed regenerate prompt to the LLM and ask it to recreate the diagram as a clean SVG or PNG with a transparent background, tightly cropped.
   

This approach kinda works, but the results are extremely inconsistent. LLMs really struggle to accurately recreate mathematical diagrams, graphs, probability scales, coordinate grids, and similar visuals — especially when precision matters (e.g. tick marks, angles, and distances).

I feel like I’ve hit a wall here.

Has anyone successfully solved this problem before? Specifically, I’m looking for a more reliable way to:

• Detect exactly which visuals are needed to answer questions in exam PDFs.
• Extract or recreate those visuals with pixel-perfect accuracy (spinners, graphs, scales, number lines, etc.).
• Ideally, integrate this into an automated pipeline without having to manually crop images page by page.

Any recommendations for tools, libraries, or workflows that can handle this better than an LLM would be hugely appreciated.

#!/usr/bin/env python3
import argparse
import base64
import json
import os
import re
import time
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, List, Tuple

import fitz
import requests

from openai import AzureOpenAI
from PIL import Image

MODEL = "gpt-5"
CLIENT = AzureOpenAI(
api_key=os.getenv("AZURE_EU_OPENAI_KEY"),
api_version="2025-04-01-preview",
azure_endpoint=os.getenv("AZURE_EU_OPENAI_ENDPOINT"),
)


def ensure_dir(p: str) -> None:
Path(p).mkdir(parents=True, exist_ok=True)


def judge_image_b64(
b64_jpg: str, stem_name: str, page_num: int, file_path: str
) -> Dict[str, Any]:
sys_judge = (
"You judge exam pages for ALL non-text visuals REQUIRED to answer any question or sub-question on the page. "
"Capture ANY visual that could affect answering, including but not limited to: spinners, pie charts, graphs, axes, coordinate grids, plotted points, number lines, probability scales, timelines, rulers, flowcharts, maps, tables with symbols, diagrams with labels or arrows, bespoke illustrations, or any drawn shapes conveying information. "
"Use the full question reference when possible (e.g., 6a, 6b, 2(ii)); if uncertain, use the main number like 6 or 'unknown'. "
"Return ONLY JSON in one of these forms: "
'{"outcome":"NO","details":""} '
"or "
'{"outcome":"YES","items":[{"question":"<number or number+subpart or \\"unknown\\">","summary":"<short description of the exact visual>",'
'"regenerate_prompt":"<fully specified redraw prompt for ONLY this visual>"}...]}. '
"If multiple visuals exist, include one item per visual ordered top-to-bottom then left-to-right. "
"Each regenerate_prompt must be sufficient to recreate ONLY that visual as a clean vector (SVG preferred; PNG acceptable), transparent background, uniform stroke weights, cropped to tight bounds, no extra text."
)

generic_template = (
"From the provided page image, for question {q} specifically, isolate and recreate ONLY this visual: {summary}. "
"Produce a 1:1 replica of the source visual. Preserve exact geometry, aspect ratio, relative distances, angles, radii, coordinates, tick spacing, grid spacing, and label placement. "
"For number lines, axes, or probability scales, exactly match baseline length, tick count and positions, and numeric/label placement. "
"For charts, diagrams, shapes, or spinners, exactly match segment counts, angles, coordinates/paths, and pointer orientation. "
"Do not include any surrounding page text or unrelated elements. "
"Render as a clean vector (SVG preferred; PNG acceptable) on a fully transparent background. "
"Scale strokes proportionally so the visual appears identical at the output size. "
"Center it, crop tightly to the visual’s bounds with zero padding, and output only the final image with no captions or extra text."
)

scale_template = (
"From the provided page image, for question {q} specifically, isolate and recreate ONLY this linear scale: {summary}. "
"Produce a 1:1 replica of the source scale. Preserve exact geometry, baseline length, orientation, aspect ratio, tick count, tick spacing and positions (including minor and major ticks), and the exact numeric/label text, placement, and alignment. "
"Include arrows, end caps, markers, braces, and any anchors exactly as shown when present. "
"Do not include any surrounding page text or unrelated elements. "
"Render as a clean vector (SVG preferred; PNG acceptable) on a fully transparent background. "
"Scale stroke widths proportionally to match the original appearance at the output size. "
"Center the scale, crop tightly to its bounds with zero padding, and output only the final image with no captions or extra text."
)

user_text = (
f"Stem: {stem_name}\nPage: {page_num}\nImage file: {file_path}\n"
"Task: Inspect the attached image of the page and respond strictly with the required JSON."
)

r = CLIENT.chat.completions.create(
model=MODEL,
messages=[
{"role": "system", "content": sys_judge},
{
"role": "user",
"content": [
{"type": "text", "text": user_text},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{b64_jpg}"},
},
],
},
],
)

content = (r.choices[0].message.content or "").strip()
m = re.search(r"\{[\s\S]*\}\s*$", content)
data = json.loads(m.group(0) if m else content)

if isinstance(data, dict) and data.get("outcome") == "YES":
items = data.get("items") or []
out_items = []
for it in items:
q = str(it.get("question") or "unknown")
summary = str(it.get("summary") or "").strip()
if re.search(
r"(probability\s*scale|number\s*line|timeline|axis|axes|scale|ruler)",
summary,
re.I,
):
regen = scale_template.format(
q=q, summary=summary if summary else "the scale"
)
else:
regen = generic_template.format(
q=q, summary=summary if summary else "the visual"
)
it["regenerate_prompt"] = regen
out_items.append(it)
data["items"] = out_items

return data


def split_to_jpegs(
pdf_path: str,
out_dir: str,
prefix: str,
digits: int,
start: int,
end: int,
dpi: int,
quality: int,
subsampling: int,
progressive: bool,
optimize: bool,
overwrite: bool,
) -> List[Dict[str, Any]]:
ensure_dir(out_dir)
pages: List[Dict[str, Any]] = []
file_stem = Path(pdf_path).stem
with fitz.open(pdf_path) as doc:
total = len(doc)
s = max(1, start or 1)
e = min(total, end or total)
zoom = dpi / 72.0
mat = fitz.Matrix(zoom, zoom)
for pno in range(s - 1, e):
page_num = pno + 1
page = doc[pno]
pm = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB, alpha=False)
im = Image.frombytes("RGB", (pm.width, pm.height), pm.samples)
name = f"{file_stem}_{prefix}{str(page_num).zfill(digits)}.jpg"
out_path = Path(out_dir) / name
if not out_path.exists() or overwrite:
title = f"{file_stem} page {page_num}"
exif = im.getexif()
exif[0x010E] = title
im.save(
out_path,
format="JPEG",
quality=quality,
subsampling=subsampling,
progressive=progressive,
optimize=optimize,
dpi=(dpi, dpi),
exif=exif.tobytes(),
)
pages.append(
{
"page": page_num,
"title": f"{file_stem} page {page_num}",
"file": str(out_path),
"width_px": pm.width,
"height_px": pm.height,
"dpi": dpi,
"quality": quality,
"subsampling": subsampling,
"progressive": progressive,
"stem": file_stem,
}
)
return pages


def load_existing_judgments(judge_path: Path) -> Tuple[List[Dict[str, Any]], set]:
if not judge_path.exists():
return [], set()
try:
with open(judge_path, "r", encoding="utf-8") as f:
arr = json.load(f)
pages = {
int(x.get("page")) for x in arr if isinstance(x, dict) and "page" in x
}
return arr, pages
except Exception:
return [], set()


def append_and_flush(judge_path: Path, row: Dict[str, Any]) -> None:
existing, _ = load_existing_judgments(judge_path)
existing.append(row)
with open(judge_path, "w", encoding="utf-8") as f:
json.dump(existing, f, ensure_ascii=False, indent=2)


def analyze_pages(pages: List[Dict[str, Any]], base_out_dir: str) -> None:
if not pages:
return
stem_dir = Path(base_out_dir)
ensure_dir(str(stem_dir))
judge_path = stem_dir / f"{pages[0].get('stem','pages')}_judge.json"
_, done_pages = load_existing_judgments(judge_path)
for r in pages:
page_num = int(r.get("page"))
file_path = r.get("file")
if page_num in done_pages:
continue
with open(file_path, "rb") as f:
b64_jpg = base64.b64encode(f.read()).decode("ascii")
result = judge_image_b64(b64_jpg, r.get("stem", "pages"), page_num, file_path)
row = {
"stem": r.get("stem", "pages"),
"page": page_num,
"file": file_path,
"llm_judgment": result,
}
append_and_flush(judge_path, row)
print(json.dumps(row, ensure_ascii=False))


def retry(max_attempts: int = 5, backoff_seconds: float = 2.0):
def deco(fn):
def wrapped(*args, **kwargs):
attempt = 0
while True:
try:
return fn(*args, **kwargs)
except Exception:
attempt += 1
if attempt >= max_attempts:
raise
time.sleep(backoff_seconds * attempt)

return wrapped

return deco


@retry()
def generate_image(
prompt: str,
output_path: Path,
size: str,
quality: str,
timeout: int = 3600,
ref_b64: str | None = None,
) -> None:
api_key = os.environ.get("AZURE_OPENAI_IMAGE_API_KEY")
if not api_key:
raise RuntimeError("AZURE_OPENAI_IMAGE_API_KEY is not set")
ensure_dir(str(output_path.parent))
if ref_b64:
url = f"{os.getenv('AZURE_EU_OPENAI_URL')}/openai/deployments/gpt-image-1/images/edits?api-version=2025-04-01-preview"
img_bytes = base64.b64decode(ref_b64)
files = {
"image[]": ("reference.jpg", img_bytes, "image/jpeg"),
}
data = {
"prompt": prompt,
"size": size,
"quality": quality,
"n": "1",
"background": "transparent",
"output_compression": "100",
"output_format": "png",
}
headers = {"api-key": api_key}
response = requests.post(
url, headers=headers, files=files, data=data, timeout=timeout
)
response.raise_for_status()
b64_json = response.json()["data"][0]["b64_json"]
img = Image.open(BytesIO(base64.b64decode(b64_json)))
img.save(output_path, format="PNG")
else:
url = f"{os.getenv('AZURE_EU_OPENAI_URL')}/openai/deployments/gpt-image-1/images/edits?api-version=2025-04-01-preview"
headers = {"Content-Type": "application/json", "api-key": api_key}
body = {
"prompt": prompt,
"size": size,
"quality": quality,
"output_compression": 100,
"output_format": "png",
"n": 1,
"background": "transparent",
}
response = requests.post(url, headers=headers, json=body, timeout=timeout)
response.raise_for_status()
b64_json = response.json()["data"][0]["b64_json"]
img = Image.open(BytesIO(base64.b64decode(b64_json)))
img.save(output_path, format="PNG")


def format_stem_for_output(stem_name: str) -> str:
m = re.match(r"^(\d)_?f(\d{4})$", stem_name, re.I)
if m:
return f"{m.group(1)}f{m.group(2)}"
return stem_name


def sanitize_question(q: str) -> str:
s = re.sub(r"\s+", "", q)
s = s.replace("(", "").replace(")", "").replace("[", "").replace("]", "")
s = re.sub(r"[^A-Za-z0-9_-]+", "", s)
return s.lower() or "unknown"


def collect_page_image_b64_map(pages: List[Dict[str, Any]]) -> Dict[int, str]:
m: Dict[int, str] = {}
for r in pages:
p = int(r["page"])
with open(r["file"], "rb") as f:
m[p] = base64.b64encode(f.read()).decode("ascii")
return m


def generate_from_judgments(
base_out_dir: str,
stem_name: str,
prefix: str,
size: str,
quality: str,
overwrite: bool,
) -> None:
stem_dir = Path(base_out_dir)
judge_path = stem_dir / f"{stem_name}_judge.json"
if not judge_path.exists():
return
with open(judge_path, "r", encoding="utf-8") as f:
judgments = json.load(f)
stem_formatted = format_stem_for_output(stem_name)
for row in judgments:
llm = row.get("llm_judgment") or {}
if not isinstance(llm, dict):
continue
if llm.get("outcome") != "YES":
continue
page_num = int(row.get("page"))
items = llm.get("items") or []
for it in items:
q = sanitize_question(str(it.get("question", "unknown")))
regen = str(it.get("regenerate_prompt", "")).strip()
if not regen:
continue
filename = f"{stem_formatted}_{prefix}{str(page_num).zfill(3)}_{q}.png"
out_path = Path(base_out_dir) / filename
if out_path.exists() and not overwrite:
continue
generate_image(regen, out_path, size=size, quality=quality)


if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument(
"--pdf",
default="/data/gcse_maths_edexcel/1fjune2017.pdf",
)
ap.add_argument(
"--out",
default="/data/gcse_maths_edexcel/done",
)
ap.add_argument("--prefix", default="p")
ap.add_argument("--digits", type=int, default=3)
ap.add_argument("--start", type=int, default=None)
ap.add_argument("--end", type=int, default=None)
ap.add_argument("--dpi", type=int, default=300)
ap.add_argument("--quality", type=int, default=90)
ap.add_argument("--subsampling", type=int, choices=[-1, 0, 1, 2], default=0)
ap.add_argument("--progressive", action="store_true")
ap.add_argument("--optimize", action="store_true")
ap.add_argument("--overwrite", action="store_true")
ap.add_argument("--img-size", default="1024x1536")
ap.add_argument("--img-quality", default="high", choices=["high", "standard"])
ap.add_argument("--skip-judge", action="store_true")
parsed_args = ap.parse_args()

page_rows = split_to_jpegs(
parsed_args.pdf,
parsed_args.out,
parsed_args.prefix,
parsed_args.digits,
parsed_args.start,
parsed_args.end,
parsed_args.dpi,
parsed_args.quality,
parsed_args.subsampling,
parsed_args.progressive,
parsed_args.optimize,
parsed_args.overwrite,
)

if not parsed_args.skip_judge:
analyze_pages(page_rows, parsed_args.out)

if page_rows:
doc_stem = page_rows[0]["stem"]
generate_from_judgments(
base_out_dir=parsed_args.out,
stem_name=doc_stem,
prefix=parsed_args.prefix,
size=parsed_args.img_size,
quality=parsed_args.img_quality,
overwrite=parsed_args.overwrite,
)

Continue reading...