diff --git a/src/ocrmypdf/_exec/tesseract.py b/src/ocrmypdf/_exec/tesseract.py index 6d9f060d..47d31e3e 100644 --- a/src/ocrmypdf/_exec/tesseract.py +++ b/src/ocrmypdf/_exec/tesseract.py @@ -6,6 +6,7 @@ from __future__ import annotations import logging +import os import re from contextlib import suppress from enum import IntEnum @@ -27,6 +28,15 @@ from ocrmypdf.subprocess import get_version, run log = logging.getLogger(__name__) +def _tesseract_env(omp_thread_limit: int | None) -> dict[str, str] | None: + """Create environment dict with OMP_THREAD_LIMIT set for Tesseract subprocesses.""" + if omp_thread_limit is None: + return None + env = os.environ.copy() + env['OMP_THREAD_LIMIT'] = str(omp_thread_limit) + return env + + class ThresholdingMethod(IntEnum): """Tesseract thresholding methods for image binarization.""" @@ -166,7 +176,10 @@ def _parse_tesseract_output(binary_output: bytes) -> dict[str, str]: def get_orientation( - input_file: Path, engine_mode: int | None, timeout: float + input_file: Path, + engine_mode: int | None, + timeout: float, + omp_thread_limit: int | None = None, ) -> OrientationConfidence: args_tesseract = tess_base_args(['osd'], engine_mode) + [ '--psm', @@ -176,7 +189,14 @@ def get_orientation( ] try: - p = run(args_tesseract, stdout=PIPE, stderr=STDOUT, timeout=timeout, check=True) + p = run( + args_tesseract, + stdout=PIPE, + stderr=STDOUT, + timeout=timeout, + check=True, + env=_tesseract_env(omp_thread_limit), + ) except TimeoutExpired: return OrientationConfidence(angle=0, confidence=0.0) except CalledProcessError as e: @@ -210,7 +230,11 @@ def _is_empty_page_error(exc): def get_deskew( - input_file: Path, languages: list[str], engine_mode: int | None, timeout: float + input_file: Path, + languages: list[str], + engine_mode: int | None, + timeout: float, + omp_thread_limit: int | None = None, ) -> float: """Gets angle to deskew this page, in degrees.""" args_tesseract = tess_base_args(languages, engine_mode) + [ @@ -221,7 +245,14 @@ def get_deskew( ] try: - p = run(args_tesseract, stdout=PIPE, stderr=STDOUT, timeout=timeout, check=True) + p = run( + args_tesseract, + stdout=PIPE, + stderr=STDOUT, + timeout=timeout, + check=True, + env=_tesseract_env(omp_thread_limit), + ) except TimeoutExpired: return 0.0 except CalledProcessError as e: @@ -308,6 +339,7 @@ def generate_hocr( thresholding: ThresholdingMethod, user_words, user_patterns, + omp_thread_limit: int | None = None, ) -> None: """Generate a hOCR file, which must be converted to PDF.""" prefix = output_hocr.with_suffix('') @@ -331,7 +363,14 @@ def generate_hocr( args_tesseract.extend([fspath(input_file), fspath(prefix), 'hocr', 'txt']) args_tesseract.extend(tessconfig) try: - p = run(args_tesseract, stdout=PIPE, stderr=STDOUT, timeout=timeout, check=True) + p = run( + args_tesseract, + stdout=PIPE, + stderr=STDOUT, + timeout=timeout, + check=True, + env=_tesseract_env(omp_thread_limit), + ) stdout = p.stdout except TimeoutExpired: # Generate a HOCR file with no recognized text if tesseract times out @@ -374,6 +413,7 @@ def generate_pdf( thresholding: ThresholdingMethod, user_words, user_patterns, + omp_thread_limit: int | None = None, ) -> None: """Generate a PDF using Tesseract's internal PDF generator. @@ -404,7 +444,14 @@ def generate_pdf( args_tesseract.extend([fspath(input_file), fspath(prefix), 'pdf', 'txt']) args_tesseract.extend(tessconfig) try: - p = run(args_tesseract, stdout=PIPE, stderr=STDOUT, timeout=timeout, check=True) + p = run( + args_tesseract, + stdout=PIPE, + stderr=STDOUT, + timeout=timeout, + check=True, + env=_tesseract_env(omp_thread_limit), + ) stdout = p.stdout with suppress(FileNotFoundError): prefix.with_suffix('.txt').replace(output_text) diff --git a/src/ocrmypdf/builtin_plugins/tesseract_ocr.py b/src/ocrmypdf/builtin_plugins/tesseract_ocr.py index 447fc5d5..c9ae6aea 100644 --- a/src/ocrmypdf/builtin_plugins/tesseract_ocr.py +++ b/src/ocrmypdf/builtin_plugins/tesseract_ocr.py @@ -94,6 +94,13 @@ class TesseractOptions(BaseModel): user_patterns: Annotated[ str | None, Field(description="Path to Tesseract user patterns file") ] = None + omp_thread_limit: Annotated[ + int | None, + Field( + description="Calculated OMP_THREAD_LIMIT for Tesseract subprocesses", + exclude=True, + ), + ] = None @classmethod def add_arguments_to_parser(cls, parser, namespace: str = 'tesseract'): @@ -334,9 +341,10 @@ def validate(pdfinfo, options): if not os.environ.get('OMP_THREAD_LIMIT', '').isnumeric(): jobs = options.jobs or available_cpu_count() tess_threads = clamp(jobs // len(pdfinfo), 1, 3) - os.environ['OMP_THREAD_LIMIT'] = str(tess_threads) else: tess_threads = int(os.environ['OMP_THREAD_LIMIT']) + # Store the thread limit in options - it will be passed to subprocess env + options.tesseract.omp_thread_limit = tess_threads log.debug("Using Tesseract OpenMP thread limit %d", tess_threads) if ( @@ -408,6 +416,7 @@ class TesseractOcrEngine(OcrEngine): input_file, engine_mode=options.tesseract.oem, timeout=options.tesseract.non_ocr_timeout, + omp_thread_limit=options.tesseract.omp_thread_limit, ) @staticmethod @@ -417,6 +426,7 @@ class TesseractOcrEngine(OcrEngine): languages=options.languages, engine_mode=options.tesseract.oem, timeout=options.tesseract.non_ocr_timeout, + omp_thread_limit=options.tesseract.omp_thread_limit, ) @staticmethod @@ -433,6 +443,7 @@ class TesseractOcrEngine(OcrEngine): thresholding=options.tesseract.thresholding, user_words=options.tesseract.user_words, user_patterns=options.tesseract.user_patterns, + omp_thread_limit=options.tesseract.omp_thread_limit, ) @staticmethod @@ -449,6 +460,7 @@ class TesseractOcrEngine(OcrEngine): thresholding=options.tesseract.thresholding, user_words=options.tesseract.user_words, user_patterns=options.tesseract.user_patterns, + omp_thread_limit=options.tesseract.omp_thread_limit, ) diff --git a/src/ocrmypdf/subprocess/__init__.py b/src/ocrmypdf/subprocess/__init__.py index c9a0700f..23aa1612 100644 --- a/src/ocrmypdf/subprocess/__init__.py +++ b/src/ocrmypdf/subprocess/__init__.py @@ -23,13 +23,13 @@ from ocrmypdf.exceptions import MissingDependencyError log = logging.getLogger(__name__) Args = Sequence[Path | str] -OsEnviron = os._Environ # pylint: disable=protected-access +Environ = Mapping[str, str] | os._Environ # pylint: disable=protected-access def run( args: Args, *, - env: OsEnviron | None = None, + env: Environ | None = None, logs_errors_to_stdout: bool = False, check: bool = False, **kwargs, @@ -81,7 +81,7 @@ def run_polling_stderr( *, callback: Callable[[str], None], check: bool = False, - env: OsEnviron | None = None, + env: Environ | None = None, **kwargs, ) -> CompletedProcess: """Run a process like ``ocrmypdf.subprocess.run``, and poll stderr. @@ -116,8 +116,8 @@ def run_polling_stderr( def _fix_process_args( - args: Args, env: OsEnviron | None, kwargs -) -> tuple[Args, OsEnviron, logging.Logger, bool]: + args: Args, env: Environ | None, kwargs +) -> tuple[Args, Environ, logging.Logger, bool]: if not env: env = os.environ @@ -142,7 +142,7 @@ def get_version( *, version_arg: str = '--version', regex=r'(\d+(\.\d+)*)', - env: OsEnviron | None = None, + env: Environ | None = None, ) -> str: """Get the version of the specified program.