Refactor tesseract thresholding to use enum type
Replace integer-based thresholding parameter with ThresholdingMethod enum for improved type safety. The CLI still accepts the same string values (auto, otsu, adaptive-otsu, sauvola) but internally uses a strongly-typed enum. This makes the code more maintainable and catches type errors at development time.
This commit is contained in:
@@ -8,6 +8,7 @@ from __future__ import annotations
|
||||
import logging
|
||||
import re
|
||||
from contextlib import suppress
|
||||
from enum import IntEnum
|
||||
from math import pi
|
||||
from os import fspath
|
||||
from pathlib import Path
|
||||
@@ -26,11 +27,21 @@ from ocrmypdf.subprocess import get_version, run
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ThresholdingMethod(IntEnum):
|
||||
"""Tesseract thresholding methods for image binarization."""
|
||||
|
||||
AUTO = 0
|
||||
OTSU = 0 # Alias for AUTO - uses Tesseract's default (legacy Otsu)
|
||||
ADAPTIVE_OTSU = 1
|
||||
SAUVOLA = 2
|
||||
|
||||
|
||||
# Legacy dictionary for backward compatibility
|
||||
TESSERACT_THRESHOLDING_METHODS: dict[str, int] = {
|
||||
'auto': 0,
|
||||
'otsu': 0,
|
||||
'adaptive-otsu': 1,
|
||||
'sauvola': 2,
|
||||
'auto': ThresholdingMethod.AUTO,
|
||||
'otsu': ThresholdingMethod.OTSU,
|
||||
'adaptive-otsu': ThresholdingMethod.ADAPTIVE_OTSU,
|
||||
'sauvola': ThresholdingMethod.SAUVOLA,
|
||||
}
|
||||
|
||||
|
||||
@@ -294,7 +305,7 @@ def generate_hocr(
|
||||
tessconfig: list[str],
|
||||
timeout: float,
|
||||
pagesegmode: int,
|
||||
thresholding: int,
|
||||
thresholding: ThresholdingMethod,
|
||||
user_words,
|
||||
user_patterns,
|
||||
) -> None:
|
||||
@@ -306,7 +317,7 @@ def generate_hocr(
|
||||
if pagesegmode is not None:
|
||||
args_tesseract.extend(['--psm', str(pagesegmode)])
|
||||
|
||||
if thresholding != 0 and has_thresholding():
|
||||
if thresholding != ThresholdingMethod.AUTO and has_thresholding():
|
||||
args_tesseract.extend(['-c', f'thresholding_method={thresholding}'])
|
||||
|
||||
if user_words:
|
||||
@@ -360,7 +371,7 @@ def generate_pdf(
|
||||
tessconfig: list[str],
|
||||
timeout: float,
|
||||
pagesegmode: int,
|
||||
thresholding: int,
|
||||
thresholding: ThresholdingMethod,
|
||||
user_words,
|
||||
user_patterns,
|
||||
) -> None:
|
||||
@@ -376,7 +387,7 @@ def generate_pdf(
|
||||
|
||||
args_tesseract.extend(['-c', 'textonly_pdf=1'])
|
||||
|
||||
if thresholding != 0 and has_thresholding():
|
||||
if thresholding != ThresholdingMethod.AUTO and has_thresholding():
|
||||
args_tesseract.extend(['-c', f'thresholding_method={thresholding}'])
|
||||
|
||||
if user_words:
|
||||
|
||||
@@ -14,8 +14,9 @@ from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
from ocrmypdf import hookimpl
|
||||
from ocrmypdf._exec import tesseract
|
||||
from ocrmypdf._exec.tesseract import ThresholdingMethod
|
||||
from ocrmypdf._jobcontext import PageContext
|
||||
from ocrmypdf.cli import numeric, str_to_int
|
||||
from ocrmypdf.cli import numeric
|
||||
from ocrmypdf.exceptions import BadArgsError, MissingDependencyError
|
||||
from ocrmypdf.helpers import available_cpu_count, clamp
|
||||
from ocrmypdf.imageops import calculate_downsample, downsample_image
|
||||
@@ -25,6 +26,34 @@ from ocrmypdf.subprocess import check_external_program
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _thresholding_method_converter(value: str) -> ThresholdingMethod:
|
||||
"""Convert string argument to ThresholdingMethod enum.
|
||||
|
||||
Args:
|
||||
value: String name of thresholding method (auto, otsu, adaptive-otsu, sauvola)
|
||||
|
||||
Returns:
|
||||
ThresholdingMethod enum value
|
||||
|
||||
Raises:
|
||||
argparse.ArgumentTypeError: If value is not a valid thresholding method
|
||||
"""
|
||||
method_map = {
|
||||
'auto': ThresholdingMethod.AUTO,
|
||||
'otsu': ThresholdingMethod.OTSU,
|
||||
'adaptive-otsu': ThresholdingMethod.ADAPTIVE_OTSU,
|
||||
'sauvola': ThresholdingMethod.SAUVOLA,
|
||||
}
|
||||
if value.lower() not in method_map:
|
||||
import argparse
|
||||
|
||||
valid = ', '.join(method_map.keys())
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"Invalid thresholding method '{value}'. Must be one of: {valid}"
|
||||
)
|
||||
return method_map[value.lower()]
|
||||
|
||||
|
||||
class TesseractOptions(BaseModel):
|
||||
"""Options specific to Tesseract OCR engine."""
|
||||
|
||||
@@ -39,8 +68,9 @@ class TesseractOptions(BaseModel):
|
||||
int | None, Field(ge=0, le=3, description="Set Tesseract OCR engine mode")
|
||||
] = None
|
||||
thresholding: Annotated[
|
||||
int | None, Field(description="Set Tesseract input image thresholding mode")
|
||||
] = None
|
||||
ThresholdingMethod,
|
||||
Field(description="Set Tesseract input image thresholding mode"),
|
||||
] = ThresholdingMethod.AUTO
|
||||
timeout: Annotated[
|
||||
float, Field(ge=0, description="Timeout for OCR operations in seconds")
|
||||
] = 180.0
|
||||
@@ -115,16 +145,16 @@ class TesseractOptions(BaseModel):
|
||||
tess.add_argument(
|
||||
f'--{namespace}-thresholding',
|
||||
action='store',
|
||||
type=str_to_int(tesseract.TESSERACT_THRESHOLDING_METHODS),
|
||||
type=_thresholding_method_converter,
|
||||
default='auto',
|
||||
metavar='METHOD',
|
||||
dest=f'{namespace}_thresholding',
|
||||
help=(
|
||||
"Set Tesseract 5.0+ input image thresholding mode. This may improve OCR "
|
||||
"results on low quality images or those that contain high contrast color. "
|
||||
"legacy-otsu is the Tesseract default; adaptive-otsu is an improved Otsu "
|
||||
"algorithm with improved sort for background color changes; sauvola is "
|
||||
"based on local standard deviation."
|
||||
"Set Tesseract 5.0+ input image thresholding mode. This may improve "
|
||||
"OCR results on low quality images or those that contain high "
|
||||
"contrast color. Options: auto, otsu, adaptive-otsu, sauvola. "
|
||||
"auto/otsu is the Tesseract default (legacy Otsu); adaptive-otsu "
|
||||
"is an improved Otsu algorithm with improved sort for background "
|
||||
"color changes; sauvola is based on local standard deviation."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -226,10 +256,7 @@ class TesseractOptions(BaseModel):
|
||||
@model_validator(mode='after')
|
||||
def validate_downsample_consistency(self):
|
||||
"""Validate downsample options are consistent."""
|
||||
if (
|
||||
self.downsample_above != 32767
|
||||
and not self.downsample_large_images
|
||||
):
|
||||
if self.downsample_above != 32767 and not self.downsample_large_images:
|
||||
log.warning(
|
||||
"The --tesseract-downsample-above argument will have no effect unless "
|
||||
"--tesseract-downsample-large-images is also given."
|
||||
@@ -283,7 +310,10 @@ def check_options(options):
|
||||
)
|
||||
|
||||
# Check version-specific feature compatibility
|
||||
if not tesseract.has_thresholding() and options.tesseract.thresholding != 0:
|
||||
if (
|
||||
not tesseract.has_thresholding()
|
||||
and options.tesseract.thresholding != ThresholdingMethod.AUTO
|
||||
):
|
||||
log.warning(
|
||||
"The installed version of Tesseract does not support changes to its "
|
||||
"thresholding method. The --tesseract-threshold argument will be "
|
||||
|
||||
Reference in New Issue
Block a user