From 01dc8e23ff57fdd8aea2d47600f92b36b9cbcbec Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Thu, 13 Apr 2023 23:27:58 -0700 Subject: [PATCH] As you were --- src/ocrmypdf/builtin_plugins/tesseract_ocr.py | 34 +++++ src/ocrmypdf/imageops.py | 121 ++++++++++++++++++ src/ocrmypdf/pluginspec.py | 27 ++-- tests/test_imageops.py | 28 ++++ 4 files changed, 200 insertions(+), 10 deletions(-) create mode 100644 src/ocrmypdf/imageops.py create mode 100644 tests/test_imageops.py diff --git a/src/ocrmypdf/builtin_plugins/tesseract_ocr.py b/src/ocrmypdf/builtin_plugins/tesseract_ocr.py index 84799b49..533b9e8e 100644 --- a/src/ocrmypdf/builtin_plugins/tesseract_ocr.py +++ b/src/ocrmypdf/builtin_plugins/tesseract_ocr.py @@ -8,10 +8,14 @@ from __future__ import annotations import logging import os +from PIL import Image + from ocrmypdf import hookimpl from ocrmypdf._exec import tesseract +from ocrmypdf._jobcontext import PageContext from ocrmypdf.cli import numeric, str_to_int from ocrmypdf.helpers import clamp +from ocrmypdf.imageops import calculate_downsample, downsample_image from ocrmypdf.pluginspec import OcrEngine from ocrmypdf.subprocess import check_external_program @@ -85,6 +89,19 @@ def add_options(parser): "because these operations are not as expensive as OCR." ), ) + tess.add_argument( + '--tesseract-downsample-large-images', + action='store_true', + help=( + "Downsample large images before OCR. Tesseract has an upper limit on the " + "size images it will support. If this argument is given, OCRmyPDF will " + "downsample large images to fit Tesseract. This may reduce OCR quality, " + "on large images the most desirable text is usually larger. If this " + "parameter is not supplied, Tesseract will error out and produce no OCR " + "on the page in question. This argument should be used with a high value " + "of --tesseract-timeout to ensure Tesseract has enough to time." + ), + ) tess.add_argument( '--user-words', metavar='FILE', @@ -145,6 +162,23 @@ def validate(pdfinfo, options): log.debug("Using Tesseract OpenMP thread limit %d", tess_threads) +@hookimpl +def filter_ocr_image(page: PageContext, image: Image.Image) -> Image.Image: + """Filter the image before OCR. + + Tesseract cannot handle images with more than 32767 pixels in either axis, + or more than 2**31 bytes. This function resizes the image to fit within + those limits. + """ + options = page.options + if options.tesseract_downsample_large_images: + factor = calculate_downsample( + image, max_size=(32767, 32767), max_bytes=(2**31) - 1 + ) + image = downsample_image(image, factor) + return image + + class TesseractOcrEngine(OcrEngine): """Implements OCR with Tesseract.""" diff --git a/src/ocrmypdf/imageops.py b/src/ocrmypdf/imageops.py new file mode 100644 index 00000000..75ab7ffa --- /dev/null +++ b/src/ocrmypdf/imageops.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: 2023 James R. Barlow +# SPDX-License-Identifier: MPL-2.0 + +"""OCR-related image manipulation.""" + +import logging +from math import ceil, sqrt + +from PIL import Image + +log = logging.getLogger(__name__) + + +def bytes_per_pixel(mode: str) -> int: + """ + Return the number of padded bytes per pixel for a given PIL image mode. + + In RGB mode we assume 4 bytes per pixel, which is the case for most + consumers. + """ + if mode in ('1', 'L', 'P'): + return 1 + if mode in ('LA', 'PA', 'La') or mode.startswith('I;16'): + return 2 + return 4 + + +def calculate_downsample( + image: Image.Image, + *, + max_size: tuple[int, int] | None = None, + max_pixels: int | None = None, + max_bytes: int | None = None, +) -> float: + """ + Calculate the scaling factor required to downsample an image to fit within + the given limits. + + If no limit is exceeded, 1.0 is returned. + + Args: + image: The image to downsample. + max_size: The maximum width and height of the image. + max_pixels: The maximum number of pixels in the image. Some image consumers + limit the total number of pixels as some value other than width*height. + max_bytes: The maximum number of bytes in the image. RGB is counted as 4 + bytes; all other modes are counted as 1 byte. + """ + scaling_factor = 1.0 + + if max_size is not None: + major_axis = max(image.size) + if major_axis > max(max_size): + log.debug("Resizing image to fit Tesseract image size limit") + scaling_factor = max(max_size) / major_axis + + if max_pixels is not None: + if image.size[0] * image.size[1] * scaling_factor * scaling_factor > max_pixels: + log.debug("Resizing image to fit image pixel limit") + scaling_factor *= sqrt( + max_pixels / (image.size[0] * image.size[1] * scaling_factor) + ) + + if max_bytes is not None: + bpp = bytes_per_pixel(image.mode) + # stride = bytes per line + stride = ceil(image.size[0] * scaling_factor) * bpp + height = ceil(image.size[1] * scaling_factor) + size = stride * height + if size > max_bytes: + log.debug("Resizing image to fit image byte size limit") + scaling_factor *= sqrt((max_bytes - 1) / size) + scaled_bytes_per_line = ceil(image.size[0] * scaling_factor) * bpp + height = ceil(image.size[1] * scaling_factor) + size = scaled_bytes_per_line * height + assert size <= max_bytes, f"{size} > {max_bytes}" + + return scaling_factor + + +def downsample_image( + image: Image.Image, + scaling_factor: float, + *, + resample_mode: Image.Resampling = Image.Resampling.BICUBIC, + reducing_gap: int = 3, +) -> Image.Image: + """ + Downsample an image to fit within the given limits. + + The DPI is adjusted to match the new size, which is how we can ensure the + OCR is positioned correctly. + + Args: + image: The image to downsample + scaling_factor: The scaling factor to apply to the image, calculated using + calculate_downsample(). + resample_mode: The resampling mode to use when downsampling. + reducing_gap: The reducing gap to use when downsampling (for larger + reductions). + """ + if scaling_factor == 1.0: + return image + if scaling_factor > 1.0 or scaling_factor <= 0: + raise ValueError("scaling_factor must be <= 1.0 and > 0") + + original_dpi = image.info['dpi'] + image = image.resize( + ( + ceil(image.size[0] * scaling_factor), + ceil(image.size[1] * scaling_factor), + ), + resample=resample_mode, + reducing_gap=reducing_gap, + ) + image.info['dpi'] = ( + original_dpi[0] * scaling_factor, + original_dpi[1] * scaling_factor, + ) + log.debug(f"Rescaled image to {image.size} pixels and {image.info['dpi']} dpi") + return image diff --git a/src/ocrmypdf/pluginspec.py b/src/ocrmypdf/pluginspec.py index afb46194..b23fa0d6 100644 --- a/src/ocrmypdf/pluginspec.py +++ b/src/ocrmypdf/pluginspec.py @@ -228,19 +228,26 @@ def filter_ocr_image(page: PageContext, image: Image.Image) -> Image.Image: """Called to filter the image before it is sent to OCR. This is the image that OCR sees, not what the user sees when they view the - PDF. If ``redo_ocr`` is enabled, portions of the image will be masked so - they are not shown to OCR. The main use of this hook is expected to be hiding - content from OCR. + PDF. In certain modes such as ``--redo-ocr``, portions of the image may be + masked out to hide them from OCR. + + The main uses of this hook are expected to be hiding content from OCR, + conditioning images to OCR better with filters, and adjusting images to + match any constraints imposed by the OCR engine. The input image may be color, grayscale, or monochrome, and the - output image may differ. The pixel width and height of the - output image must be identical to the input image, or misalignment between - the OCR text layer and visual position of the text will occur. Likewise, - the output must be a faithful representation of the input, or alignment - errors may occurs. + output image may differ. For example, if you know that a custom OCR engine + does not care about the color of the text, you could convert the image to + it to grayscale or monochrome. - Tesseract OCR only deals with monochrome images, and internally converts - non-monochrome images to OCR. + Generally speaking, the output image should be a faithful representation of + of the input image. You *may* change the pixel width and height of the + the input image, but you must not change the aspect ratio, and you must + calculate the DPI of the output image based on the new pixel width and + height or the OCR text layer will be misaligned with the visual position. + + The built-in Tesseract OCR engine uses this hook itself to downsample + very large images to fit its constraints. Note: This hook will be called from child processes. Modifying global state diff --git a/tests/test_imageops.py b/tests/test_imageops.py new file mode 100644 index 00000000..7f75a8fa --- /dev/null +++ b/tests/test_imageops.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: 2023 James R. Barlow +# SPDX-License-Identifier: MPL-2.0 + +from __future__ import annotations + +from ocrmypdf.imageops import bytes_per_pixel, calculate_downsample, downsample_image +from PIL import Image + + +def test_bytes_per_pixel(): + assert bytes_per_pixel('RGB') == 4 + assert bytes_per_pixel('RGBA') == 4 + assert bytes_per_pixel('LA') == 2 + assert bytes_per_pixel('L') == 1 + + +def test_calculate_downsample(): + im = Image.new('RGB', (100, 100)) + assert calculate_downsample(im, max_size=(50, 50)) == 0.5 + assert calculate_downsample(im, max_pixels=2500) == 0.5 + assert calculate_downsample(im, max_bytes=10000) == 0.25 + + +def test_downsample_image(): + im = Image.new('RGB', (100, 100)) + im.info['dpi'] = (300, 300) + assert downsample_image(im, 0.5).size == (50, 50) + assert im.info['dpi'] == (150, 150)