diff --git a/doctr/cli/main.py b/doctr/cli/main.py index cadc634551..6c5d5d1aaf 100644 --- a/doctr/cli/main.py +++ b/doctr/cli/main.py @@ -17,7 +17,10 @@ from doctr.models import ocr_predictor from doctr.version import __version__ -logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) +logger = logging.getLogger(__name__) + +# Name of the handler installed by the CLI, so that repeated `main` calls do not stack handlers +_CLI_HANDLER_NAME = "doctr-cli" # Canonical export formats and their aliases, mirroring `Document.export_as` FORMAT_ALIASES: dict[str, str] = { @@ -52,6 +55,29 @@ READING_DIRECTIONS = ["auto", "ltr", "rtl", "ttb-rtl", "ttb-ltr"] +def _setup_logging(quiet: bool = False) -> None: + """Configure the logging of the `doctr` package for the CLI + + Only the library logger is configured, so that the CLI output stays free of the third-party + logs the root logger would otherwise pick up. + + Args: + quiet: whether only errors should be logged + """ + doctr_logger = logging.getLogger("doctr") + + # `main` can be called several times in the same process, so the previous handler is replaced + for handler in list(doctr_logger.handlers): + if handler.name == _CLI_HANDLER_NAME: + doctr_logger.removeHandler(handler) + + handler = logging.StreamHandler() + handler.set_name(_CLI_HANDLER_NAME) + handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s")) + doctr_logger.addHandler(handler) + doctr_logger.setLevel(logging.ERROR if quiet else logging.INFO) + + def _resolve_format(args: argparse.Namespace) -> str: """Resolve the canonical export format, inferring it from the output extension if needed @@ -86,15 +112,15 @@ def _load_document(args: argparse.Namespace) -> list[np.ndarray]: pages.extend(DocumentFile.from_pdf(input_path, **pdf_kwargs)) else: pages.extend(DocumentFile.from_images(input_path)) - logging.info(f"Document loaded successfully from {input_path}") + logger.info(f"Document loaded successfully from {input_path}") except FileNotFoundError: - logging.error(f"File not found: {input_path}") + logger.error(f"File not found: {input_path}") sys.exit(1) except ValueError: - logging.error(f"File could not be read as a valid image or PDF: {input_path}") + logger.error(f"File could not be read as a valid image or PDF: {input_path}") sys.exit(1) except Exception as e: - logging.error(f"Error occurred while loading the document: {e}") + logger.error(f"Error occurred while loading the document: {e}") sys.exit(1) return pages @@ -114,7 +140,7 @@ def _resolve_device(device: str | None) -> torch.device: try: return torch.device(device) except (RuntimeError, ValueError) as e: - logging.error(f"Invalid device '{device}': {e}") + logger.error(f"Invalid device '{device}': {e}") sys.exit(1) @@ -131,7 +157,7 @@ def _set_thresholds(model: Any, args: argparse.Namespace) -> None: det_model = getattr(getattr(model, "det_predictor", None), "model", None) postprocessor = getattr(det_model, "postprocessor", None) if postprocessor is None: - logging.warning("The detection model exposes no postprocessor: --bin_thresh & --box_thresh are ignored") + logger.warning("The detection model exposes no postprocessor: --bin_thresh & --box_thresh are ignored") return if args.bin_thresh is not None: @@ -155,11 +181,11 @@ def _to_device(model: Any, args: argparse.Namespace) -> Any: model = model.to(device) if device.type == "cuda" and torch.cuda.get_device_capability(device) >= (8, 0): model = model.bfloat16() - logging.info(f"Model loaded on {device} with bfloat16 precision") + logger.info(f"Model loaded on {device} with bfloat16 precision") except (RuntimeError, AssertionError, ValueError) as e: - logging.error(f"Could not load the model on device '{device}': {e}") + logger.error(f"Could not load the model on device '{device}': {e}") sys.exit(1) - logging.info(f"Model loaded on {device}") + logger.info(f"Model loaded on {device}") return model @@ -175,7 +201,7 @@ def _build_predictor(args: argparse.Namespace) -> Any: # Region masking is resolved by the layout model, so it has to be enabled along with `ignore_regions` detect_layout = args.detect_layout or bool(args.ignore_regions) if detect_layout and not args.detect_layout: - logging.info("--ignore_regions requires the layout model: layout detection has been enabled") + logger.info("--ignore_regions requires the layout model: layout detection has been enabled") kwargs: dict[str, Any] = { "det_arch": args.det_arch, @@ -261,7 +287,7 @@ def _save_results(result: Any, fmt: str, args: argparse.Namespace) -> None: try: exported = result.export_as(fmt, **_export_kwargs(fmt, args)) except Exception as e: - logging.error(f"Results could not be exported as '{fmt}': {e}") + logger.error(f"Results could not be exported as '{fmt}': {e}") sys.exit(1) try: @@ -270,19 +296,19 @@ def _save_results(result: Any, fmt: str, args: argparse.Namespace) -> None: paths = _xml_paths(args.output, len(exported)) for path, (xml_bytes, _) in zip(paths, exported): path.write_bytes(xml_bytes) - logging.info(f"Results saved to {', '.join(str(path) for path in paths)}") + logger.info(f"Results saved to {', '.join(str(path) for path in paths)}") else: with open(args.output, "w", encoding="utf-8") as f: if fmt == "json": json.dump(exported, f, indent=args.indent, ensure_ascii=False) else: f.write(exported) - logging.info(f"Results saved to {args.output}") + logger.info(f"Results saved to {args.output}") except FileNotFoundError: - logging.error(f"Could not write output file at given path: {args.output}") + logger.error(f"Could not write output file at given path: {args.output}") sys.exit(1) except Exception as e: - logging.error(f"Results could not be saved: {e}") + logger.error(f"Results could not be saved: {e}") sys.exit(1) @@ -290,8 +316,7 @@ def main(argv=None): """Main function for the docTR CLI tool""" # parse command-line arguments and set up the model args = _parse_args(argv) - if args.quiet: - logging.getLogger().setLevel(logging.ERROR) + _setup_logging(args.quiet) fmt = _resolve_format(args) model = _build_predictor(args) @@ -300,7 +325,7 @@ def main(argv=None): doc = _load_document(args) # perform OCR - logging.info("Performing OCR...") + logger.info("Performing OCR...") result = model(doc) # save results to the requested format diff --git a/doctr/file_utils.py b/doctr/file_utils.py index b0f44bf0e2..28e738a0ca 100644 --- a/doctr/file_utils.py +++ b/doctr/file_utils.py @@ -8,6 +8,8 @@ __all__ = ["requires_package", "CLASS_NAME"] +logger = logging.getLogger(__name__) + CLASS_NAME: str = "words" ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} @@ -22,7 +24,7 @@ def requires_package(name: str, extra_message: str | None = None) -> None: # pr """ try: _pkg_version = importlib.metadata.version(name) - logging.info(f"{name} version {_pkg_version} available.") + logger.info(f"{name} version {_pkg_version} available.") except importlib.metadata.PackageNotFoundError: raise ImportError( f"\n\n{extra_message if extra_message is not None else ''} " diff --git a/doctr/models/factory/hub.py b/doctr/models/factory/hub.py index 5d75535472..f039ceae63 100644 --- a/doctr/models/factory/hub.py +++ b/doctr/models/factory/hub.py @@ -24,6 +24,8 @@ __all__ = ["login_to_hub", "push_to_hf_hub", "from_hub", "_save_model_and_config_for_hf_hub"] +logger = logging.getLogger(__name__) + AVAILABLE_ARCHS = { "classification": models.classification.zoo.ARCHS + models.classification.zoo.ORIENTATION_ARCHS, @@ -38,7 +40,7 @@ def login_to_hub() -> None: # pragma: no cover """Login to huggingface hub""" access_token = get_token() if access_token is not None: - logging.info("Huggingface Hub token found and valid") + logger.info("Huggingface Hub token found and valid") login(token=access_token) else: login() diff --git a/doctr/models/utils/pytorch.py b/doctr/models/utils/pytorch.py index bc4b68b968..3982be3d37 100644 --- a/doctr/models/utils/pytorch.py +++ b/doctr/models/utils/pytorch.py @@ -25,6 +25,8 @@ "_CompiledModule", ] +logger = logging.getLogger(__name__) + # torch compiled model type _CompiledModule = torch._dynamo.eval_frame.OptimizedModule @@ -58,7 +60,7 @@ def load_pretrained_params( **kwargs: additional arguments to be passed to `doctr.utils.data.download_from_url` """ if path_or_url is None: - logging.warning("No model URL or Path provided, using default initialization.") + logger.warning("No model URL or Path provided, using default initialization.") return archive_path = ( @@ -212,7 +214,7 @@ def export_model_to_onnx( verbose=False, **kwargs, ) - logging.info(f"Model exported to {model_name}.onnx") + logger.info(f"Model exported to {model_name}.onnx") return f"{model_name}.onnx" @@ -466,7 +468,7 @@ def _constrain_logits( if verbose: # pragma: no cover kept = sum(char in allowed for char in vocab) - logging.info( + logger.info( f"add_whitelist: {type(reco_model).__name__} - kept {kept}/{vocab_size} vocabulary " f"characters, forbade {vocab_size - kept}" + (f", reassigned {reassigned} to a nearest allowed character." if strategy == "nearest" else ".") diff --git a/doctr/utils/data.py b/doctr/utils/data.py index ba735423df..f8c09dc60c 100644 --- a/doctr/utils/data.py +++ b/doctr/utils/data.py @@ -17,6 +17,8 @@ __all__ = ["download_from_url"] +logger = logging.getLogger(__name__) + # matches bfd8deac from resnet18-bfd8deac.ckpt HASH_REGEX = re.compile(r"-([a-f0-9]*)\.") @@ -84,7 +86,7 @@ def download_from_url( file_path = folder_path.joinpath(file_name) # Check file existence if file_path.is_file() and (hash_prefix is None or _check_integrity(file_path, hash_prefix)): - logging.info(f"Using downloaded & verified file: {file_path}") + logger.info(f"Using downloaded & verified file: {file_path}") return file_path try: @@ -98,7 +100,7 @@ def download_from_url( error_message += ( ". You can change default cache directory using 'DOCTR_CACHE_DIR' environment variable if needed." ) - logging.error(error_message) + logger.error(error_message) raise # Download the file try: diff --git a/doctr/utils/fonts.py b/doctr/utils/fonts.py index ac93a707e8..56cd76aae4 100644 --- a/doctr/utils/fonts.py +++ b/doctr/utils/fonts.py @@ -11,6 +11,8 @@ __all__ = ["get_font"] +logger = logging.getLogger(__name__) + _FONT_CANDIDATES: dict[str, tuple[str, ...]] = { "Linux": ( "DejaVuSans.ttf", @@ -67,7 +69,7 @@ def get_font(font_family: str | None = None, font_size: int = 13) -> ImageFont.F try: return ImageFont.load_default(size=font_size) except TypeError: # pragma: no cover - logging.warning( + logger.warning( "Unable to load any recommended font family. Loading default PIL font, " "font size issues may be expected. " "To prevent this, it is recommended to specify the value of 'font_family'." diff --git a/doctr/utils/reconstitution.py b/doctr/utils/reconstitution.py index 62d7dc3797..ad99c903b5 100644 --- a/doctr/utils/reconstitution.py +++ b/doctr/utils/reconstitution.py @@ -15,6 +15,8 @@ __all__ = ["synthesize_page", "synthesize_kie_page"] +logger = logging.getLogger(__name__) + class _Word(NamedTuple): """A word to render, with its text, position, size and rotation.""" @@ -33,14 +35,14 @@ def _cached_font(font_family: str | None, font_size: int) -> ImageFont.FreeTypeF try: return get_font(font_family, max(font_size, 1)) except Exception: # pragma: no cover - logging.warning(f"Could not load font '{font_family}', falling back to the default font") + logger.warning(f"Could not load font '{font_family}', falling back to the default font") return get_font(None, max(font_size, 1)) @lru_cache(maxsize=1) def _warn_rotation_once() -> None: # pragma: no cover # lru_cache is thread-safe "warn once" semantics without a mutable global - logging.warning("Polygons with larger rotations may lead to slightly inaccurate rendering") + logger.warning("Polygons with larger rotations may lead to slightly inaccurate rendering") def _points(geometry: Any) -> list[tuple[float, float]] | None: @@ -168,7 +170,7 @@ def _draw_word( # Anchors are rejected by bitmap fonts, which would otherwise leave the page blank d.text(xy, anyascii(text), font=font, fill=fill) except Exception: - logging.warning(f"Could not render word: {text}") + logger.warning(f"Could not render word: {text}") def _paste_word( @@ -599,7 +601,7 @@ def _render( try: size = _entry_font_size(entry, words, polygon, w, h, font_family, min_font_size, max_font_size) except Exception as exc: - logging.warning(f"Could not size entry: {exc}") + logger.warning(f"Could not size entry: {exc}") continue prepared.append((entry, polygon, angle, box, words, size)) @@ -663,7 +665,7 @@ def _render( else: _synthesize_value(response, entry, polygon, angle, w, h, size, font_family, text_color, bold) except Exception as exc: - logging.warning(f"Could not render entry: {exc}") + logger.warning(f"Could not render entry: {exc}") if draw_proba: _draw_confidence(response, entry, box, font_family) diff --git a/pyproject.toml b/pyproject.toml index fab8ff82f3..d01a031d73 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -179,7 +179,7 @@ select = [ "E", "W", "F", "I", "N", "Q", "C4", "T10", "LOG", "D101", "D103", "D201","D202","D207","D208","D214","D215","D300","D301","D417", "D419", "D207" # pydocstyle ] -ignore = ["E402", "E203", "F403", "E731", "N812", "N817", "C408", "LOG015"] +ignore = ["E402", "E203", "F403", "E731", "N812", "N817", "C408"] [tool.ruff.lint.isort] known-first-party = ["doctr", "app", "utils"] diff --git a/references/classification/train_character.py b/references/classification/train_character.py index 22a9298ecb..efc159294a 100644 --- a/references/classification/train_character.py +++ b/references/classification/train_character.py @@ -255,7 +255,7 @@ def main(args): elif torch.cuda.is_available(): args.device = 0 else: - logging.warning("No accessible GPU, targe device set to CPU.") + logging.warning("No accessible GPU, targe device set to CPU.") # noqa: LOG015 if torch.cuda.is_available(): torch.cuda.set_device(args.device) model = model.cuda() diff --git a/references/classification/train_orientation.py b/references/classification/train_orientation.py index 0cbbf05616..3f9ca4800f 100644 --- a/references/classification/train_orientation.py +++ b/references/classification/train_orientation.py @@ -264,7 +264,7 @@ def main(args): elif torch.cuda.is_available(): args.device = 0 else: - logging.warning("No accessible GPU, targe device set to CPU.") + logging.warning("No accessible GPU, targe device set to CPU.") # noqa: LOG015 if torch.cuda.is_available(): torch.cuda.set_device(args.device) model = model.cuda() diff --git a/references/detection/train.py b/references/detection/train.py index 503f9c9bde..dcc3ef8845 100644 --- a/references/detection/train.py +++ b/references/detection/train.py @@ -236,7 +236,7 @@ def main(args): elif torch.cuda.is_available(): device = torch.device("cuda", 0) else: - logging.warning("No accessible GPU, target device set to CPU.") + logging.warning("No accessible GPU, target device set to CPU.") # noqa: LOG015 device = torch.device("cpu") slack_token = os.getenv("TQDM_SLACK_TOKEN") diff --git a/references/layout/train.py b/references/layout/train.py index b421573544..383cbb6d1c 100644 --- a/references/layout/train.py +++ b/references/layout/train.py @@ -238,7 +238,7 @@ def main(args): elif torch.cuda.is_available(): device = torch.device("cuda", 0) else: - logging.warning("No accessible GPU, target device set to CPU.") + logging.warning("No accessible GPU, target device set to CPU.") # noqa: LOG015 device = torch.device("cpu") slack_token = os.getenv("TQDM_SLACK_TOKEN") diff --git a/references/recognition/train.py b/references/recognition/train.py index 913f2ed017..ca86e704e4 100644 --- a/references/recognition/train.py +++ b/references/recognition/train.py @@ -216,7 +216,7 @@ def main(args): elif torch.cuda.is_available(): device = torch.device("cuda", 0) else: - logging.warning("No accessible GPU, target device set to CPU.") + logging.warning("No accessible GPU, target device set to CPU.") # noqa: LOG015 device = torch.device("cpu") slack_token = os.getenv("TQDM_SLACK_TOKEN") diff --git a/references/table/train.py b/references/table/train.py index 0505e60a0f..3701dd0e9f 100644 --- a/references/table/train.py +++ b/references/table/train.py @@ -200,7 +200,7 @@ def main(args): elif torch.cuda.is_available(): device = torch.device("cuda", 0) else: - logging.warning("No accessible GPU, target device set to CPU.") + logging.warning("No accessible GPU, target device set to CPU.") # noqa: LOG015 device = torch.device("cpu") slack_token = os.getenv("TQDM_SLACK_TOKEN") diff --git a/tests/common/test_utils_data.py b/tests/common/test_utils_data.py index 733aee266c..04a6da5a5b 100644 --- a/tests/common/test_utils_data.py +++ b/tests/common/test_utils_data.py @@ -1,3 +1,4 @@ +import logging import os from pathlib import PosixPath from unittest.mock import patch @@ -25,22 +26,22 @@ def test_download_from_url_customizing_cache_dir(mkdir_mock, urlretrieve_mock): @patch.dict(os.environ, {"HOME": "/"}, clear=True) @patch("pathlib.Path.mkdir", side_effect=OSError) -@patch("logging.error") -def test_download_from_url_error_creating_directory(logging_mock, mkdir_mock): - with pytest.raises(OSError): - download_from_url("test_url") - logging_mock.assert_called_with( +def test_download_from_url_error_creating_directory(mkdir_mock, caplog): + with caplog.at_level(logging.ERROR, logger="doctr.utils.data"): + with pytest.raises(OSError): + download_from_url("test_url") + assert ( "Failed creating cache directory at /.cache/doctr." " You can change default cache directory using 'DOCTR_CACHE_DIR' environment variable if needed." - ) + ) in caplog.text @patch.dict(os.environ, {"HOME": "/", "DOCTR_CACHE_DIR": "/test"}, clear=True) @patch("pathlib.Path.mkdir", side_effect=OSError) -@patch("logging.error") -def test_download_from_url_error_creating_directory_with_env_var(logging_mock, mkdir_mock): - with pytest.raises(OSError): - download_from_url("test_url") - logging_mock.assert_called_with( +def test_download_from_url_error_creating_directory_with_env_var(mkdir_mock, caplog): + with caplog.at_level(logging.ERROR, logger="doctr.utils.data"): + with pytest.raises(OSError): + download_from_url("test_url") + assert ( "Failed creating cache directory at /test using path from 'DOCTR_CACHE_DIR' environment variable." - ) + ) in caplog.text