From 23951c9e380b31719503fc01d43625049f537478 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Fri, 13 Oct 2023 03:25:12 -0700 Subject: [PATCH] Working HOCR folder to PDF converter --- src/ocrmypdf/__init__.py | 4 +- src/ocrmypdf/_pipelines/common.py | 12 ++++- src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py | 12 ++--- src/ocrmypdf/api.py | 59 ++++++++++++++++++++++ tests/test_api.py | 16 ++++++ 5 files changed, 92 insertions(+), 11 deletions(-) diff --git a/src/ocrmypdf/__init__.py b/src/ocrmypdf/__init__.py index d1e797dc..5f017798 100644 --- a/src/ocrmypdf/__init__.py +++ b/src/ocrmypdf/__init__.py @@ -11,7 +11,7 @@ from ocrmypdf import helpers, hocrtransform, pdfa, pdfinfo from ocrmypdf._concurrent import Executor from ocrmypdf._jobcontext import PageContext, PdfContext from ocrmypdf._version import PROGRAM_NAME, __version__ -from ocrmypdf.api import Verbosity, configure_logging, ocr, pdf_to_hocr +from ocrmypdf.api import Verbosity, configure_logging, hocr_to_ocr_pdf, ocr, pdf_to_hocr from ocrmypdf.exceptions import ( BadArgsError, DpiError, @@ -40,6 +40,7 @@ __all__ = [ 'ExitCode', 'ExitCodeException', 'helpers', + 'hocr_to_ocr_pdf', 'hocrtransform', 'hookimpl', 'InputFileError', @@ -49,6 +50,7 @@ __all__ = [ 'OrientationConfidence', 'OutputFileAccessError', 'PageContext', + 'pdf_to_hocr', 'pdfa', 'PdfContext', 'pdfinfo', diff --git a/src/ocrmypdf/_pipelines/common.py b/src/ocrmypdf/_pipelines/common.py index 5fef133c..41de222b 100644 --- a/src/ocrmypdf/_pipelines/common.py +++ b/src/ocrmypdf/_pipelines/common.py @@ -108,7 +108,11 @@ class HOCRResult: def __getstate__(self): """Return state values to be pickled.""" return { - k: (str(v) if k in ('pdf_page_from_image', 'hocr', 'textpdf') else v) + k: ( + ('Path://' + str(v)) + if k in ('pdf_page_from_image', 'hocr', 'textpdf') and v is not None + else v + ) for k, v in self.__dict__.items() } @@ -116,7 +120,11 @@ class HOCRResult: """Restore state from the unpickled state values.""" self.__dict__.update( { - k: (Path(v) if k in ('pdf_page_from_image', 'hocr', 'textpdf') else v) + k: ( + Path(v.removeprefix('Path://')) + if k in ('pdf_page_from_image', 'hocr', 'textpdf') and v is not None + else v + ) for k, v in state.items() } ) diff --git a/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py b/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py index 542789b3..46657e26 100644 --- a/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py +++ b/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py @@ -114,26 +114,22 @@ def run_hocr_to_ocr_pdf_pipeline( plugin_manager: OcrmypdfPluginManager, ) -> ExitCode: with manage_work_folder( - work_folder=options.input_folder, retain=True, print_location=False + work_folder=options.work_folder, retain=True, print_location=False ) as work_folder: executor = setup_pipeline(options, plugin_manager) origin_pdf = work_folder / 'origin.pdf' - shutil.copy2(options.input_file, origin_pdf) # Gather pdfinfo and create context pdfinfo = get_pdfinfo( - options.input_file, + origin_pdf, executor=executor, detailed_analysis=options.redo_ocr, progbar=options.progress_bar, max_workers=options.jobs if not options.use_threads else 1, # To help debug check_pages=options.pages, ) - context = PdfContext( - options, work_folder, options.input_file, pdfinfo, plugin_manager - ) - # Validate options are okay for this pdf - validate_pdfinfo_options(context) + context = PdfContext(options, work_folder, origin_pdf, pdfinfo, plugin_manager) + plugin_manager.hook.check_options(options=options) optimize_messages = exec_hocr_to_ocr_pdf(context, executor) return report_output_pdf(options, origin_pdf, optimize_messages) diff --git a/src/ocrmypdf/api.py b/src/ocrmypdf/api.py index bd527bce..d7577ccc 100644 --- a/src/ocrmypdf/api.py +++ b/src/ocrmypdf/api.py @@ -20,6 +20,7 @@ from warnings import warn import pluggy from ocrmypdf._logging import PageNumberFilter +from ocrmypdf._pipelines.hocr_to_ocr_pdf import run_hocr_to_ocr_pdf_pipeline from ocrmypdf._plugin_manager import get_plugin_manager from ocrmypdf._sync import run_hocr_pipeline, run_pipeline, run_pipeline_cli from ocrmypdf._validation import check_options @@ -454,6 +455,63 @@ def pdf_to_hocr( return run_hocr_pipeline(options=options, plugin_manager=plugin_manager) +def hocr_to_ocr_pdf( + work_folder: Path, + output_file: Path, + *, + jobs: int | None = None, + use_threads: bool | None = None, + optimize: int | None = None, + jpg_quality: int | None = None, + png_quality: int | None = None, + jbig2_lossy: bool | None = None, + jbig2_page_group_size: int | None = None, + jbig2_threshold: float | None = None, + pdfa_image_compression: str | None = None, + color_conversion_strategy: str | None = None, + fast_web_view: float | None = None, + plugin_manager=None, + plugins: Iterable[StrPath] | None = None, + **kwargs, +): + """Run OCRmyPDF and produces an output folder containing hOCR files.""" + # No new variable names should be assigned until these two steps are run + create_options_kwargs = { + k: v + for k, v in locals().items() + if k not in {'work_folder', 'output_pdf', 'kwargs'} + } + create_options_kwargs.update(kwargs) + + parser = get_parser() + + with _api_lock: + # We can't allow multiple ocrmypdf.ocr() threads to run in parallel, because + # they might install different plugins, and generally speaking we have areas + # of code that use global state. + + if not plugin_manager: + plugin_manager = get_plugin_manager(plugins) + plugin_manager.hook.add_options(parser=parser) # pylint: disable=no-member + + cmdline, deferred = _kwargs_to_cmdline( + defer_kwargs={'work_folder', 'output_file', 'plugins'}, + **create_options_kwargs, + ) + cmdline.append(str(work_folder)) + cmdline.append(str(output_file)) + parser.enable_api_mode() + options = parser.parse_args(cmdline) + for keyword, val in deferred.items(): + setattr(options, keyword, val) + delattr(options, 'input_file') + setattr(options, 'work_folder', work_folder) + + return run_hocr_to_ocr_pdf_pipeline( + options=options, plugin_manager=plugin_manager + ) + + __all__ = [ 'PageNumberFilter', 'Verbosity', @@ -466,4 +524,5 @@ __all__ = [ 'pdf_to_hocr', 'run_pipeline', 'run_pipeline_cli', + 'hocr_to_ocr_pdf', ] diff --git a/tests/test_api.py b/tests/test_api.py index 613c49f5..d7a0a200 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -39,3 +39,19 @@ def test_hocr_api(resources: Path, outdir: Path): assert (outdir / '000006_ocr_hocr.hocr').exists() assert not (outdir / '000004_ocr_hocr.hocr').exists() + + +def test_hocr_to_pdf_api(resources: Path, outdir: Path, outpdf: Path): + ocrmypdf.pdf_to_hocr( + resources / 'ccitt.pdf', + outdir, + language='eng', + skip_text=True, + plugins=['tests/plugins/tesseract_cache.py'], + ) + assert (outdir / '000001_ocr_hocr.hocr').exists() + hocr = (outdir / '000001_ocr_hocr.hocr').read_text(encoding='utf-8') + mangled = hocr.replace('the', 'hocr') + (outdir / '000001_ocr_hocr.hocr').write_text(mangled, encoding='utf-8') + + ocrmypdf.hocr_to_ocr_pdf(outdir, outpdf)