mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-24 21:41:04 +02:00
fix: restore PyPDF2 as the PDF text extractor in the SDK path
Text extraction (and therefore tree building) now matches the CLI / pre-SDK default. Image extraction keeps PyMuPDF — the only part that needs it — with  references appended per page.
This commit is contained in:
parent
b2b0901579
commit
f2b407f62f
2 changed files with 63 additions and 70 deletions
|
|
@ -205,7 +205,7 @@ Omitting `doc_ids` queries the **entire collection** and lets the agent pick whi
|
||||||
|
|
||||||
# ⚙️ Package Usage
|
# ⚙️ Package Usage
|
||||||
|
|
||||||
> **Note:** This package uses standard PDF parsing. The CLI parses PDFs with PyPDF2 (as in previous releases), while the SDK's local mode parses with PyMuPDF and also extracts images — extracted text can differ slightly, so the two entry points may produce different trees for the same PDF. For use cases with complex PDFs, our [cloud service](https://pageindex.ai/developer) (via MCP and API) offers enhanced OCR, tree building, and retrieval.
|
> **Note:** This package uses standard PDF parsing. For use cases with complex PDFs, our [cloud service](https://pageindex.ai/developer) (via MCP and API) offers enhanced OCR, tree building, and retrieval.
|
||||||
|
|
||||||
You can follow these steps to generate a PageIndex tree from a PDF document.
|
You can follow these steps to generate a PageIndex tree from a PDF document.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,4 @@
|
||||||
|
import PyPDF2
|
||||||
import pymupdf
|
import pymupdf
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from .protocol import ContentNode, ParsedDocument
|
from .protocol import ContentNode, ParsedDocument
|
||||||
|
|
@ -15,39 +16,42 @@ class PdfParser:
|
||||||
path = Path(file_path)
|
path = Path(file_path)
|
||||||
model = kwargs.get("model")
|
model = kwargs.get("model")
|
||||||
images_dir = kwargs.get("images_dir")
|
images_dir = kwargs.get("images_dir")
|
||||||
|
|
||||||
|
# Images are extracted with PyMuPDF (PyPDF2 cannot); text stays with
|
||||||
|
# PyPDF2 below so the extracted text — and therefore the tree — matches
|
||||||
|
# the CLI / pre-SDK default.
|
||||||
|
page_images: dict[int, list[dict]] = {}
|
||||||
|
if images_dir:
|
||||||
|
with pymupdf.open(str(path)) as doc:
|
||||||
|
for i, page in enumerate(doc):
|
||||||
|
images = self._extract_page_images(page, i + 1, images_dir)
|
||||||
|
if images:
|
||||||
|
page_images[i + 1] = images
|
||||||
|
|
||||||
|
reader = PyPDF2.PdfReader(str(path))
|
||||||
nodes = []
|
nodes = []
|
||||||
|
for i, page in enumerate(reader.pages):
|
||||||
|
page_num = i + 1
|
||||||
|
content = page.extract_text() or ""
|
||||||
|
images = page_images.get(page_num)
|
||||||
|
if images:
|
||||||
|
refs = "\n".join(f"" for img in images)
|
||||||
|
content = f"{content}\n{refs}" if content else refs
|
||||||
|
|
||||||
with pymupdf.open(str(path)) as doc:
|
tokens = count_tokens(content, model=model)
|
||||||
page_count = doc.page_count
|
nodes.append(ContentNode(
|
||||||
for i, page in enumerate(doc):
|
content=content,
|
||||||
page_num = i + 1
|
tokens=tokens,
|
||||||
if images_dir:
|
index=page_num,
|
||||||
content, images = self._extract_page_with_images(
|
images=images,
|
||||||
doc, page, page_num, images_dir)
|
))
|
||||||
else:
|
|
||||||
content = page.get_text()
|
|
||||||
images = None
|
|
||||||
|
|
||||||
tokens = count_tokens(content, model=model)
|
|
||||||
nodes.append(ContentNode(
|
|
||||||
content=content or "",
|
|
||||||
tokens=tokens,
|
|
||||||
index=page_num,
|
|
||||||
images=images if images else None,
|
|
||||||
))
|
|
||||||
|
|
||||||
return ParsedDocument(doc_name=path.stem, nodes=nodes,
|
return ParsedDocument(doc_name=path.stem, nodes=nodes,
|
||||||
metadata={"page_count": page_count})
|
metadata={"page_count": len(reader.pages)})
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _extract_page_with_images(doc, page, page_num: int,
|
def _extract_page_images(page, page_num: int, images_dir: str) -> list[dict]:
|
||||||
images_dir: str) -> tuple[str, list[dict]]:
|
"""Save a page's images to disk and return their metadata."""
|
||||||
"""Extract text and images from a page, preserving their relative order.
|
|
||||||
|
|
||||||
Uses get_text("dict") to iterate blocks in reading order.
|
|
||||||
Text blocks become text; image blocks are saved to disk and replaced
|
|
||||||
with an inline placeholder: 
|
|
||||||
"""
|
|
||||||
images_path = Path(images_dir)
|
images_path = Path(images_dir)
|
||||||
images_path.mkdir(parents=True, exist_ok=True)
|
images_path.mkdir(parents=True, exist_ok=True)
|
||||||
# Store an absolute path so the  reference resolves
|
# Store an absolute path so the  reference resolves
|
||||||
|
|
@ -55,53 +59,42 @@ class PdfParser:
|
||||||
# break as soon as the query runs from a different directory.)
|
# break as soon as the query runs from a different directory.)
|
||||||
abs_images_path = images_path.resolve()
|
abs_images_path = images_path.resolve()
|
||||||
|
|
||||||
parts: list[str] = []
|
|
||||||
images: list[dict] = []
|
images: list[dict] = []
|
||||||
img_idx = 0
|
img_idx = 0
|
||||||
|
|
||||||
for block in page.get_text("dict")["blocks"]:
|
for block in page.get_text("dict")["blocks"]:
|
||||||
if block["type"] == 0: # text block
|
if block["type"] != 1: # image blocks only
|
||||||
lines = []
|
continue
|
||||||
for line in block["lines"]:
|
width = block.get("width", 0)
|
||||||
spans_text = "".join(span["text"] for span in line["spans"])
|
height = block.get("height", 0)
|
||||||
lines.append(spans_text)
|
if width < _MIN_IMAGE_SIZE or height < _MIN_IMAGE_SIZE:
|
||||||
parts.append("\n".join(lines))
|
continue
|
||||||
|
|
||||||
elif block["type"] == 1: # image block
|
image_bytes = block.get("image")
|
||||||
width = block.get("width", 0)
|
if not image_bytes:
|
||||||
height = block.get("height", 0)
|
continue
|
||||||
if width < _MIN_IMAGE_SIZE or height < _MIN_IMAGE_SIZE:
|
|
||||||
continue
|
|
||||||
|
|
||||||
image_bytes = block.get("image")
|
try:
|
||||||
if not image_bytes:
|
pix = pymupdf.Pixmap(image_bytes)
|
||||||
continue
|
# n includes the alpha channel, so a plain RGBA pixmap also
|
||||||
|
# has n==4 — subtract alpha before comparing. Without this,
|
||||||
|
# a CMYK image with no alpha (n==4, same as RGBA) skips the
|
||||||
|
# RGB conversion, and pix.save() as .png then raises
|
||||||
|
# "unsupported colorspace for 'png'", silently dropping the
|
||||||
|
# image via the bare except below.
|
||||||
|
if pix.n - pix.alpha >= 4:
|
||||||
|
pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
|
||||||
|
filename = f"p{page_num}_img{img_idx}.png"
|
||||||
|
pix.save(str(images_path / filename))
|
||||||
|
pix = None
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
|
||||||
try:
|
images.append({
|
||||||
pix = pymupdf.Pixmap(image_bytes)
|
"path": str(abs_images_path / filename),
|
||||||
# n includes the alpha channel, so a plain RGBA pixmap also
|
"width": width,
|
||||||
# has n==4 — subtract alpha before comparing. Without this,
|
"height": height,
|
||||||
# a CMYK image with no alpha (n==4, same as RGBA) skips the
|
})
|
||||||
# RGB conversion, and pix.save() as .png then raises
|
img_idx += 1
|
||||||
# "unsupported colorspace for 'png'", silently dropping the
|
|
||||||
# image via the bare except below.
|
|
||||||
if pix.n - pix.alpha >= 4:
|
|
||||||
pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
|
|
||||||
filename = f"p{page_num}_img{img_idx}.png"
|
|
||||||
save_path = images_path / filename
|
|
||||||
pix.save(str(save_path))
|
|
||||||
pix = None
|
|
||||||
except Exception:
|
|
||||||
continue
|
|
||||||
|
|
||||||
img_path = str(abs_images_path / filename)
|
return images
|
||||||
images.append({
|
|
||||||
"path": img_path,
|
|
||||||
"width": width,
|
|
||||||
"height": height,
|
|
||||||
})
|
|
||||||
parts.append(f"")
|
|
||||||
img_idx += 1
|
|
||||||
|
|
||||||
content = "\n".join(parts)
|
|
||||||
return content, images
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue