# 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 pathlib import Path from shutil import copyfileobj, copystat 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 HocrTransform 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 if effective_rotation % 180 == 90: pagesize = pagesize[1], pagesize[0] # This create a single page PDF with open(image, 'rb') as imfile, open(output_file, 'wb') as pdf: log.debug('convert') layout_fun = img2pdf.get_layout_fun(pagesize) img2pdf.convert( imfile, layout_fun=layout_fun, outputstream=pdf, **IMG2PDF_KWARGS ) log.debug('convert done') 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_mode = options.pdf_renderer == 'hocrdebug' hocrtransform = HocrTransform(hocr_filename=hocr, dpi=dpi.to_scalar()) # square hocrtransform.to_pdf( out_filename=output_file, image_filename=None, invisible_text=True if not debug_mode 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 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) # Attempt to copy file attributes from input to output if original_file: with suppress(OSError): # Copy original file's permissions, ownership, etc. if possible copystat(original_file, output_file) # Set output file's modification time to now Path(output_file).touch(exist_ok=True)