1067 lines
39 KiB
Python
1067 lines
39 KiB
Python
# SPDX-FileCopyrightText: 2018-2022 James R. Barlow
|
|
# SPDX-FileCopyrightText: 2019 Martin Wind
|
|
# SPDX-License-Identifier: MPL-2.0
|
|
|
|
"""OCRmyPDF page processing pipeline functions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import re
|
|
import sys
|
|
from collections.abc import Iterable, Iterator, Sequence
|
|
from contextlib import suppress
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
from shutil import copyfileobj
|
|
from typing import Any, BinaryIO, TypeVar, cast
|
|
|
|
import img2pdf
|
|
import pikepdf
|
|
from PIL import Image, ImageColor, ImageDraw
|
|
|
|
from ocrmypdf._concurrent import Executor
|
|
from ocrmypdf._exec import unpaper
|
|
from ocrmypdf._jobcontext import PageContext, PdfContext
|
|
from ocrmypdf._metadata import repair_docinfo_nuls
|
|
from ocrmypdf.exceptions import (
|
|
DigitalSignatureError,
|
|
DpiError,
|
|
EncryptedPdfError,
|
|
InputFileError,
|
|
PriorOcrFoundError,
|
|
TaggedPDFError,
|
|
UnsupportedImageFormatError,
|
|
)
|
|
from ocrmypdf.helpers import IMG2PDF_KWARGS, Resolution, safe_symlink
|
|
from ocrmypdf.hocrtransform import DebugRenderOptions, HocrTransform
|
|
from ocrmypdf.hocrtransform._font import Courier
|
|
from ocrmypdf.pdfa import generate_pdfa_ps
|
|
from ocrmypdf.pdfinfo import Colorspace, Encoding, PageInfo, PdfInfo
|
|
from ocrmypdf.pluginspec import OrientationConfidence
|
|
|
|
T = TypeVar("T")
|
|
log = logging.getLogger(__name__)
|
|
|
|
VECTOR_PAGE_DPI = 400
|
|
|
|
|
|
def triage_image_file(input_file: Path, output_file: Path, options) -> None:
|
|
"""Triage the input image file.
|
|
|
|
If the input file is an image, check its resolution and convert it to PDF.
|
|
|
|
Args:
|
|
input_file: The path to the input file.
|
|
output_file: The path to the output file.
|
|
options: An object containing the options passed to the OCRmyPDF command.
|
|
|
|
Raises:
|
|
UnsupportedImageFormatError: If the input file is not a supported image format.
|
|
DpiError: If the input image has no resolution (DPI) in its metadata or if the
|
|
resolution is not credible.
|
|
"""
|
|
log.info("Input file is not a PDF, checking if it is an image...")
|
|
try:
|
|
im = Image.open(input_file)
|
|
except OSError as e:
|
|
# Recover the original filename
|
|
log.error(str(e).replace(str(input_file), str(options.input_file)))
|
|
raise UnsupportedImageFormatError() from e
|
|
|
|
with im:
|
|
log.info("Input file is an image")
|
|
if 'dpi' in im.info:
|
|
if im.info['dpi'] <= (96, 96) and not options.image_dpi:
|
|
log.info("Image size: (%d, %d)", *im.size)
|
|
log.info("Image resolution: (%d, %d)", *im.info['dpi'])
|
|
raise DpiError(
|
|
"Input file is an image, but the resolution (DPI) is "
|
|
"not credible. Estimate the resolution at which the "
|
|
"image was scanned and specify it using --image-dpi."
|
|
)
|
|
elif not options.image_dpi:
|
|
log.info("Image size: (%d, %d)", *im.size)
|
|
raise DpiError(
|
|
"Input file is an image, but has no resolution (DPI) "
|
|
"in its metadata. Estimate the resolution at which "
|
|
"image was scanned and specify it using --image-dpi."
|
|
)
|
|
|
|
if im.mode in ('RGBA', 'LA'):
|
|
raise UnsupportedImageFormatError(
|
|
"The input image has an alpha channel. Remove the alpha "
|
|
"channel first."
|
|
)
|
|
|
|
if 'iccprofile' not in im.info:
|
|
if im.mode == 'RGB':
|
|
log.info("Input image has no ICC profile, assuming sRGB")
|
|
elif im.mode == 'CMYK':
|
|
raise UnsupportedImageFormatError(
|
|
"Input CMYK image has no ICC profile, not usable"
|
|
)
|
|
|
|
try:
|
|
log.info("Image seems valid. Try converting to PDF...")
|
|
layout_fun = img2pdf.default_layout_fun
|
|
if options.image_dpi:
|
|
layout_fun = img2pdf.get_fixed_dpi_layout_fun(
|
|
Resolution(options.image_dpi, options.image_dpi)
|
|
)
|
|
with open(output_file, 'wb') as outf:
|
|
img2pdf.convert(
|
|
os.fspath(input_file),
|
|
layout_fun=layout_fun,
|
|
outputstream=outf,
|
|
**IMG2PDF_KWARGS,
|
|
)
|
|
log.info("Successfully converted to PDF, processing...")
|
|
except img2pdf.ImageOpenError as e:
|
|
raise UnsupportedImageFormatError() from e
|
|
|
|
|
|
def _pdf_guess_version(input_file: Path, search_window=1024) -> str:
|
|
"""Try to find version signature at start of file.
|
|
|
|
Not robust enough to deal with appended files.
|
|
|
|
Returns empty string if not found, indicating file is probably not PDF.
|
|
"""
|
|
with open(input_file, 'rb') as f:
|
|
signature = f.read(search_window)
|
|
m = re.search(br'%PDF-(\d\.\d)', signature)
|
|
if m:
|
|
return m.group(1).decode('ascii')
|
|
return ''
|
|
|
|
|
|
def triage(
|
|
original_filename: str, input_file: Path, output_file: Path, options
|
|
) -> Path:
|
|
"""Triage the input file. We can handle PDFs and images."""
|
|
try:
|
|
if _pdf_guess_version(input_file):
|
|
if options.image_dpi:
|
|
log.warning(
|
|
"Argument --image-dpi is being ignored because the "
|
|
"input file is a PDF, not an image."
|
|
)
|
|
# Origin file is a pdf create a symlink with pdf extension
|
|
safe_symlink(input_file, output_file)
|
|
return output_file
|
|
except OSError as e:
|
|
log.debug(f"Temporary file was at: {input_file}")
|
|
msg = str(e).replace(str(input_file), original_filename)
|
|
raise InputFileError(msg) from e
|
|
|
|
triage_image_file(input_file, output_file, options)
|
|
return output_file
|
|
|
|
|
|
def get_pdfinfo(
|
|
input_file,
|
|
*,
|
|
executor: Executor,
|
|
detailed_analysis: bool = False,
|
|
progbar: bool = False,
|
|
max_workers: int | None = None,
|
|
use_threads: bool = True,
|
|
check_pages=None,
|
|
) -> PdfInfo:
|
|
"""Get the PDF info."""
|
|
try:
|
|
return PdfInfo(
|
|
input_file,
|
|
detailed_analysis=detailed_analysis,
|
|
progbar=progbar,
|
|
max_workers=max_workers,
|
|
use_threads=use_threads,
|
|
check_pages=check_pages,
|
|
executor=executor,
|
|
)
|
|
except pikepdf.PasswordError as e:
|
|
raise EncryptedPdfError() from e
|
|
except pikepdf.PdfError as e:
|
|
raise InputFileError() from e
|
|
|
|
|
|
def validate_pdfinfo_options(context: PdfContext) -> None:
|
|
"""Validate the PDF info options."""
|
|
pdfinfo = context.pdfinfo
|
|
options = context.options
|
|
|
|
if pdfinfo.needs_rendering:
|
|
raise InputFileError(
|
|
"This PDF contains dynamic XFA forms created by Adobe LiveCycle "
|
|
"Designer and can only be read by Adobe Acrobat or Adobe Reader."
|
|
)
|
|
if pdfinfo.has_signature:
|
|
if options.invalidate_digital_signatures:
|
|
log.warning("All digital signatures will be invalidated")
|
|
else:
|
|
raise DigitalSignatureError()
|
|
if pdfinfo.has_acroform:
|
|
if options.redo_ocr:
|
|
raise InputFileError(
|
|
"This PDF has a user fillable form. --redo-ocr is not "
|
|
"currently possible on such files."
|
|
)
|
|
else:
|
|
log.warning(
|
|
"This PDF has a fillable form. "
|
|
"Chances are it is a pure digital "
|
|
"document that does not need OCR."
|
|
)
|
|
if not options.force_ocr:
|
|
log.info(
|
|
"Use the option --force-ocr to produce an image of the "
|
|
"form and all filled form fields. The output PDF will be "
|
|
"'flattened' and will no longer be fillable."
|
|
)
|
|
if pdfinfo.is_tagged:
|
|
if options.force_ocr or options.skip_text or options.redo_ocr:
|
|
log.warning(
|
|
"This PDF is marked as a Tagged PDF. This often indicates "
|
|
"that the PDF was generated from an office document and does "
|
|
"not need OCR. PDF pages processed by OCRmyPDF may not be "
|
|
"tagged correctly."
|
|
)
|
|
else:
|
|
raise TaggedPDFError()
|
|
context.plugin_manager.hook.validate(pdfinfo=pdfinfo, options=options)
|
|
|
|
|
|
def _vector_page_dpi(pageinfo: PageInfo) -> int:
|
|
"""Get a DPI to use for vector pages, if the page has vector content."""
|
|
return VECTOR_PAGE_DPI if pageinfo.has_vector or pageinfo.has_text else 0
|
|
|
|
|
|
def get_page_square_dpi(
|
|
page_context: PageContext, image_dpi: Resolution | None = None
|
|
) -> Resolution:
|
|
"""Get the DPI when we require xres == yres, scaled to physical units.
|
|
|
|
Page DPI includes UserUnit scaling.
|
|
"""
|
|
pageinfo = page_context.pageinfo
|
|
options = page_context.options
|
|
if not image_dpi:
|
|
image_dpi = pageinfo.dpi
|
|
xres = image_dpi.x or 0.0
|
|
yres = image_dpi.y or 0.0
|
|
userunit = float(pageinfo.userunit) or 1.0
|
|
units = float(
|
|
max(
|
|
(xres * userunit) or VECTOR_PAGE_DPI,
|
|
(yres * userunit) or VECTOR_PAGE_DPI,
|
|
_vector_page_dpi(pageinfo),
|
|
options.oversample or 0.0,
|
|
)
|
|
)
|
|
return Resolution(units, units)
|
|
|
|
|
|
def get_canvas_square_dpi(
|
|
page_context: PageContext, image_dpi: Resolution | None = None
|
|
) -> Resolution:
|
|
"""Get the DPI when we require xres == yres, in Postscript units.
|
|
|
|
Canvas DPI is independent of PDF UserUnit scaling, which is
|
|
used to describe situations where the PDF user space is not 1:1 with
|
|
the physical units of the page.
|
|
"""
|
|
pageinfo = page_context.pageinfo
|
|
options = page_context.options
|
|
if not image_dpi:
|
|
image_dpi = pageinfo.dpi
|
|
units = float(
|
|
max(
|
|
image_dpi.x or VECTOR_PAGE_DPI,
|
|
image_dpi.y or VECTOR_PAGE_DPI,
|
|
_vector_page_dpi(pageinfo),
|
|
options.oversample or 0.0,
|
|
)
|
|
)
|
|
return Resolution(units, units)
|
|
|
|
|
|
def is_ocr_required(page_context: PageContext) -> bool:
|
|
"""Check if the page needs to be OCR'd."""
|
|
pageinfo = page_context.pageinfo
|
|
options = page_context.options
|
|
|
|
ocr_required = True
|
|
|
|
if options.pages and pageinfo.pageno not in options.pages:
|
|
log.debug(f"skipped {pageinfo.pageno} as requested by --pages {options.pages}")
|
|
ocr_required = False
|
|
elif pageinfo.has_text:
|
|
if not options.force_ocr and not (options.skip_text or options.redo_ocr):
|
|
raise PriorOcrFoundError(
|
|
"page already has text! - aborting (use --force-ocr to force OCR; "
|
|
" see also help for the arguments --skip-text and --redo-ocr"
|
|
)
|
|
elif options.force_ocr:
|
|
log.info("page already has text! - rasterizing text and running OCR anyway")
|
|
ocr_required = True
|
|
elif options.redo_ocr:
|
|
if pageinfo.has_corrupt_text:
|
|
log.warning(
|
|
"some text on this page cannot be mapped to characters: "
|
|
"consider using --force-ocr instead"
|
|
)
|
|
else:
|
|
log.info("redoing OCR")
|
|
ocr_required = True
|
|
elif options.skip_text:
|
|
log.info("skipping all processing on this page")
|
|
ocr_required = False
|
|
elif not pageinfo.images and not options.lossless_reconstruction:
|
|
# We found a page with no images and no text. That means it may
|
|
# have vector art that the user wants to OCR. If we determined
|
|
# lossless reconstruction is not possible then we have to rasterize
|
|
# the image. So if OCR is being forced, take that to mean YES, go
|
|
# ahead and rasterize. If not forced, then pretend there's no text
|
|
# on the page at all so we don't lose anything.
|
|
# This could be made smarter by explicitly searching for vector art.
|
|
if options.force_ocr and options.oversample:
|
|
# The user really wants to reprocess this file
|
|
log.info(
|
|
"page has no images - "
|
|
f"rasterizing at {options.oversample} DPI because "
|
|
"--force-ocr --oversample was specified"
|
|
)
|
|
elif options.force_ocr:
|
|
# Warn the user they might not want to do this
|
|
log.warning(
|
|
"page has no images - "
|
|
"all vector content will be "
|
|
f"rasterized at {VECTOR_PAGE_DPI} DPI, losing some resolution and "
|
|
"likely increasing file size. Use --oversample to adjust the "
|
|
"DPI."
|
|
)
|
|
else:
|
|
log.info(
|
|
"page has no images - "
|
|
"skipping all processing on this page to avoid losing detail. "
|
|
"Use --force-ocr if you wish to perform OCR on pages that "
|
|
"have vector content."
|
|
)
|
|
ocr_required = False
|
|
|
|
if ocr_required and options.skip_big and pageinfo.images:
|
|
pixel_count = pageinfo.width_pixels * pageinfo.height_pixels
|
|
if pixel_count > (options.skip_big * 1_000_000):
|
|
ocr_required = False
|
|
log.warning(
|
|
"page too big, skipping OCR "
|
|
f"({(pixel_count / 1_000_000):.1f} MPixels > "
|
|
f"{options.skip_big:.1f} MPixels --skip-big)"
|
|
)
|
|
return ocr_required
|
|
|
|
|
|
def rasterize_preview(input_file: Path, page_context: PageContext) -> Path:
|
|
"""Generate a lower quality preview image."""
|
|
output_file = page_context.get_path('rasterize_preview.jpg')
|
|
canvas_dpi = Resolution(300.0, 300.0).take_min(
|
|
[get_canvas_square_dpi(page_context)]
|
|
)
|
|
page_dpi = Resolution(300.0, 300.0).take_min([get_page_square_dpi(page_context)])
|
|
page_context.plugin_manager.hook.rasterize_pdf_page(
|
|
input_file=input_file,
|
|
output_file=output_file,
|
|
raster_device='jpeggray',
|
|
raster_dpi=canvas_dpi,
|
|
pageno=page_context.pageinfo.pageno + 1,
|
|
page_dpi=page_dpi,
|
|
rotation=0,
|
|
filter_vector=False,
|
|
stop_on_soft_error=not page_context.options.continue_on_soft_render_error,
|
|
)
|
|
return output_file
|
|
|
|
|
|
def describe_rotation(
|
|
page_context: PageContext, orient_conf: OrientationConfidence, correction: int
|
|
) -> str:
|
|
"""Describe the page rotation we are going to perform (or not perform)."""
|
|
direction = {0: '⇧', 90: '⇨', 180: '⇩', 270: '⇦'}
|
|
turns = {0: ' ', 90: '⬏', 180: '↻', 270: '⬑'}
|
|
|
|
existing_rotation = page_context.pageinfo.rotation
|
|
action = ''
|
|
if orient_conf.confidence >= page_context.options.rotate_pages_threshold:
|
|
if correction != 0:
|
|
action = 'will rotate ' + turns[correction]
|
|
else:
|
|
action = 'rotation appears correct'
|
|
else:
|
|
if correction != 0:
|
|
action = 'confidence too low to rotate'
|
|
else:
|
|
action = 'no change'
|
|
|
|
facing = ''
|
|
|
|
if existing_rotation != 0:
|
|
facing = f"with existing rotation {direction.get(existing_rotation, '?')}, "
|
|
facing += f"page is facing {direction.get(orient_conf.angle, '?')}"
|
|
|
|
return f"{facing}, confidence {orient_conf.confidence:.2f} - {action}"
|
|
|
|
|
|
def get_orientation_correction(preview: Path, page_context: PageContext) -> int:
|
|
"""Work out orientation correction for each page.
|
|
|
|
We ask Ghostscript to draw a preview page, which will rasterize with the
|
|
current /Rotate applied, and then ask OCR which way the page is
|
|
oriented. If the value of /Rotate is correct (e.g., a user already
|
|
manually fixed rotation), then OCR will say the page is pointing
|
|
up and the correction is zero. Otherwise, the orientation found by
|
|
OCR represents the clockwise rotation, or the counterclockwise
|
|
correction to rotation.
|
|
|
|
When we draw the real page for OCR, we rotate it by the CCW correction,
|
|
which points it (hopefully) upright. _graft.py takes care of the orienting
|
|
the image and text layers.
|
|
"""
|
|
orient_conf = page_context.plugin_manager.hook.get_ocr_engine().get_orientation(
|
|
preview, page_context.options
|
|
)
|
|
|
|
correction = orient_conf.angle % 360
|
|
log.info(describe_rotation(page_context, orient_conf, correction))
|
|
if (
|
|
orient_conf.confidence >= page_context.options.rotate_pages_threshold
|
|
and correction != 0
|
|
):
|
|
return correction
|
|
|
|
return 0
|
|
|
|
|
|
def calculate_image_dpi(page_context: PageContext) -> Resolution:
|
|
"""Calculate the DPI for the page image."""
|
|
pageinfo = page_context.pageinfo
|
|
dpi_profile = pageinfo.page_dpi_profile()
|
|
if dpi_profile and dpi_profile.average_to_max_dpi_ratio < 0.8:
|
|
image_dpi = Resolution(dpi_profile.weighted_dpi, dpi_profile.weighted_dpi)
|
|
else:
|
|
image_dpi = pageinfo.dpi
|
|
return image_dpi
|
|
|
|
|
|
def calculate_raster_dpi(page_context: PageContext):
|
|
"""Calculate the DPI for rasterization."""
|
|
# Produce the page image with square resolution or else deskew and OCR
|
|
# will not work properly.
|
|
image_dpi = calculate_image_dpi(page_context)
|
|
dpi_profile = page_context.pageinfo.page_dpi_profile()
|
|
canvas_dpi = get_canvas_square_dpi(page_context, image_dpi)
|
|
page_dpi = get_page_square_dpi(page_context, image_dpi)
|
|
if dpi_profile and dpi_profile.average_to_max_dpi_ratio < 0.8:
|
|
log.warning(
|
|
"Weight average image DPI is %0.1f, max DPI is %0.1f. "
|
|
"The discrepancy may indicate a high detail region on this page, "
|
|
"but could also indicate a problem with the input PDF file. "
|
|
"Page image will be rendered at %0.1f DPI.",
|
|
dpi_profile.weighted_dpi,
|
|
dpi_profile.max_dpi,
|
|
canvas_dpi.to_scalar(),
|
|
)
|
|
return canvas_dpi, page_dpi
|
|
|
|
|
|
def rasterize(
|
|
input_file: Path,
|
|
page_context: PageContext,
|
|
correction: int = 0,
|
|
output_tag: str = '',
|
|
remove_vectors: bool | None = None,
|
|
) -> Path:
|
|
"""Rasterize a PDF page to a PNG image.
|
|
|
|
Args:
|
|
input_file: The input PDF file path.
|
|
page_context: The page context object.
|
|
correction: The orientation correction angle. Defaults to 0.
|
|
output_tag: The output tag. Defaults to ''.
|
|
remove_vectors: Whether to remove vectors. Defaults to None, which means
|
|
the value from the page context options will be used. If the value
|
|
is True or False, it will override the page context options.
|
|
|
|
Returns:
|
|
Path: The output PNG file path.
|
|
"""
|
|
colorspaces = ['pngmono', 'pnggray', 'png256', 'png16m']
|
|
device_idx = 0
|
|
|
|
if remove_vectors is None:
|
|
remove_vectors = page_context.options.remove_vectors
|
|
|
|
output_file = page_context.get_path(f'rasterize{output_tag}.png')
|
|
pageinfo = page_context.pageinfo
|
|
|
|
def at_least(colorspace):
|
|
return max(device_idx, colorspaces.index(colorspace))
|
|
|
|
for image in pageinfo.images:
|
|
if image.type_ != 'image':
|
|
continue # ignore masks
|
|
if image.bpc > 1:
|
|
if image.color == Colorspace.index:
|
|
device_idx = at_least('png256')
|
|
elif image.color == Colorspace.gray:
|
|
device_idx = at_least('pnggray')
|
|
else:
|
|
device_idx = at_least('png16m')
|
|
|
|
if pageinfo.has_vector:
|
|
log.debug("Page has vector content, using png16m")
|
|
device_idx = at_least('png16m')
|
|
|
|
device = colorspaces[device_idx]
|
|
|
|
log.debug(f"Rasterize with {device}, rotation {correction}")
|
|
|
|
canvas_dpi, page_dpi = calculate_raster_dpi(page_context)
|
|
|
|
page_context.plugin_manager.hook.rasterize_pdf_page(
|
|
input_file=input_file,
|
|
output_file=output_file,
|
|
raster_device=device,
|
|
raster_dpi=canvas_dpi,
|
|
page_dpi=page_dpi,
|
|
pageno=pageinfo.pageno + 1,
|
|
rotation=correction,
|
|
filter_vector=remove_vectors,
|
|
stop_on_soft_error=not page_context.options.continue_on_soft_render_error,
|
|
)
|
|
return output_file
|
|
|
|
|
|
def preprocess_remove_background(input_file: Path, page_context: PageContext) -> Path:
|
|
"""Remove the background from the input image (temporarily disabled)."""
|
|
if any(image.bpc > 1 for image in page_context.pageinfo.images):
|
|
raise NotImplementedError("--remove-background is temporarily not implemented")
|
|
# output_file = page_context.get_path('pp_rm_bg.png')
|
|
# leptonica.remove_background(input_file, output_file)
|
|
# return output_file
|
|
log.info("background removal skipped on mono page")
|
|
return input_file
|
|
|
|
|
|
def preprocess_deskew(input_file: Path, page_context: PageContext) -> Path:
|
|
"""Deskews the input image using the OCR engine and saves the output to a file.
|
|
|
|
Args:
|
|
input_file: The input image file to deskew.
|
|
page_context: The context of the page being processed.
|
|
|
|
Returns:
|
|
Path: The path to the deskewed image file.
|
|
"""
|
|
output_file = page_context.get_path('pp_deskew.png')
|
|
dpi = get_page_square_dpi(page_context, calculate_image_dpi(page_context))
|
|
|
|
ocr_engine = page_context.plugin_manager.hook.get_ocr_engine()
|
|
deskew_angle_degrees = ocr_engine.get_deskew(input_file, page_context.options)
|
|
|
|
with Image.open(input_file) as im:
|
|
# According to Pillow docs, .rotate() will automatically use Image.NEAREST
|
|
# resampling if image is mode '1' or 'P'
|
|
deskewed = im.rotate(
|
|
deskew_angle_degrees,
|
|
resample=Image.Resampling.BICUBIC,
|
|
fillcolor=ImageColor.getcolor('white', mode=im.mode), # type: ignore
|
|
)
|
|
deskewed.save(output_file, dpi=dpi)
|
|
|
|
return output_file
|
|
|
|
|
|
def preprocess_clean(input_file: Path, page_context: PageContext) -> Path:
|
|
"""Clean the input image using unpaper."""
|
|
output_file = page_context.get_path('pp_clean.png')
|
|
dpi = get_page_square_dpi(page_context, calculate_image_dpi(page_context))
|
|
return unpaper.clean(
|
|
input_file,
|
|
output_file,
|
|
dpi=dpi.to_scalar(),
|
|
unpaper_args=page_context.options.unpaper_args,
|
|
)
|
|
|
|
|
|
def create_ocr_image(image: Path, page_context: PageContext) -> Path:
|
|
"""Create the image we send for OCR.
|
|
|
|
Might not be the same as the display image depending on preprocessing.
|
|
This image will never be shown to the user.
|
|
"""
|
|
output_file = page_context.get_path('ocr.png')
|
|
options = page_context.options
|
|
with Image.open(image) as im:
|
|
log.debug('resolution %r', im.info['dpi'])
|
|
|
|
if not options.force_ocr:
|
|
# Do not mask text areas when forcing OCR, because we need to OCR
|
|
# all text areas
|
|
mask = None # Exclude both visible and invisible text from OCR
|
|
if options.redo_ocr:
|
|
mask = True # Mask visible text, but not invisible text
|
|
|
|
draw = ImageDraw.ImageDraw(im)
|
|
for textarea in page_context.pageinfo.get_textareas(
|
|
visible=mask, corrupt=None
|
|
):
|
|
# Calculate resolution based on the image size and page dimensions
|
|
# without regard whatever resolution is in pageinfo (may differ or
|
|
# be None)
|
|
bbox = [float(v) for v in textarea]
|
|
xyscale = tuple(float(coord) / 72.0 for coord in im.info['dpi'])
|
|
pixcoords = (
|
|
bbox[0] * xyscale[0],
|
|
im.height - bbox[3] * xyscale[1],
|
|
bbox[2] * xyscale[0],
|
|
im.height - bbox[1] * xyscale[1],
|
|
)
|
|
log.debug('blanking %r', pixcoords)
|
|
draw.rectangle(pixcoords, fill='white')
|
|
# draw.rectangle(pixcoords, outline='pink')
|
|
|
|
filter_im = page_context.plugin_manager.hook.filter_ocr_image(
|
|
page=page_context, image=im
|
|
)
|
|
if filter_im is not None:
|
|
im = filter_im
|
|
|
|
# Pillow requires integer DPI
|
|
dpi = tuple(round(coord) for coord in im.info['dpi'])
|
|
im.save(output_file, dpi=dpi)
|
|
return output_file
|
|
|
|
|
|
def ocr_engine_hocr(input_file: Path, page_context: PageContext) -> tuple[Path, Path]:
|
|
"""Run the OCR engine and generate hOCR output."""
|
|
hocr_out = page_context.get_path('ocr_hocr.hocr')
|
|
hocr_text_out = page_context.get_path('ocr_hocr.txt')
|
|
options = page_context.options
|
|
|
|
ocr_engine = page_context.plugin_manager.hook.get_ocr_engine()
|
|
ocr_engine.generate_hocr(
|
|
input_file=input_file,
|
|
output_hocr=hocr_out,
|
|
output_text=hocr_text_out,
|
|
options=options,
|
|
)
|
|
return hocr_out, hocr_text_out
|
|
|
|
|
|
def should_visible_page_image_use_jpg(pageinfo: PageInfo) -> bool:
|
|
"""Determines whether the visible page image should be saved as a JPEG.
|
|
|
|
If all images were JPEGs originally, permit a JPEG as output.
|
|
|
|
Args:
|
|
pageinfo: The PageInfo object containing information about the page.
|
|
|
|
Returns:
|
|
A boolean indicating whether the visible page image should be saved as a JPEG.
|
|
"""
|
|
return bool(pageinfo.images) and all(
|
|
im.enc == Encoding.jpeg for im in pageinfo.images
|
|
)
|
|
|
|
|
|
def create_visible_page_jpg(image: Path, page_context: PageContext) -> Path:
|
|
"""Create a visible page image in JPEG format.
|
|
|
|
This is intended to be used when all images on the page were originally JPEGs.
|
|
"""
|
|
output_file = page_context.get_path('visible.jpg')
|
|
with Image.open(image) as im:
|
|
# At this point the image should be a .png, but deskew, unpaper
|
|
# might have removed the DPI information. In this case, fall back to
|
|
# square DPI used to rasterize. When the preview image was
|
|
# rasterized, it was also converted to square resolution, which is
|
|
# what we want to give to the OCR engine, so keep it square.
|
|
if 'dpi' in im.info:
|
|
dpi = Resolution(*im.info['dpi'])
|
|
else:
|
|
# Fallback to page-implied DPI
|
|
dpi = get_page_square_dpi(page_context, calculate_image_dpi(page_context))
|
|
|
|
# Pillow requires integer DPI
|
|
im.save(output_file, format='JPEG', dpi=dpi.to_int())
|
|
return output_file
|
|
|
|
|
|
def create_pdf_page_from_image(
|
|
image: Path, page_context: PageContext, orientation_correction: int
|
|
) -> Path:
|
|
"""Create a PDF page from a page image."""
|
|
# We rasterize a square DPI version of each page because most image
|
|
# processing tools don't support rectangular DPI. Use the square DPI as it
|
|
# accurately describes the image. It would be possible to resample the image
|
|
# at this stage back to non-square DPI to more closely resemble the input,
|
|
# except that the hocr renderer does not understand non-square DPI. The
|
|
# sandwich renderer would be fine.
|
|
output_file = page_context.get_path('visible.pdf')
|
|
|
|
pageinfo = page_context.pageinfo
|
|
pagesize = 72.0 * float(pageinfo.width_inches), 72.0 * float(pageinfo.height_inches)
|
|
effective_rotation = (pageinfo.rotation - orientation_correction) % 360
|
|
swap_axis = effective_rotation % 180 == 90
|
|
if swap_axis:
|
|
pagesize = pagesize[1], pagesize[0]
|
|
|
|
# Create a new single page PDF to hold
|
|
bio = BytesIO()
|
|
with open(image, 'rb') as imfile:
|
|
log.debug('convert')
|
|
|
|
layout_fun = img2pdf.get_layout_fun(pagesize)
|
|
img2pdf.convert(
|
|
imfile,
|
|
layout_fun=layout_fun,
|
|
outputstream=bio,
|
|
engine=img2pdf.Engine.pikepdf,
|
|
rotation=img2pdf.Rotation.ifvalid,
|
|
)
|
|
log.debug('convert done')
|
|
|
|
# img2pdf does not generate boxes correctly, so we fix them
|
|
bio.seek(0)
|
|
fix_pagepdf_boxes(bio, output_file, page_context, swap_axis=swap_axis)
|
|
|
|
output_file = page_context.plugin_manager.hook.filter_pdf_page(
|
|
page=page_context, image_filename=image, output_pdf=output_file
|
|
)
|
|
return output_file
|
|
|
|
|
|
def render_hocr_page(hocr: Path, page_context: PageContext) -> Path:
|
|
"""Render the hOCR page to a PDF."""
|
|
options = page_context.options
|
|
output_file = page_context.get_path('ocr_hocr.pdf')
|
|
if hocr.stat().st_size == 0:
|
|
# If hOCR file is empty (skipped page marker), create an empty PDF file
|
|
output_file.touch()
|
|
return output_file
|
|
|
|
dpi = get_page_square_dpi(page_context, calculate_image_dpi(page_context))
|
|
debug_kwargs = {}
|
|
if options.pdf_renderer == 'hocrdebug':
|
|
debug_kwargs = dict(
|
|
debug_render_options=DebugRenderOptions(
|
|
render_baseline=True,
|
|
render_triangle=True,
|
|
render_line_bbox=False,
|
|
render_word_bbox=True,
|
|
render_paragraph_bbox=False,
|
|
render_space_bbox=False,
|
|
),
|
|
font=Courier(),
|
|
)
|
|
HocrTransform(
|
|
hocr_filename=hocr, dpi=dpi.to_scalar(), **debug_kwargs # square
|
|
).to_pdf(
|
|
out_filename=output_file,
|
|
image_filename=None,
|
|
invisible_text=True if not debug_kwargs else False,
|
|
)
|
|
return output_file
|
|
|
|
|
|
def ocr_engine_textonly_pdf(
|
|
input_image: Path, page_context: PageContext
|
|
) -> tuple[Path, Path]:
|
|
"""Run the OCR engine and generate a text-only PDF (will look blank)."""
|
|
output_pdf = page_context.get_path('ocr_tess.pdf')
|
|
output_text = page_context.get_path('ocr_tess.txt')
|
|
options = page_context.options
|
|
|
|
ocr_engine = page_context.plugin_manager.hook.get_ocr_engine()
|
|
ocr_engine.generate_pdf(
|
|
input_file=input_image,
|
|
output_pdf=output_pdf,
|
|
output_text=output_text,
|
|
options=options,
|
|
)
|
|
return output_pdf, output_text
|
|
|
|
|
|
def _offset_rect(rect: tuple[float, float, float, float], offset: tuple[float, float]):
|
|
"""Offset a rectangle by a given amount."""
|
|
return (
|
|
rect[0] + offset[0],
|
|
rect[1] + offset[1],
|
|
rect[2] + offset[0],
|
|
rect[3] + offset[1],
|
|
)
|
|
|
|
|
|
def fix_pagepdf_boxes(
|
|
infile: Path | BinaryIO,
|
|
out_file: Path,
|
|
page_context: PageContext,
|
|
swap_axis: bool = False,
|
|
) -> Path:
|
|
"""Fix the bounding boxes in a single page PDF.
|
|
|
|
The single page PDF is created with a normal MediaBox with its lower left corner
|
|
at (0, 0). infile is the single page PDF. page_context.mediabox has the original
|
|
file's mediabox, which may have a different origin. We needto adjust the other
|
|
boxes in the single page PDF to match the effect they had on the original page.
|
|
|
|
When correcting page rotation, we create a single page PDF that is correctly
|
|
rotated instead of an incorrectly rotated and then setting page.Rotate on it.
|
|
If rotation is either 90 or 270 degrees, then this function can be called
|
|
with swap_axis to swap the X and Y coordinates of all the boxes.
|
|
|
|
We are not concerned with solving degenerate cases where the boxes overlap or
|
|
or express invalid rectangles. We merely pass the boxes, producing a
|
|
transformation equivalent to the change made by constructing a new page image.
|
|
"""
|
|
with pikepdf.open(infile) as pdf:
|
|
for page in pdf.pages:
|
|
# page.BleedBox = page_context.pageinfo.bleedbox
|
|
# page.ArtBox = page_context.pageinfo.artbox
|
|
mediabox = page_context.pageinfo.mediabox
|
|
offset = mediabox[0], mediabox[1]
|
|
cropbox = _offset_rect(page_context.pageinfo.cropbox, offset)
|
|
trimbox = _offset_rect(page_context.pageinfo.trimbox, offset)
|
|
|
|
if swap_axis:
|
|
cropbox = cropbox[1], cropbox[0], cropbox[3], cropbox[2]
|
|
trimbox = trimbox[1], trimbox[0], trimbox[3], trimbox[2]
|
|
page.CropBox = cropbox
|
|
page.TrimBox = trimbox
|
|
pdf.save(out_file)
|
|
return pdf
|
|
|
|
|
|
def generate_postscript_stub(context: PdfContext) -> Path:
|
|
"""Generates a PostScript file stub for the given PDF context.
|
|
|
|
Args:
|
|
context: The PDF context to generate the PostScript file stub for.
|
|
|
|
Returns:
|
|
Path: The path to the generated PostScript file stub.
|
|
"""
|
|
output_file = context.get_path('pdfa.ps')
|
|
generate_pdfa_ps(output_file)
|
|
return output_file
|
|
|
|
|
|
def convert_to_pdfa(input_pdf: Path, input_ps_stub: Path, context: PdfContext) -> Path:
|
|
"""Converts the given PDF to PDF/A.
|
|
|
|
Args:
|
|
input_pdf: The input PDF file path (presumably not PDF/A).
|
|
input_ps_stub: The input PostScript file path, containing instructions
|
|
for the PDF/A generator to use.
|
|
context: The PDF context.
|
|
"""
|
|
options = context.options
|
|
input_pdfinfo = context.pdfinfo
|
|
fix_docinfo_file = context.get_path('fix_docinfo.pdf')
|
|
output_file = context.get_path('pdfa.pdf')
|
|
|
|
# If the DocumentInfo record contains NUL characters, Ghostscript will
|
|
# produce XMP metadata which contains invalid XML entities (�).
|
|
# NULs in DocumentInfo seem to be common since older Acrobats included them.
|
|
# pikepdf can deal with this, but we make the world a better place by
|
|
# stamping them out as soon as possible.
|
|
with pikepdf.open(input_pdf) as pdf_file:
|
|
if repair_docinfo_nuls(pdf_file):
|
|
pdf_file.save(fix_docinfo_file)
|
|
else:
|
|
safe_symlink(input_pdf, fix_docinfo_file)
|
|
|
|
context.plugin_manager.hook.generate_pdfa(
|
|
pdf_version=input_pdfinfo.min_version,
|
|
pdf_pages=[fix_docinfo_file],
|
|
pdfmark=input_ps_stub,
|
|
output_file=output_file,
|
|
context=context,
|
|
pdfa_part=options.output_type[-1], # is pdfa-1, pdfa-2, or pdfa-3
|
|
progressbar_class=(
|
|
context.plugin_manager.hook.get_progressbar_class()
|
|
if options.progress_bar
|
|
else None
|
|
),
|
|
stop_on_soft_error=not options.continue_on_soft_render_error,
|
|
)
|
|
|
|
return output_file
|
|
|
|
|
|
def should_linearize(working_file: Path, context: PdfContext) -> bool:
|
|
"""Determine whether the PDF should be linearized.
|
|
|
|
For smaller files, linearization is not worth the effort.
|
|
"""
|
|
filesize = os.stat(working_file).st_size
|
|
if filesize > (context.options.fast_web_view * 1_000_000):
|
|
return True
|
|
return False
|
|
|
|
|
|
def get_pdf_save_settings(output_type: str) -> dict[str, Any]:
|
|
"""Get pikepdf.Pdf.save settings for the given output type.
|
|
|
|
Essentially, don't use features that are incompatible with a given
|
|
PDF/A specification.
|
|
"""
|
|
if output_type == 'pdfa-1':
|
|
# Trigger recompression to ensure object streams are removed, because
|
|
# Acrobat complains about them in PDF/A-1b validation.
|
|
return dict(
|
|
preserve_pdfa=True,
|
|
compress_streams=True,
|
|
stream_decode_level=pikepdf.StreamDecodeLevel.generalized,
|
|
object_stream_mode=pikepdf.ObjectStreamMode.disable,
|
|
)
|
|
else:
|
|
return dict(
|
|
preserve_pdfa=True,
|
|
compress_streams=True,
|
|
object_stream_mode=(pikepdf.ObjectStreamMode.generate),
|
|
)
|
|
|
|
|
|
def _file_size_ratio(
|
|
input_file: Path, output_file: Path
|
|
) -> tuple[float | None, float | None]:
|
|
"""Calculate ratio of input to output file sizes and percentage savings.
|
|
|
|
Args:
|
|
input_file (Path): The path to the input file.
|
|
output_file (Path): The path to the output file.
|
|
|
|
Returns:
|
|
tuple[float | None, float | None]: A tuple containing the file size
|
|
ratio and the percentage savings achieved by the output file size
|
|
compared to the input file size.
|
|
"""
|
|
input_size = input_file.stat().st_size
|
|
output_size = output_file.stat().st_size
|
|
if output_size == 0:
|
|
return None, None
|
|
ratio = input_size / output_size
|
|
savings = 1 - output_size / input_size
|
|
return ratio, savings
|
|
|
|
|
|
def optimize_pdf(
|
|
input_file: Path, context: PdfContext, executor: Executor
|
|
) -> tuple[Path, Sequence[str]]:
|
|
"""Optimize the given PDF file."""
|
|
output_file = context.get_path('optimize.pdf')
|
|
output_pdf, messages = context.plugin_manager.hook.optimize_pdf(
|
|
input_pdf=input_file,
|
|
output_pdf=output_file,
|
|
context=context,
|
|
executor=executor,
|
|
linearize=should_linearize(input_file, context),
|
|
)
|
|
|
|
ratio, savings = _file_size_ratio(input_file, output_file)
|
|
if ratio:
|
|
log.info(f"Image optimization ratio: {ratio:.2f} savings: {(savings):.1%}")
|
|
ratio, savings = _file_size_ratio(context.origin, output_file)
|
|
if ratio:
|
|
log.info(f"Total file size ratio: {ratio:.2f} savings: {(savings):.1%}")
|
|
return output_pdf, messages
|
|
|
|
|
|
def enumerate_compress_ranges(
|
|
iterable: Iterable[T],
|
|
) -> Iterator[tuple[tuple[int, int], T | None]]:
|
|
"""Enumerate the ranges of non-empty elements in an iterable.
|
|
|
|
Compresses consecutive ranges of length 1 into single elements.
|
|
|
|
Args:
|
|
iterable: An iterable of elements to enumerate.
|
|
|
|
Yields:
|
|
A tuple containing a range of indices and the corresponding element.
|
|
If the element is None, the range represents a skipped range of indices.
|
|
"""
|
|
skipped_from, index = None, None
|
|
for index, txt_file in enumerate(iterable):
|
|
index += 1
|
|
if txt_file:
|
|
if skipped_from is not None:
|
|
yield (skipped_from, index - 1), None
|
|
skipped_from = None
|
|
yield (index, index), txt_file
|
|
else:
|
|
if skipped_from is None:
|
|
skipped_from = index
|
|
if skipped_from is not None:
|
|
yield (skipped_from, index), None
|
|
|
|
|
|
def merge_sidecars(txt_files: Iterable[Path | None], context: PdfContext) -> Path:
|
|
"""Merge the page sidecar files into a single file.
|
|
|
|
Sidecar files are created by the OCR engine and contain the text for each
|
|
page in the PDF. This function merges the sidecar files into a single file
|
|
and returns the path to the merged file.
|
|
"""
|
|
output_file = context.get_path('sidecar.txt')
|
|
with open(output_file, 'w', encoding="utf-8") as stream:
|
|
for (from_, to_), txt_file in enumerate_compress_ranges(txt_files):
|
|
if from_ != 1:
|
|
stream.write('\f') # Form feed between pages for all pages after first
|
|
if txt_file:
|
|
txt = txt_file.read_text(encoding="utf-8")
|
|
# Some versions of Tesseract add a form feed at the end and
|
|
# others don't. Remove it if it exists, since we add one manually.
|
|
stream.write(txt.removesuffix('\f'))
|
|
else:
|
|
if from_ != to_:
|
|
pages = f'{from_}-{to_}'
|
|
else:
|
|
pages = f'{from_}'
|
|
stream.write(f'[OCR skipped on page(s) {pages}]')
|
|
return output_file
|
|
|
|
|
|
def copy_final(
|
|
input_file: Path, output_file: str | Path | BinaryIO, original_file: Path | None
|
|
) -> None:
|
|
"""Copy the final temporary file to the output destination.
|
|
|
|
Args:
|
|
input_file (Path): The intermediate input file to copy.
|
|
output_file (str | Path | BinaryIO): The output file to copy to.
|
|
original_file: The original file to copy attributes from.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
log.debug('%s -> %s', input_file, output_file)
|
|
with input_file.open('rb') as input_stream:
|
|
if output_file == '-':
|
|
copyfileobj(input_stream, sys.stdout.buffer) # type: ignore[misc]
|
|
sys.stdout.flush()
|
|
elif hasattr(output_file, 'writable'):
|
|
output_stream = cast(BinaryIO, output_file)
|
|
copyfileobj(input_stream, output_stream) # type: ignore[misc]
|
|
with suppress(AttributeError):
|
|
output_stream.flush()
|
|
else:
|
|
# At this point we overwrite the output_file specified by the user
|
|
# use copyfileobj because then we use open() to create the file and
|
|
# get the appropriate umask, ownership, etc.
|
|
with open(output_file, 'w+b') as output_stream:
|
|
copyfileobj(input_stream, output_stream)
|