Merge feature to downsample very large images

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:
+        size = calculate_downsample(
+            image, max_size=(32767, 32767), max_bytes=(2**31) - 1
+        )
+        image = downsample_image(image, size)
+    return image
+
+
 class TesseractOcrEngine(OcrEngine):
     """Implements OCR with Tesseract."""
 

src/ocrmypdf/imageops.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: 2023 James R. Barlow
+# SPDX-License-Identifier: MPL-2.0
+
+"""OCR-related image manipulation."""
+
+import logging
+from math import ceil, floor, 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,
+) -> tuple[int, int]:
+    """
+    Calculate the new image size required to downsample an image to fit within
+    the given limits.
+
+    If no limit is exceeded, the input image's size 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.
+    """
+    size = image.size
+
+    if max_size is not None:
+        major_axis = max(image.size)
+        size_factor = max(max_size) / major_axis
+        if size_factor < 1.0:
+            log.debug("Resizing image to fit Tesseract image size limit")
+            size = floor(size[0] * size_factor), floor(size[1] * size_factor)
+
+    if max_pixels is not None:
+        if size[0] * size[1] > max_pixels:
+            log.debug("Resizing image to fit image pixel limit")
+            pixels_factor = sqrt(max_pixels / (image.size[0] * image.size[1]))
+            size = floor(size[0] * pixels_factor), floor(size[1] * pixels_factor)
+
+    if max_bytes is not None:
+        bpp = bytes_per_pixel(image.mode)
+        # stride = bytes per line
+        stride = size[0] * bpp
+        height = size[1]
+        if stride * height > max_bytes:
+            log.debug("Resizing image to fit image byte size limit")
+            bytes_factor = sqrt((max_bytes) / (stride * height))
+            scaled_stride = floor(stride * bytes_factor)
+            scaled_height = floor(height * bytes_factor)
+            size = ceil(scaled_stride / bpp), scaled_height
+            assert (size[0] * bpp * size[1]) <= max_bytes
+
+    return size
+
+
+def downsample_image(
+    image: Image.Image,
+    new_size: tuple[int, int],
+    *,
+    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 new_size == image.size:
+        return image
+
+    original_size = image.size
+    original_dpi = image.info['dpi']
+    image = image.resize(
+        new_size,
+        resample=resample_mode,
+        reducing_gap=reducing_gap,
+    )
+    image.info['dpi'] = (
+        round(original_dpi[0] * new_size[0] / original_size[0]),
+        round(original_dpi[1] * new_size[1] / original_size[1]),
+    )
+    log.debug(f"Rescaled image to {image.size} pixels and {image.info['dpi']} dpi")
+    return image

src/ocrmypdf/pluginspec.py

@@ -236,19 +236,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

tests/test_imageops.py

@@ -0,0 +1,30 @@
+# 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)) == (50, 50)
+    assert calculate_downsample(im, max_pixels=2500) == (50, 50)
+    assert calculate_downsample(im, max_bytes=10000) == (50, 50)
+    assert calculate_downsample(im, max_bytes=100000) == (100, 100)
+
+
+def test_downsample_image():
+    im = Image.new('RGB', (100, 100))
+    im.info['dpi'] = (300, 300)
+    ds = downsample_image(im, (50, 50))
+    assert ds.size == (50, 50)
+    assert ds.info['dpi'] == (150, 150)