diff --git a/ocrmypdf/main.py b/ocrmypdf/main.py index f7f4dcab..99f64602 100755 --- a/ocrmypdf/main.py +++ b/ocrmypdf/main.py @@ -170,13 +170,16 @@ advanced = parser.add_argument_group( "Advanced", "Advanced options for power users") advanced.add_argument( - '--tesseract-config', default=[], type=list, action='append', + '--tesseract-config', action='append', metavar='CFG', default=[], help="additional Tesseract configuration files") +advanced.add_argument( + '--tesseract-pagesegmode', action='store', type=int, metavar='PSM', + help="set Tesseract page segmentation mode (see tesseract --help)") advanced.add_argument( '--pdf-renderer', choices=['auto', 'tesseract', 'hocr'], default='auto', help='choose OCR PDF renderer') advanced.add_argument( - '--tesseract-timeout', default=180.0, type=float, + '--tesseract-timeout', default=180.0, type=float, metavar='SECONDS', help='give up on OCR after the timeout, but copy the preprocessed page ' 'into the final output') @@ -529,6 +532,7 @@ def ocr_tesseract_hocr( timeout=options.tesseract_timeout, pageinfo_getter=partial(get_pageinfo, input_file, pdfinfo, pdfinfo_lock), + pagesegmode=options.tesseract_pagesegmode, log=log ) @@ -635,6 +639,7 @@ def tesseract_ocr_and_render_pdf( language=options.language, tessconfig=options.tesseract_config, timeout=options.tesseract_timeout, + pagesegmode=options.tesseract_pagesegmode, log=log) diff --git a/ocrmypdf/tesseract.py b/ocrmypdf/tesseract.py index 6e66213e..e1084555 100644 --- a/ocrmypdf/tesseract.py +++ b/ocrmypdf/tesseract.py @@ -77,17 +77,24 @@ def languages(): def generate_hocr(input_file, output_hocr, language: list, tessconfig: list, - timeout: float, pageinfo_getter, log): + timeout: float, pageinfo_getter, pagesegmode: int, log): badxml = os.path.splitext(output_hocr)[0] + '.badxml' args_tesseract = [ get_program('tesseract'), - '-l', '+'.join(language), + '-l', '+'.join(language) + ] + + if pagesegmode is not None: + args_tesseract.extend(['-psm', str(pagesegmode)]) + + args_tesseract.extend([ input_file, badxml, 'hocr' - ] + tessconfig + ] + tessconfig) + print(args_tesseract) p = Popen(args_tesseract, close_fds=True, stdout=PIPE, stderr=PIPE, universal_newlines=True) try: @@ -135,7 +142,7 @@ def generate_hocr(input_file, output_hocr, language: list, tessconfig: list, def generate_pdf(input_image, skip_pdf, output_pdf, language: list, - tessconfig: list, timeout: float, log): + tessconfig: list, timeout: float, pagesegmode: int, log): '''Use Tesseract to render a PDF. input_image -- image to analyze @@ -148,11 +155,17 @@ def generate_pdf(input_image, skip_pdf, output_pdf, language: list, args_tesseract = [ get_program('tesseract'), - '-l', '+'.join(language), + '-l', '+'.join(language) + ] + + if pagesegmode is not None: + args_tesseract.extend(['-psm', str(pagesegmode)]) + + args_tesseract.extend([ input_image, os.path.splitext(output_pdf)[0], # Tesseract appends suffix 'pdf' - ] + tessconfig + ] + tessconfig) p = Popen(args_tesseract, close_fds=True, stdout=PIPE, stderr=PIPE, universal_newlines=True) diff --git a/tests/spoof/tesseract_cache.py b/tests/spoof/tesseract_cache.py index df83a67f..f6586808 100755 --- a/tests/spoof/tesseract_cache.py +++ b/tests/spoof/tesseract_cache.py @@ -37,6 +37,11 @@ def main(): m.update(lang.encode()) except ValueError: pass + try: + psm = sys.argv[sys.argv.index('-psm') + 1] + m.update(psm.encode()) + except ValueError: + pass input_file = sys.argv[-3] output_file = sys.argv[-2] diff --git a/tests/test_main.py b/tests/test_main.py index 07db9820..3a652dfc 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -56,8 +56,8 @@ def check_ocrmypdf(input_basename, output_basename, *args, env=None): input_file = _make_input(input_basename) output_file = _make_output(output_basename) - sh, _, err = run_ocrmypdf_sh(input_file, output_file, *args, env=env) - assert sh.returncode == 0, err + sh, out, err = run_ocrmypdf_sh(input_file, output_file, *args, env=env) + assert sh.returncode == 0, dict(stdout=out, stderr=err) assert os.path.exists(output_file), "Output file not created" assert os.stat(output_file).st_size > 100, "PDF too small or empty" return output_file @@ -348,3 +348,18 @@ def test_encrypted(): p, out, err = run_ocrmypdf_env('skew-encrypted.pdf', 'wont_be_created.pdf') assert p.returncode == ExitCode.input_file assert out.find('password') + + +@pytest.mark.parametrize('renderer', [ + 'hocr', + 'tesseract', + ]) +def test_pagesegmode(renderer, spoof_tesseract_cache): + check_ocrmypdf( + 'skew.pdf', 'test_psm_%s.pdf' % renderer, + '--tesseract-pagesegmode', '7', + '-v', '1', + '--pdf-renderer', renderer, env=spoof_tesseract_cache) + + +