Files
pdf/gateway/app/services/convert/layout/images.py
T

504 lines
18 KiB
Python
Raw Normal View History

"""Page image / display-list helpers for classification and table rulings."""
from __future__ import annotations
import io
from contextlib import suppress
from dataclasses import dataclass
from typing import Any
# Identity transform, in PDF's [a b c d e f] form.
_IDENTITY = (1.0, 0.0, 0.0, 1.0, 0.0, 0.0)
# Below this on-page size (points) an image is furniture — a rule, a bullet, a
# spacer — not a figure worth carrying into the output.
MIN_FIGURE_POINTS = 12.0
def _mat_mul(m: tuple, n: tuple) -> tuple:
"""PDF matrix product m x n (row-vector convention)."""
a1, b1, c1, d1, e1, f1 = m
a2, b2, c2, d2, e2, f2 = n
return (
a1 * a2 + b1 * c2,
a1 * b2 + b1 * d2,
c1 * a2 + d1 * c2,
c1 * b2 + d1 * d2,
e1 * a2 + f1 * c2 + e2,
e1 * b2 + f1 * d2 + f2,
)
def _apply(m: tuple, x: float, y: float) -> tuple[float, float]:
a, b, c, d, e, f = m
return (a * x + c * y + e, b * x + d * y + f)
@dataclass
class ImagePlacement:
"""Where an image XObject is actually painted on the page.
``name`` is the XObject resource name, so the placement can be matched back
to the image data. Coordinates are PDF user space (origin bottom-left).
"""
name: str
x: float
y: float
w: float
h: float
px_w: int = 0
px_h: int = 0
@property
def area(self) -> float:
return max(self.w, 0.0) * max(self.h, 0.0)
def image_placements(page) -> list[ImagePlacement]:
"""True on-page rectangles for every image drawn by this page.
Memoised on the page object: the coverage estimate and the figure
extractor both need it, and walking the content stream twice per page is
pure waste.
An image XObject is painted into the unit square, so the CTM in effect at
its ``Do`` operator *is* its placement. Walking ``q``/``Q``/``cm``/``Do``
recovers that exactly.
This replaces guessing placement from pixel dimensions, which produced a
fabricated rectangle: a 1200x1200 logo and a full-page scan have similar
pixel counts but completely different footprints, and treating the logo as
a full-page raster silently dropped it from the output.
"""
cached = getattr(page, "_dq_image_placements", None)
if cached is not None:
return cached
out: list[ImagePlacement] = []
try:
from pypdf.generic import ContentStream
resources = page.get("/Resources")
xobjects = {}
if resources is not None:
xo = resources.get_object().get("/XObject")
if xo is not None:
xobjects = xo.get_object()
if not xobjects:
return out
contents = page.get_contents()
if contents is None:
return out
stream = ContentStream(contents, getattr(page, "pdf", None))
ctm = _IDENTITY
stack: list[tuple] = []
for operands, operator in stream.operations:
if operator == b"q":
stack.append(ctm)
elif operator == b"Q":
ctm = stack.pop() if stack else _IDENTITY
elif operator == b"cm" and len(operands) >= 6:
try:
ctm = _mat_mul(tuple(float(v) for v in operands[:6]), ctm)
except (TypeError, ValueError):
continue
elif operator == b"Do" and operands:
name = str(operands[0])
obj = xobjects.get(name)
if obj is None:
continue
try:
obj = obj.get_object()
except Exception:
continue
if obj.get("/Subtype") != "/Image":
continue
corners = [_apply(ctm, cx, cy) for cx, cy in ((0, 0), (1, 0), (1, 1), (0, 1))]
xs = [c[0] for c in corners]
ys = [c[1] for c in corners]
out.append(
ImagePlacement(
name=name,
x=min(xs),
y=min(ys),
w=max(xs) - min(xs),
h=max(ys) - min(ys),
px_w=int(obj.get("/Width", 0) or 0),
px_h=int(obj.get("/Height", 0) or 0),
)
)
except Exception:
return out
with suppress(Exception):
page._dq_image_placements = out
return out
def image_blocks_from_pypdf_page(page) -> list[dict]:
"""Image footprints on a pypdf page, in PDF user-space units.
Uses the real content-stream placement, falling back to the old pixel-ratio
approximation only when the content stream cannot be walked.
"""
placements = image_placements(page)
if placements:
return [
{"x": p.x, "y": p.y, "w": p.w, "h": p.h, "width": p.w, "height": p.h, "name": p.name}
for p in placements
]
blocks: list[dict] = []
try:
if "/Resources" not in page or "/XObject" not in page["/Resources"]:
return blocks
xobjects = page["/Resources"]["/XObject"].get_object()
mediabox = page.mediabox
page_w = float(mediabox.width)
page_h = float(mediabox.height)
for _name in xobjects:
obj = xobjects[_name]
if obj.get("/Subtype") != "/Image":
continue
w = float(obj.get("/Width", page_w))
h = float(obj.get("/Height", page_h))
# Without CTM we approximate: large images cover most of the page
scale = min(page_w / max(w, 1), page_h / max(h, 1))
blocks.append({"w": w * scale, "h": h * scale, "width": w * scale, "height": h * scale})
except Exception:
return blocks
return blocks
def image_blocks_from_engine_page(page: Any) -> list[dict]:
blocks: list[dict] = []
for meth in ("extract_images", "list_images", "get_images"):
if not hasattr(page, meth):
continue
try:
raw = getattr(page, meth)()
except Exception:
continue
if not raw:
continue
for item in raw:
if isinstance(item, dict):
blocks.append(item)
else:
blocks.append(
{
"w": float(getattr(item, "w", getattr(item, "width", 0)) or 0),
"h": float(getattr(item, "h", getattr(item, "height", 0)) or 0),
}
)
if blocks:
return blocks
return blocks
def display_list_ops_from_engine_page(page: Any) -> list[dict]:
if not hasattr(page, "extract_display_list"):
return []
try:
raw = page.extract_display_list()
except Exception:
return []
if isinstance(raw, str):
raw_str = raw.strip()
if raw_str.startswith(("[", "{")):
try:
import json
raw = json.loads(raw_str)
except Exception:
pass
if isinstance(raw, list):
ops: list[dict] = []
for x in raw:
if isinstance(x, dict):
op_dict = dict(x)
if "args" in op_dict and isinstance(op_dict["args"], (list, tuple)):
args = op_dict["args"]
if len(args) >= 4 and str(op_dict.get("op") or "").lower() in ("re", "rect"):
try:
rx, ry, rw, rh = float(args[0]), float(args[1]), float(args[2]), float(args[3])
op_dict.setdefault("x", rx)
op_dict.setdefault("y", ry)
op_dict.setdefault("w", rw)
op_dict.setdefault("h", rh)
op_dict.setdefault("x0", rx)
op_dict.setdefault("y0", ry)
op_dict.setdefault("x1", rx + rw)
op_dict.setdefault("y1", ry + rh)
except (ValueError, TypeError):
pass
ops.append(op_dict)
else:
ops.append({"type": "line"})
return ops
if isinstance(raw, str):
# Best-effort parse of simple "line x0 y0 x1 y1" lines
ops: list[dict] = []
for line in raw.splitlines():
parts = line.strip().split()
if len(parts) >= 5 and parts[0].lower() in ("line", "vline"):
try:
ops.append(
{
"type": parts[0].lower(),
"x0": float(parts[1]),
"y0": float(parts[2]),
"x1": float(parts[3]),
"y1": float(parts[4]),
}
)
except ValueError:
continue
return ops
return []
def _colorspace_mode(cs) -> str | None:
"""Map PDF ColorSpace to Pillow mode when possible."""
if cs is None:
return "RGB"
name = str(cs)
if "DeviceGray" in name or name.endswith("/DeviceGray"):
return "L"
if "DeviceRGB" in name or name.endswith("/DeviceRGB"):
return "RGB"
if "DeviceCMYK" in name or name.endswith("/DeviceCMYK"):
return "CMYK"
# ICCBased / Indexed: try RGB bytes if Length matches later
if "ICCBased" in name:
return "RGB"
return None
def _decode_xobject_png(obj) -> bytes | None:
"""One image XObject as PNG bytes, or None when it cannot be decoded."""
from PIL import Image
from pypdf.generic import NameObject
width = int(obj.get("/Width", 0))
height = int(obj.get("/Height", 0))
try:
data = obj.get_data()
except Exception:
return None
filt = obj.get("/Filter")
img = None
try:
if filt == "/DCTDecode" or (isinstance(filt, list) and NameObject("/DCTDecode") in filt):
img = Image.open(io.BytesIO(data))
elif filt == "/FlateDecode" or str(filt) == "/FlateDecode" or (
isinstance(filt, list) and any("Flate" in str(f) for f in filt)
):
mode = _colorspace_mode(obj.get("/ColorSpace")) or "RGB"
expected = width * height * (1 if mode == "L" else 4 if mode == "CMYK" else 3)
if len(data) < expected:
# Sometimes get_data already expands; try open as encoded
try:
img = Image.open(io.BytesIO(data))
except Exception:
return None
else:
img = Image.frombytes(mode, (width, height), data[:expected])
else:
img = Image.open(io.BytesIO(data))
except Exception:
return None
if img is None:
return None
buf = io.BytesIO()
img.convert("RGB").save(buf, format="PNG")
return buf.getvalue()
def decode_page_images(
page, *, warnings: list[str] | None = None, only: set[str] | None = None
) -> dict[str, bytes]:
"""Decodable image XObjects on the page, keyed by resource name.
Keyed by name so a placement recovered from the content stream can be
matched to its pixels: a page with a chart and a photo needs both, and the
old "largest image only" extraction returned one of them.
``only`` restricts decoding to the named XObjects. Decoding is the
expensive part, so a caller that already knows which images it will keep
should say so rather than decoding a full-page raster to discard it.
"""
out: dict[str, bytes] = {}
try:
if "/Resources" not in page or "/XObject" not in page["/Resources"]:
return out
xobjects = page["/Resources"]["/XObject"].get_object()
image_count = 0
for name in xobjects:
if only is not None and str(name) not in only:
continue
obj = xobjects[name]
if obj.get("/Subtype") != "/Image":
continue
image_count += 1
png = _decode_xobject_png(obj)
if png:
out[str(name)] = png
if not out and image_count and warnings is not None:
warnings.append(
f"page has {image_count} image XObject(s) but decode failed "
f"({image_count} failure(s))."
)
except Exception as exc:
if warnings is not None:
warnings.append(f"embedded image extract failed: {exc}")
return out
def embedded_page_image_png(page, *, warnings: list[str] | None = None) -> bytes | None:
"""Largest embedded image on a pypdf page as PNG bytes, if any.
"Largest" is by pixel area, as the OCR path needs the highest-resolution
raster — not by encoded size, which favours noisy images over big ones.
"""
images = decode_page_images(page, warnings=warnings)
if not images:
return None
def _area(png: bytes) -> int:
try:
from PIL import Image
with Image.open(io.BytesIO(png)) as im:
return im.width * im.height
except Exception:
return len(png)
return max(images.values(), key=_area)
def figure_blocks_from_pypdf_page(
page,
*,
reading_order: int = 0,
warnings: list[str] | None = None,
max_page_coverage: float = 0.38,
max_bytes: int = 450_000,
) -> list:
"""Extract embedded images as IDM figure blocks (digital path).
Every image on the page is emitted, at the rectangle the content stream
actually paints it into. Two things this fixes:
*multiple figures per page*
The previous implementation extracted only the single largest image, so
a page with a chart and a photograph lost one of them outright.
*real placement instead of a fabricated box*
Position and size came from a hard-coded guess, which destroyed both
reading order relative to the text and the aspect ratio. Coverage was
likewise inferred from a pixel-count ladder, so a high-resolution logo
was mistaken for a full-page scan and dropped.
Near-full-page rasters are still skipped: those belong to OCR/scan handling
and would otherwise embed a multi-MB page image per page.
"""
from app.services.convert.idm.model import BBox, Block, BlockType
out: list = []
try:
mediabox = page.mediabox
pw, ph = float(mediabox.width), float(mediabox.height)
page_area = max(pw * ph, 1.0)
placements = image_placements(page)
if not placements:
# No usable content stream: fall back to the single largest image
# with a conservative centred box rather than dropping everything.
png = embedded_page_image_png(page, warnings=warnings)
if not png or len(png) > max_bytes:
return out
out.append(
Block(
type=BlockType.figure,
text="[Image]",
bbox=BBox(x=72, y=ph * 0.3, w=min(pw - 144, 400), h=min(ph * 0.4, 300)),
image_png=png,
reading_order=reading_order,
)
)
return out
# Decide from geometry *before* decoding. Decoding a full-page raster
# only to discard it costs seconds per page on an image-heavy document,
# and the placement already says it will be discarded.
keep: list[ImagePlacement] = []
skipped_large = 0
for placement in placements:
if placement.w < MIN_FIGURE_POINTS or placement.h < MIN_FIGURE_POINTS:
# Rules, bullets and spacer pixels are not figures.
continue
if placement.area / page_area >= max_page_coverage:
skipped_large += 1
continue
keep.append(placement)
if keep:
wanted = {p.name for p in keep}
images = decode_page_images(page, warnings=warnings, only=wanted)
# Top-to-bottom, then left-to-right: the order a reader meets them.
keep.sort(key=lambda p: (-(p.y + p.h), p.x))
for offset, placement in enumerate(keep):
png = images.get(placement.name)
if not png:
continue
if len(png) > max_bytes:
skipped_large += 1
continue
out.append(
Block(
type=BlockType.figure,
text="[Image]",
bbox=BBox(x=placement.x, y=placement.y, w=placement.w, h=placement.h),
image_png=png,
reading_order=reading_order + offset,
)
)
if skipped_large and warnings is not None:
warnings.append(
f"Skipped {skipped_large} large/full-page embedded image(s); prefer OCR text."
)
except Exception as exc:
if warnings is not None:
warnings.append(f"figure block extract failed: {exc}")
return out
return out
def figure_blocks_from_engine_images(
img_meta: list[dict], png_blobs: list[bytes], *, reading_order: int = 0
) -> list:
from app.services.convert.idm.model import BBox, Block, BlockType
out: list = []
for i, png in enumerate(png_blobs):
meta = img_meta[i] if i < len(img_meta) else {}
w = float(meta.get("w") or meta.get("width") or 300)
h = float(meta.get("h") or meta.get("height") or 200)
x = float(meta.get("x", 72))
y = float(meta.get("y", 200))
out.append(
Block(
type=BlockType.figure,
text=f"[Image {i + 1}]",
bbox=BBox(x=x, y=y, w=w, h=h),
image_png=png,
reading_order=reading_order + i,
)
)
return out