From ea69e868ed95a335b362a3708628c0372cb7abb8 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Tue, 11 Jan 2022 10:44:38 -0800 Subject: [PATCH] unpaper: issue warning if image too large to clean --- src/ocrmypdf/_exec/unpaper.py | 33 +++++++++++++++++++++++++++------ src/ocrmypdf/_pipeline.py | 3 +-- tests/test_unpaper.py | 28 +++++++++++++++++++++++----- 3 files changed, 51 insertions(+), 13 deletions(-) diff --git a/src/ocrmypdf/_exec/unpaper.py b/src/ocrmypdf/_exec/unpaper.py index aec365c2..d478a8b4 100644 --- a/src/ocrmypdf/_exec/unpaper.py +++ b/src/ocrmypdf/_exec/unpaper.py @@ -22,14 +22,28 @@ from typing import List, Optional, Tuple, Union from PIL import Image from ocrmypdf.exceptions import MissingDependencyError, SubprocessOutputError -from ocrmypdf.subprocess import get_version -from ocrmypdf.subprocess import run as external_run +from ocrmypdf.subprocess import get_version, run + +UNPAPER_IMAGE_PIXEL_LIMIT = 256 * 1024 * 1024 DecFloat = Union[Decimal, float] log = logging.getLogger(__name__) +class UnpaperImageTooLargeError(Exception): + def __init__( + self, + w, + h, + message="Image with size {}x{} is too large for cleaning with 'unpaper'.", + ): + self.w = w + self.h = h + self.message = message.format(w, h) + super().__init__(self.message) + + def version() -> str: return get_version('unpaper') @@ -38,6 +52,8 @@ def _setup_unpaper_io(tmpdir: Path, input_file: Path) -> Tuple[Path, Path]: SUFFIXES = {'1': '.pbm', 'L': '.pgm', 'RGB': '.ppm'} with Image.open(input_file) as im: im_modified = False + if im.width * im.height >= UNPAPER_IMAGE_PIXEL_LIMIT: + raise UnpaperImageTooLargeError(w=im.width, h=im.height) if im.mode not in SUFFIXES: log.info("Converting image to other colorspace") try: @@ -68,7 +84,7 @@ def _setup_unpaper_io(tmpdir: Path, input_file: Path) -> Tuple[Path, Path]: return input_pnm, output_pnm -def run( +def run_unpaper( input_file: Path, output_file: Path, *, dpi: DecFloat, mode_args: List[str] ) -> None: args_unpaper = ['unpaper', '-v', '--dpi', str(round(dpi, 6))] + mode_args @@ -84,7 +100,7 @@ def run( # This should ensure that a user cannot clobber some other file with # their unpaper arguments (whether intentionally or otherwise) args_unpaper.extend([os.fspath(input_pnm), os.fspath(output_pnm)]) - external_run( + run( args_unpaper, close_fds=True, check=True, @@ -117,7 +133,7 @@ def clean( *, dpi: DecFloat, unpaper_args: Optional[List[str]] = None, -): +) -> Path: default_args = [ '--layout', 'none', @@ -131,4 +147,9 @@ def clean( ] if not unpaper_args: unpaper_args = default_args - run(input_file, output_file, dpi=dpi, mode_args=unpaper_args) + try: + run_unpaper(input_file, output_file, dpi=dpi, mode_args=unpaper_args) + return output_file + except UnpaperImageTooLargeError as e: + log.warning(str(e)) + return input_file diff --git a/src/ocrmypdf/_pipeline.py b/src/ocrmypdf/_pipeline.py index dc414d7b..8349d9fa 100644 --- a/src/ocrmypdf/_pipeline.py +++ b/src/ocrmypdf/_pipeline.py @@ -494,13 +494,12 @@ def preprocess_deskew(input_file: Path, page_context: PageContext): def preprocess_clean(input_file: Path, page_context: PageContext): output_file = page_context.get_path('pp_clean.png') dpi = get_page_square_dpi(page_context.pageinfo, page_context.options) - unpaper.clean( + return unpaper.clean( input_file, output_file, dpi=dpi.x, unpaper_args=page_context.options.unpaper_args, ) - return output_file def create_ocr_image(image: Path, page_context: PageContext): diff --git a/tests/test_unpaper.py b/tests/test_unpaper.py index e39b6bf6..44b0e5fd 100644 --- a/tests/test_unpaper.py +++ b/tests/test_unpaper.py @@ -5,19 +5,24 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. +import logging from os import fspath from unittest.mock import patch import pytest +from PIL import Image +from ocrmypdf._exec import unpaper from ocrmypdf._plugin_manager import get_parser_options_plugins from ocrmypdf._validation import check_options from ocrmypdf.exceptions import ExitCode, MissingDependencyError -from .conftest import check_ocrmypdf, have_unpaper, run_ocrmypdf +from .conftest import check_ocrmypdf, have_unpaper, ocrmypdf_exec, run_ocrmypdf # pylint: disable=redefined-outer-name +needs_unpaper = pytest.mark.skipif(not have_unpaper(), reason="requires unpaper") + def test_no_unpaper(resources, no_outpdf): input_ = fspath(resources / "c02-22.pdf") @@ -45,7 +50,7 @@ def test_old_unpaper(resources, no_outpdf): mock.assert_called() -@pytest.mark.skipif(not have_unpaper(), reason="requires unpaper") +@needs_unpaper def test_clean(resources, outpdf): check_ocrmypdf( resources / "skew.pdf", @@ -56,7 +61,7 @@ def test_clean(resources, outpdf): ) -@pytest.mark.skipif(not have_unpaper(), reason="requires unpaper") +@needs_unpaper def test_unpaper_args_valid(resources, outpdf): check_ocrmypdf( resources / "skew.pdf", @@ -69,7 +74,7 @@ def test_unpaper_args_valid(resources, outpdf): ) -@pytest.mark.skipif(not have_unpaper(), reason="requires unpaper") +@needs_unpaper def test_unpaper_args_invalid_filename(resources, outpdf): p = run_ocrmypdf( resources / "skew.pdf", @@ -84,7 +89,7 @@ def test_unpaper_args_invalid_filename(resources, outpdf): assert p.returncode == ExitCode.bad_args -@pytest.mark.skipif(not have_unpaper(), reason="requires unpaper") +@needs_unpaper def test_unpaper_args_invalid(resources, outpdf): p = run_ocrmypdf( resources / "skew.pdf", @@ -98,3 +103,16 @@ def test_unpaper_args_invalid(resources, outpdf): # Can't tell difference between unpaper choking on bad arguments or some # other unpaper failure assert p.returncode == ExitCode.child_process_error + + +@needs_unpaper +def test_unpaper_image_too_big(resources, outdir, caplog): + with patch('ocrmypdf._exec.unpaper.UNPAPER_IMAGE_PIXEL_LIMIT', 42): + infile = resources / 'crom.png' + unpaper.clean(infile, outdir / 'out.png', dpi=300) == infile + + assert any( + 'too large for cleaning' in rec.message + for rec in caplog.get_records('call') + if rec.levelno == logging.WARNING + )