From 40cf657d872765e52a7067e7ec16eabb4175a7e2 Mon Sep 17 00:00:00 2001 From: felix Date: Fri, 9 Oct 2026 09:36:32 +0200 Subject: [PATCH 1/4] reset From 25effb416f851abff9113f42d45561fab6a68ddb Mon Sep 17 00:00:00 2001 From: felix Date: Fri, 9 Oct 2026 10:39:41 +0200 Subject: [PATCH 2/4] Add confidence aggregation --- .pre-commit-config.yaml | 2 +- docs/source/using_doctr/using_cli.rst | 1 + docs/source/using_doctr/using_models.rst | 36 ++++++ doctr/cli/main.py | 10 ++ doctr/models/_utils.py | 91 +++++++++++++- doctr/models/recognition/core.py | 15 ++- doctr/models/recognition/crnn/pytorch.py | 48 +++++-- doctr/models/recognition/master/base.py | 5 +- doctr/models/recognition/master/pytorch.py | 25 ++-- doctr/models/recognition/parseq/base.py | 5 +- doctr/models/recognition/parseq/pytorch.py | 10 +- doctr/models/recognition/predictor/_utils.py | 6 +- doctr/models/recognition/predictor/pytorch.py | 5 +- doctr/models/recognition/sar/pytorch.py | 24 +++- doctr/models/recognition/viptr/pytorch.py | 48 +++++-- doctr/models/recognition/vitstr/base.py | 5 +- doctr/models/recognition/vitstr/pytorch.py | 10 +- doctr/models/zoo.py | 27 ++++ pyproject.toml | 2 +- tests/common/test_cli.py | 9 ++ tests/common/test_models.py | 45 ++++++- .../test_models_recognition_predictor.py | 10 ++ tests/pytorch/test_models_recognition_pt.py | 119 +++++++++++++++++- tests/pytorch/test_models_zoo_pt.py | 69 ++++++++++ 24 files changed, 574 insertions(+), 53 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6c88e28176..5407043a3f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -16,7 +16,7 @@ repos: - id: no-commit-to-branch args: ['--branch', 'main'] - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.16.8 + rev: v0.16.10 hooks: - id: ruff args: [ --fix ] diff --git a/docs/source/using_doctr/using_cli.rst b/docs/source/using_doctr/using_cli.rst index e6c2fde2bd..c1f4b37c16 100644 --- a/docs/source/using_doctr/using_cli.rst +++ b/docs/source/using_doctr/using_cli.rst @@ -47,6 +47,7 @@ The CLI supports a variety of arguments to fine-tune the detection, recognition * ``--reco_bs``: Batch size used for the recognition model. *Default: 128* * ``--bin_thresh``: Binarization threshold of the detection post-processing. *Default: the value of the architecture* * ``--box_thresh``: Minimum confidence of a detected box. *Default: the value of the architecture* +* ``--confidence_aggregation``: Aggregation method of the character probabilities into the word confidence, one of ``mean``, ``min``, ``max``, ``median``, ``geometric_mean``, ``harmonic_mean``. *Default: the method of the recognition architecture* * ``--detect_orientation``: Enables automatic detection of page orientation. *Default: False* * ``--detect_language``: Enables language detection for the extracted text. *Default: False* * ``--detect_layout``: Attaches the detected layout regions to each page. *Default: False* diff --git a/docs/source/using_doctr/using_models.rst b/docs/source/using_doctr/using_models.rst index ae44678526..af4012d076 100644 --- a/docs/source/using_doctr/using_models.rst +++ b/docs/source/using_doctr/using_models.rst @@ -698,6 +698,42 @@ This is useful to detect (possible less) text regions more accurately with a hig out = predictor([input_page]) +* Modify how the word confidences are aggregated. + +The character probabilities of the recognition model are aggregated into the word confidence with `"mean"`, `"min"`, +`"max"`, `"median"`, `"geometric_mean"`, `"harmonic_mean"` or a Python callable taking the +1D array of probabilities and returning a float. The default is `"min"` for the CRNN, VIPTR, SAR and MASTER models and +`"mean"` for the ViTSTR and PARSeq models. +A mean can hide a single uncertain character: for `[0.99, 0.98, 0.12, 0.99]`, the mean is `0.77`, the geometric mean +`0.58`, the harmonic mean `0.35` and the minimum `0.12`. `"max"` and `"median"` ignore most uncertain characters, so +they hardly separate the misread words from the correct ones. + +Only the probabilities of the predicted characters are aggregated: the CTC blank frames (CRNN, VIPTR) and the +positions from the end-of-sequence token onwards (SAR, MASTER) are ignored, a CTC character has the highest +probability of its frames, and an empty prediction has a confidence of `0`. The parts of a split wide crop are +aggregated with the `split_confidence_aggregation` of the recognition predictor (`"min"` by default). + +Pass `confidence_aggregation` to `ocr_predictor`, `kie_predictor` or `recognition_predictor` (a model instance passed +to them is modified), or change it afterwards on the postprocessor of the recognition model. + +.. code:: python3 + + from doctr.models import ocr_predictor, parseq + + # Aggregate the character probabilities into the word confidence with the mean + predictor = ocr_predictor(pretrained=True, confidence_aggregation="mean") + + # Or modify the aggregation method of the recognition model afterwards + predictor = ocr_predictor('db_resnet50', 'crnn_vgg16_bn', pretrained=True) + predictor.reco_predictor.model.postprocessor.confidence_aggregation = "geometric_mean" + + # Modify the aggregation of the confidences of the parts of a split wide crop + predictor.reco_predictor.split_confidence_aggregation = "mean" + + # Or set it when instantiating a model + predictor = ocr_predictor('db_resnet50', parseq(pretrained=True, confidence_aggregation="min"), pretrained=True) + + * Disable page orientation classification If you deal with documents which contains only small rotations (~ -45 to 45 degrees), you can disable the page orientation classification to speed up the inference. diff --git a/doctr/cli/main.py b/doctr/cli/main.py index 0da54c211c..e40c1bd2a1 100644 --- a/doctr/cli/main.py +++ b/doctr/cli/main.py @@ -17,6 +17,7 @@ from doctr.io import DocumentFile from doctr.io.figures import IMAGE_FORMATS, IMAGE_MODES, FigureEncoder from doctr.models import ocr_predictor +from doctr.models._utils import _CONFIDENCE_AGGREGATIONS from doctr.version import __version__ logger = logging.getLogger(__name__) @@ -221,6 +222,7 @@ def _build_predictor(args: argparse.Namespace) -> Any: "ignore_regions": args.ignore_regions or None, "det_bs": args.det_bs, "reco_bs": args.reco_bs, + "confidence_aggregation": args.confidence_aggregation, # `_OCRPredictor` keyword args "disable_page_orientation": args.disable_page_orientation, "disable_crop_orientation": args.disable_crop_orientation, @@ -447,6 +449,14 @@ def _parse_args(argv=None): default=None, help="minimal objectness score to consider a detected box (defaults to the architecture value)", ) + parser.add_argument( + "--confidence_aggregation", + type=str, + default=None, + choices=list(_CONFIDENCE_AGGREGATIONS), + help="aggregation method of the character probabilities into the word confidence (defaults to the method of " + "the recognition architecture)", + ) parser.add_argument("--det_bs", type=int, default=2, help="batch size for detection") parser.add_argument("--reco_bs", type=int, default=128, help="batch size for recognition") parser.add_argument("--detect_orientation", action="store_true", help="automatically detect page orientation") diff --git a/doctr/models/_utils.py b/doctr/models/_utils.py index f5f1423652..a0e72159f8 100644 --- a/doctr/models/_utils.py +++ b/doctr/models/_utils.py @@ -4,17 +4,21 @@ # See LICENSE or go to for full license details. import heapq -from math import floor +from collections.abc import Callable, Sequence +from math import floor, isnan from statistics import median_low -from typing import Any +from typing import Any, Literal, TypeAlias import cv2 import numpy as np +import torch from langdetect import LangDetectException, detect_langs from doctr.utils.geometry import rotate_image __all__ = [ + "ConfidenceAggregation", + "aggregate_confidence", "estimate_orientation", "get_language", "invert_data_structure", @@ -23,6 +27,89 @@ "rectify_loc_preds", ] +ConfidenceAggregation: TypeAlias = ( + Literal["mean", "min", "max", "median", "geometric_mean", "harmonic_mean"] | Callable[[np.ndarray], float] +) + + +def _geometric_mean(scores: np.ndarray) -> float: + # numerically stable form of prod(scores) ** (1 / N), a zero score yields 0 + with np.errstate(divide="ignore"): + return float(np.exp(np.mean(np.log(scores, dtype=np.float64)))) + + +def _harmonic_mean(scores: np.ndarray) -> float: + # a zero score yields 0 + with np.errstate(divide="ignore"): + return float(scores.size / np.sum(1.0 / scores.astype(np.float64))) + + +_CONFIDENCE_AGGREGATIONS: dict[str, Callable[[np.ndarray], Any]] = { + "mean": np.mean, + "min": np.min, + "max": np.max, + "median": np.median, + "geometric_mean": _geometric_mean, + "harmonic_mean": _harmonic_mean, +} + + +def _resolve_confidence_aggregation(method: ConfidenceAggregation) -> Callable[[np.ndarray], Any]: + """Return the function of a confidence aggregation method, raise a `ValueError` if it is unknown""" + if isinstance(method, torch.nn.Module): + raise ValueError( + f"Unsupported confidence aggregation {type(method).__name__}: " + f"expected a Python callable or one of {list(_CONFIDENCE_AGGREGATIONS)}" + ) + if callable(method): + return method + # the type check comes first: an unhashable value (e.g. a list) can not be looked up + if not isinstance(method, str) or method not in _CONFIDENCE_AGGREGATIONS: + raise ValueError( + f"Unknown confidence aggregation {method!r}, expected a callable or one of {list(_CONFIDENCE_AGGREGATIONS)}" + ) + return _CONFIDENCE_AGGREGATIONS[method] + + +def _confidence_aggregation_repr(method: ConfidenceAggregation) -> str: + """Return the name of a confidence aggregation method to display it in a `repr`""" + return repr(method) if isinstance(method, str) else getattr(method, "__name__", type(method).__name__) + + +def aggregate_confidence( + scores: np.ndarray | Sequence[float], + method: ConfidenceAggregation = "mean", +) -> float: + """Aggregate several scores (e.g. the character probabilities of a word) into a single confidence score + + >>> from doctr.models._utils import aggregate_confidence + >>> aggregate_confidence([0.99, 0.98, 0.12, 0.99], "min") + 0.12 + + Available methods: + - "mean": arithmetic mean, a single low score can be hidden by many high ones + - "min": the lowest score, a single low score is enough to lower the confidence + - "max": the highest score + - "median": the median score, robust to outliers + - "geometric_mean": `exp(mean(log(scores)))`, i.e. the exponential of the average log-probability + - "harmonic_mean": strongly penalizes low scores + - a Python callable taking the 1D array of scores and returning a float + + Args: + scores: the scores to aggregate, expected in [0, 1] + method: the aggregation method + + Returns: + the aggregated confidence clipped to [0, 1], 0 if there is no score or if the aggregation is NaN (e.g. a + callable averaging an empty selection) + """ + aggregate = _resolve_confidence_aggregation(method) + scores = np.asarray(scores).ravel() + if scores.size == 0: + return 0.0 + value = float(np.asarray(aggregate(scores)).squeeze()) + return 0.0 if isnan(value) else min(max(value, 0.0), 1.0) + def get_max_width_length_ratio(contour: np.ndarray) -> float: """Get the maximum shape ratio of a contour. diff --git a/doctr/models/recognition/core.py b/doctr/models/recognition/core.py index 7b12602cd5..e3cc4d2375 100644 --- a/doctr/models/recognition/core.py +++ b/doctr/models/recognition/core.py @@ -7,6 +7,11 @@ import numpy as np from doctr.datasets import encode_sequences +from doctr.models._utils import ( + ConfidenceAggregation, + _confidence_aggregation_repr, + _resolve_confidence_aggregation, +) from doctr.utils.repr import NestedObject __all__ = ["RecognitionPostProcessor", "RecognitionModel"] @@ -41,14 +46,22 @@ class RecognitionPostProcessor(NestedObject): Args: vocab: string containing the ordered sequence of supported characters + confidence_aggregation: aggregation method of the character probabilities into the word confidence: + "mean", "min", "max", "median", "geometric_mean", "harmonic_mean" or a callable """ def __init__( self, vocab: str, + confidence_aggregation: ConfidenceAggregation = "mean", ) -> None: + _resolve_confidence_aggregation(confidence_aggregation) self.vocab = vocab + self.confidence_aggregation = confidence_aggregation self._embedding = list(self.vocab) + [""] def extra_repr(self) -> str: - return f"vocab_size={len(self.vocab)}" + return ( + f"vocab_size={len(self.vocab)}, " + f"confidence_aggregation={_confidence_aggregation_repr(self.confidence_aggregation)}" + ) diff --git a/doctr/models/recognition/crnn/pytorch.py b/doctr/models/recognition/crnn/pytorch.py index 8cfd49e960..0475ab227b 100644 --- a/doctr/models/recognition/crnn/pytorch.py +++ b/doctr/models/recognition/crnn/pytorch.py @@ -5,15 +5,16 @@ from collections.abc import Callable from copy import deepcopy -from itertools import groupby from typing import Any +import numpy as np import torch from torch import nn from torch.nn import functional as F from doctr.datasets import VOCABS, decode_sequence +from ..._utils import ConfidenceAggregation, aggregate_confidence from ...classification import mobilenet_v3_large_r, mobilenet_v3_small_r, vgg16_bn_r from ...utils import load_pretrained_params from ..core import RecognitionModel, RecognitionPostProcessor @@ -50,13 +51,22 @@ class CTCPostProcessor(RecognitionPostProcessor): Args: vocab: string containing the ordered sequence of supported characters + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ + def __init__( + self, + vocab: str, + confidence_aggregation: ConfidenceAggregation = "min", + ) -> None: + super().__init__(vocab, confidence_aggregation) + @staticmethod def ctc_best_path( logits: torch.Tensor, vocab: str = VOCABS["french"], blank: int = 0, + confidence_aggregation: ConfidenceAggregation = "min", ) -> list[tuple[str, float]]: """Implements best path decoding as shown by Graves (Dissertation, p63), highly inspired from `_. @@ -65,18 +75,27 @@ def ctc_best_path( logits: model output, shape: N x T x C vocab: vocabulary to use blank: index of blank label + confidence_aggregation: aggregation method of the character probabilities into the word confidence Returns: A list of tuples: (word, confidence) """ - # Gather the most confident characters, and assign the smallest conf among those to the sequence prob - probs = F.softmax(logits, dim=-1).max(dim=-1).values.min(dim=1).values - - # collapse best path (using itertools.groupby), map to chars, join char list to string - best_paths = torch.argmax(logits, dim=-1).detach().cpu().numpy() - words = [decode_sequence([k for k, _ in groupby(seq.tolist()) if k != blank], vocab) for seq in best_paths] - - return list(zip(words, probs.tolist())) + best_paths = logits.argmax(dim=-1).detach().cpu().numpy() + probs = F.softmax(logits.float(), dim=-1).max(dim=-1).values.detach().cpu().numpy() + + results = [] + for path, path_probs in zip(best_paths, probs): + # Collapse the repeated labels: the probability of a character is the highest one within its run + run_starts = np.flatnonzero(np.r_[True, path[1:] != path[:-1]]) + labels, label_probs = path[run_starts], np.maximum.reduceat(path_probs, run_starts) + # Remove the blanks + is_char = labels != blank + results.append(( + decode_sequence(labels[is_char].tolist(), vocab), + aggregate_confidence(label_probs[is_char], confidence_aggregation), + )) + + return results def __call__(self, logits: torch.Tensor) -> list[tuple[str, float]]: """Performs decoding of raw output with CTC and decoding of CTC predictions @@ -90,7 +109,12 @@ def __call__(self, logits: torch.Tensor) -> list[tuple[str, float]]: """ # Decode CTC - return self.ctc_best_path(logits=logits, vocab=self.vocab, blank=len(self.vocab)) + return self.ctc_best_path( + logits=logits, + vocab=self.vocab, + blank=len(self.vocab), + confidence_aggregation=self.confidence_aggregation, + ) class CRNN(RecognitionModel, nn.Module): @@ -103,6 +127,7 @@ class CRNN(RecognitionModel, nn.Module): rnn_units: number of units in the LSTM layers exportable: onnx exportable returns only logits cfg: configuration dictionary + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ _children_names: list[str] = ["feat_extractor", "decoder", "linear", "postprocessor"] @@ -115,6 +140,7 @@ def __init__( input_shape: tuple[int, int, int] = (3, 32, 128), exportable: bool = False, cfg: dict[str, Any] | None = None, + confidence_aggregation: ConfidenceAggregation = "min", ) -> None: super().__init__() self.vocab = vocab @@ -139,7 +165,7 @@ def __init__( # features units = 2 * rnn_units because bidirectional layers self.linear = nn.Linear(in_features=2 * rnn_units, out_features=len(vocab) + 1) - self.postprocessor = CTCPostProcessor(vocab=vocab) + self.postprocessor = CTCPostProcessor(vocab=vocab, confidence_aggregation=confidence_aggregation) for n, m in self.named_modules(): # Don't override the initialization of the backbone diff --git a/doctr/models/recognition/master/base.py b/doctr/models/recognition/master/base.py index 1b6112e18f..6a063f689a 100644 --- a/doctr/models/recognition/master/base.py +++ b/doctr/models/recognition/master/base.py @@ -7,6 +7,7 @@ import numpy as np from ....datasets import encode_sequences +from ..._utils import ConfidenceAggregation from ..core import RecognitionPostProcessor @@ -44,11 +45,13 @@ class _MASTERPostProcessor(RecognitionPostProcessor): Args: vocab: string containing the ordered sequence of supported characters + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( self, vocab: str, + confidence_aggregation: ConfidenceAggregation = "min", ) -> None: - super().__init__(vocab) + super().__init__(vocab, confidence_aggregation) self._embedding = list(vocab) + [""] + [""] + [""] diff --git a/doctr/models/recognition/master/pytorch.py b/doctr/models/recognition/master/pytorch.py index 0bbd2bf8fa..69105607a7 100644 --- a/doctr/models/recognition/master/pytorch.py +++ b/doctr/models/recognition/master/pytorch.py @@ -7,6 +7,7 @@ from copy import deepcopy from typing import Any +import numpy as np import torch from torch import nn from torch.nn import functional as F @@ -16,6 +17,7 @@ from doctr.models.classification import magc_resnet31 from doctr.models.modules.transformer import Decoder, PositionalEncoding +from ..._utils import ConfidenceAggregation, aggregate_confidence from ...utils import _bf16_to_float32, load_pretrained_params from .base import _MASTER, _MASTERPostProcessor @@ -49,6 +51,7 @@ class MASTER(_MASTER, nn.Module): input_shape: size of the image inputs exportable: onnx exportable returns only logits cfg: dictionary containing information about the model + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( @@ -64,6 +67,7 @@ def __init__( input_shape: tuple[int, int, int] = (3, 32, 128), # different from the paper exportable: bool = False, cfg: dict[str, Any] | None = None, + confidence_aggregation: ConfidenceAggregation = "min", ) -> None: super().__init__() @@ -88,7 +92,7 @@ def __init__( ) self.linear = nn.Linear(self.d_model, self.vocab_size + 3) - self.postprocessor = MASTERPostProcessor(vocab=self.vocab) + self.postprocessor = MASTERPostProcessor(vocab=self.vocab, confidence_aggregation=confidence_aggregation) for n, m in self.named_modules(): # Don't override the initialization of the backbone @@ -264,19 +268,24 @@ def __call__( logits: torch.Tensor, ) -> list[tuple[str, float]]: # compute pred with argmax for attention models - out_idxs = logits.argmax(-1) + out_idxs = logits.argmax(-1).detach().cpu().numpy() # N x L - probs = torch.gather(torch.softmax(logits, -1), -1, out_idxs.unsqueeze(-1)).squeeze(-1) - # Take the minimum confidence of the sequence - probs = probs.min(dim=1).values.detach().cpu() + preds_prob = torch.softmax(logits, -1).max(dim=-1)[0].detach().cpu().numpy() # Manual decoding word_values = [ - "".join(self._embedding[idx] for idx in encoded_seq).split("")[0] - for encoded_seq in out_idxs.detach().cpu().numpy() + "".join(self._embedding[idx] for idx in encoded_seq).split("")[0] for encoded_seq in out_idxs + ] + # aggregate the character probabilities of each word up to the EOS token: the number of predicted tokens is + # used since the and tokens are decoded as several characters + is_eos = out_idxs == len(self.vocab) + seq_lens = np.where(is_eos.any(axis=-1), is_eos.argmax(axis=-1), out_idxs.shape[-1]) + probs = [ + aggregate_confidence(preds_prob[i, :seq_len], self.confidence_aggregation) + for i, seq_len in enumerate(seq_lens) ] - return list(zip(word_values, probs.numpy().clip(0, 1).tolist())) + return list(zip(word_values, probs)) def _master( diff --git a/doctr/models/recognition/parseq/base.py b/doctr/models/recognition/parseq/base.py index b50f9ff845..c828f194e4 100644 --- a/doctr/models/recognition/parseq/base.py +++ b/doctr/models/recognition/parseq/base.py @@ -7,6 +7,7 @@ import numpy as np from ....datasets import encode_sequences +from ..._utils import ConfidenceAggregation from ..core import RecognitionPostProcessor @@ -44,11 +45,13 @@ class _PARSeqPostProcessor(RecognitionPostProcessor): Args: vocab: string containing the ordered sequence of supported characters + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( self, vocab: str, + confidence_aggregation: ConfidenceAggregation = "mean", ) -> None: - super().__init__(vocab) + super().__init__(vocab, confidence_aggregation) self._embedding = list(vocab) + ["", "", ""] diff --git a/doctr/models/recognition/parseq/pytorch.py b/doctr/models/recognition/parseq/pytorch.py index aec478b78f..bf1763f0c7 100644 --- a/doctr/models/recognition/parseq/pytorch.py +++ b/doctr/models/recognition/parseq/pytorch.py @@ -18,6 +18,7 @@ from doctr.datasets import VOCABS from doctr.models.modules.transformer import MultiHeadAttention, PositionwiseFeedForward +from ..._utils import ConfidenceAggregation, aggregate_confidence from ...classification import vit_s from ...utils import _bf16_to_float32, load_pretrained_params from .base import _PARSeq, _PARSeqPostProcessor @@ -120,6 +121,7 @@ class PARSeq(_PARSeq, nn.Module): input_shape: input shape of the image exportable: onnx exportable returns only logits cfg: dictionary containing information about the model + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( @@ -135,6 +137,7 @@ def __init__( input_shape: tuple[int, int, int] = (3, 32, 128), exportable: bool = False, cfg: dict[str, Any] | None = None, + confidence_aggregation: ConfidenceAggregation = "mean", ) -> None: super().__init__() self.vocab = vocab @@ -152,7 +155,7 @@ def __init__( self.pos_queries = nn.Parameter(torch.Tensor(1, self.max_length + 1, embedding_units)) # +1 for EOS self.dropout = nn.Dropout(p=dropout_prob) - self.postprocessor = PARSeqPostProcessor(vocab=self.vocab) + self.postprocessor = PARSeqPostProcessor(vocab=self.vocab, confidence_aggregation=confidence_aggregation) nn.init.trunc_normal_(self.pos_queries, std=0.02) for n, m in self.named_modules(): @@ -424,9 +427,10 @@ def __call__( "".join(self._embedding[idx] for idx in encoded_seq).split("")[0] for encoded_seq in out_idxs.detach().cpu().numpy() ] - # compute probabilties for each word up to the EOS token + # aggregate the character probabilities of each word up to the EOS token probs = [ - float(preds_prob[i, : len(word)].clip(0, 1).mean()) if word else 0.0 for i, word in enumerate(word_values) + aggregate_confidence(preds_prob[i, : len(word)], self.confidence_aggregation) + for i, word in enumerate(word_values) ] return list(zip(word_values, probs)) diff --git a/doctr/models/recognition/predictor/_utils.py b/doctr/models/recognition/predictor/_utils.py index ff53178d69..b14b12e6eb 100644 --- a/doctr/models/recognition/predictor/_utils.py +++ b/doctr/models/recognition/predictor/_utils.py @@ -8,6 +8,7 @@ import numpy as np +from ..._utils import ConfidenceAggregation, aggregate_confidence from ..utils import merge_multi_strings __all__ = ["split_crops", "remap_preds"] @@ -110,6 +111,7 @@ def remap_preds( preds: list[tuple[str, float]], crop_map: list[int | tuple[int, int, float]], overlap_ratio: float, + confidence_aggregation: ConfidenceAggregation = "min", ) -> list[tuple[str, float]]: """ Reconstruct predictions from possibly split crops. @@ -118,6 +120,7 @@ def remap_preds( preds: List of (text, confidence) tuples from each crop. crop_map: Map returned by `split_crops`. overlap_ratio: Overlap ratio used during splitting. + confidence_aggregation: Aggregation method of the confidences of the non-empty parts. Returns: List of merged (text, confidence) tuples corresponding to original crops. @@ -130,6 +133,7 @@ def remap_preds( start_idx, end_idx, last_overlap = item text_parts, confidences = zip(*preds[start_idx:end_idx]) merged_text = merge_multi_strings(list(text_parts), overlap_ratio, last_overlap) - merged_conf = min((conf for text, conf in zip(text_parts, confidences) if text), default=min(confidences)) + text_confidences = [conf for text, conf in zip(text_parts, confidences) if text] + merged_conf = aggregate_confidence(text_confidences or confidences, confidence_aggregation) remapped.append((merged_text, merged_conf)) return remapped diff --git a/doctr/models/recognition/predictor/pytorch.py b/doctr/models/recognition/predictor/pytorch.py index f15163e626..8385c07409 100644 --- a/doctr/models/recognition/predictor/pytorch.py +++ b/doctr/models/recognition/predictor/pytorch.py @@ -10,6 +10,7 @@ import torch from torch import nn +from doctr.models._utils import ConfidenceAggregation, _resolve_confidence_aggregation from doctr.models.preprocessor import PreProcessor from doctr.models.utils import set_device_and_dtype @@ -40,6 +41,7 @@ def __init__( self.critical_ar = 8 # Critical aspect ratio self.overlap_ratio = 0.5 # Ratio of overlap between neighboring crops self.target_ar = 6 # Target aspect ratio + self.split_confidence_aggregation: ConfidenceAggregation = "min" # Confidence aggregation for split crops @torch.inference_mode() def forward( @@ -49,6 +51,7 @@ def forward( ) -> list[tuple[str, float]]: if len(crops) == 0: return [] + _resolve_confidence_aggregation(self.split_confidence_aggregation) # Dimension check if any(crop.ndim != 3 for crop in crops): raise ValueError("incorrect input shape: all crops are expected to be multi-channel 2D images.") @@ -80,6 +83,6 @@ def forward( # Remap crops if self.split_wide_crops and remapped: - out = remap_preds(out, crop_map, self.overlap_ratio) + out = remap_preds(out, crop_map, self.overlap_ratio, self.split_confidence_aggregation) return out diff --git a/doctr/models/recognition/sar/pytorch.py b/doctr/models/recognition/sar/pytorch.py index 4b4bbf7d97..7f8fb03299 100644 --- a/doctr/models/recognition/sar/pytorch.py +++ b/doctr/models/recognition/sar/pytorch.py @@ -14,6 +14,7 @@ from doctr.datasets import VOCABS +from ..._utils import ConfidenceAggregation, aggregate_confidence from ...classification import resnet31 from ...utils import _bf16_to_float32, load_pretrained_params from ..core import RecognitionModel, RecognitionPostProcessor @@ -175,6 +176,7 @@ class SAR(nn.Module, RecognitionModel): dropout_prob: dropout probability of the encoder LSTM exportable: onnx exportable returns only logits cfg: dictionary containing information about the model + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( @@ -189,6 +191,7 @@ def __init__( input_shape: tuple[int, int, int] = (3, 32, 128), exportable: bool = False, cfg: dict[str, Any] | None = None, + confidence_aggregation: ConfidenceAggregation = "min", ) -> None: super().__init__() self.vocab = vocab @@ -216,7 +219,7 @@ def __init__( dropout_prob=dropout_prob, ) - self.postprocessor = SARPostProcessor(vocab=vocab) + self.postprocessor = SARPostProcessor(vocab=vocab, confidence_aggregation=confidence_aggregation) for n, m in self.named_modules(): # Don't override the initialization of the backbone @@ -320,8 +323,16 @@ class SARPostProcessor(RecognitionPostProcessor): Args: vocab: string containing the ordered sequence of supported characters + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ + def __init__( + self, + vocab: str, + confidence_aggregation: ConfidenceAggregation = "min", + ) -> None: + super().__init__(vocab, confidence_aggregation) + def __call__( self, logits: torch.Tensor, @@ -329,17 +340,20 @@ def __call__( # compute pred with argmax for attention models out_idxs = logits.argmax(-1) # N x L - probs = torch.gather(torch.softmax(logits, -1), -1, out_idxs.unsqueeze(-1)).squeeze(-1) - # Take the minimum confidence of the sequence - probs = probs.min(dim=1).values.detach().cpu() + preds_prob = torch.softmax(logits, -1).max(dim=-1)[0].detach().cpu().numpy() # Manual decoding word_values = [ "".join(self._embedding[idx] for idx in encoded_seq).split("")[0] for encoded_seq in out_idxs.detach().cpu().numpy() ] + # aggregate the character probabilities of each word up to the EOS token + probs = [ + aggregate_confidence(preds_prob[i, : len(word)], self.confidence_aggregation) + for i, word in enumerate(word_values) + ] - return list(zip(word_values, probs.numpy().clip(0, 1).tolist())) + return list(zip(word_values, probs)) def _sar( diff --git a/doctr/models/recognition/viptr/pytorch.py b/doctr/models/recognition/viptr/pytorch.py index 5f0d6da38f..ba9a56fa3e 100644 --- a/doctr/models/recognition/viptr/pytorch.py +++ b/doctr/models/recognition/viptr/pytorch.py @@ -5,9 +5,9 @@ from collections.abc import Callable from copy import deepcopy -from itertools import groupby from typing import Any +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F @@ -15,6 +15,7 @@ from doctr.datasets import VOCABS, decode_sequence +from ..._utils import ConfidenceAggregation, aggregate_confidence from ...classification import vip_tiny from ...utils import _bf16_to_float32, load_pretrained_params from ..core import RecognitionModel, RecognitionPostProcessor @@ -38,13 +39,22 @@ class VIPTRPostProcessor(RecognitionPostProcessor): Args: vocab: string containing the ordered sequence of supported characters + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ + def __init__( + self, + vocab: str, + confidence_aggregation: ConfidenceAggregation = "min", + ) -> None: + super().__init__(vocab, confidence_aggregation) + @staticmethod def ctc_best_path( logits: torch.Tensor, vocab: str = VOCABS["french"], blank: int = 0, + confidence_aggregation: ConfidenceAggregation = "min", ) -> list[tuple[str, float]]: """Implements best path decoding as shown by Graves (Dissertation, p63), highly inspired from `_. @@ -53,18 +63,27 @@ def ctc_best_path( logits: model output, shape: N x T x C vocab: vocabulary to use blank: index of blank label + confidence_aggregation: aggregation method of the character probabilities into the word confidence Returns: A list of tuples: (word, confidence) """ - # Gather the most confident characters, and assign the smallest conf among those to the sequence prob - probs = F.softmax(logits, dim=-1).max(dim=-1).values.min(dim=1).values - - # collapse best path (using itertools.groupby), map to chars, join char list to string - best_paths = torch.argmax(logits, dim=-1).detach().cpu().numpy() - words = [decode_sequence([k for k, _ in groupby(seq.tolist()) if k != blank], vocab) for seq in best_paths] - - return list(zip(words, probs.tolist())) + best_paths = logits.argmax(dim=-1).detach().cpu().numpy() + probs = F.softmax(logits.float(), dim=-1).max(dim=-1).values.detach().cpu().numpy() + + results = [] + for path, path_probs in zip(best_paths, probs): + # Collapse the repeated labels: the probability of a character is the highest one within its run + run_starts = np.flatnonzero(np.r_[True, path[1:] != path[:-1]]) + labels, label_probs = path[run_starts], np.maximum.reduceat(path_probs, run_starts) + # Remove the blanks + is_char = labels != blank + results.append(( + decode_sequence(labels[is_char].tolist(), vocab), + aggregate_confidence(label_probs[is_char], confidence_aggregation), + )) + + return results def __call__(self, logits: torch.Tensor) -> list[tuple[str, float]]: """Performs decoding of raw output with CTC and decoding of CTC predictions @@ -78,7 +97,12 @@ def __call__(self, logits: torch.Tensor) -> list[tuple[str, float]]: """ # Decode CTC - return self.ctc_best_path(logits=logits, vocab=self.vocab, blank=len(self.vocab)) + return self.ctc_best_path( + logits=logits, + vocab=self.vocab, + blank=len(self.vocab), + confidence_aggregation=self.confidence_aggregation, + ) class VIPTR(RecognitionModel, nn.Module): @@ -91,6 +115,7 @@ class VIPTR(RecognitionModel, nn.Module): input_shape: input shape of the image exportable: onnx exportable returns only logits cfg: configuration dictionary + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( @@ -100,6 +125,7 @@ def __init__( input_shape: tuple[int, int, int] = (3, 32, 128), exportable: bool = False, cfg: dict[str, Any] | None = None, + confidence_aggregation: ConfidenceAggregation = "min", ): super().__init__() self.vocab = vocab @@ -112,7 +138,7 @@ def __init__( with torch.inference_mode(): embedding_units = self.feat_extractor(torch.zeros((1, *input_shape)))["features"].shape[-1] - self.postprocessor = VIPTRPostProcessor(vocab=self.vocab) + self.postprocessor = VIPTRPostProcessor(vocab=self.vocab, confidence_aggregation=confidence_aggregation) self.head = nn.Linear(embedding_units, len(self.vocab) + 1) # +1 for PAD for n, m in self.named_modules(): diff --git a/doctr/models/recognition/vitstr/base.py b/doctr/models/recognition/vitstr/base.py index 7264f40e5c..64c76ef65a 100644 --- a/doctr/models/recognition/vitstr/base.py +++ b/doctr/models/recognition/vitstr/base.py @@ -7,6 +7,7 @@ import numpy as np from ....datasets import encode_sequences +from ..._utils import ConfidenceAggregation from ..core import RecognitionPostProcessor @@ -43,11 +44,13 @@ class _ViTSTRPostProcessor(RecognitionPostProcessor): Args: vocab: string containing the ordered sequence of supported characters + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( self, vocab: str, + confidence_aggregation: ConfidenceAggregation = "mean", ) -> None: - super().__init__(vocab) + super().__init__(vocab, confidence_aggregation) self._embedding = list(vocab) + ["", ""] diff --git a/doctr/models/recognition/vitstr/pytorch.py b/doctr/models/recognition/vitstr/pytorch.py index 6301ef2d6f..e6a281d2f7 100644 --- a/doctr/models/recognition/vitstr/pytorch.py +++ b/doctr/models/recognition/vitstr/pytorch.py @@ -14,6 +14,7 @@ from doctr.datasets import VOCABS +from ..._utils import ConfidenceAggregation, aggregate_confidence from ...classification import vit_b, vit_s from ...utils import _bf16_to_float32, load_pretrained_params from .base import _ViTSTR, _ViTSTRPostProcessor @@ -51,6 +52,7 @@ class ViTSTR(_ViTSTR, nn.Module): input_shape: input shape of the image exportable: onnx exportable returns only logits cfg: dictionary containing information about the model + confidence_aggregation: aggregation method of the character probabilities into the word confidence """ def __init__( @@ -62,6 +64,7 @@ def __init__( input_shape: tuple[int, int, int] = (3, 32, 128), # different from paper exportable: bool = False, cfg: dict[str, Any] | None = None, + confidence_aggregation: ConfidenceAggregation = "mean", ) -> None: super().__init__() self.vocab = vocab @@ -72,7 +75,7 @@ def __init__( self.feat_extractor = feature_extractor self.head = nn.Linear(embedding_units, len(self.vocab) + 1) # +1 for EOS - self.postprocessor = ViTSTRPostProcessor(vocab=self.vocab) + self.postprocessor = ViTSTRPostProcessor(vocab=self.vocab, confidence_aggregation=confidence_aggregation) def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None: """Load pretrained parameters onto the model @@ -181,9 +184,10 @@ def __call__( "".join(self._embedding[idx] for idx in encoded_seq).split("")[0] for encoded_seq in out_idxs.detach().cpu().numpy() ] - # compute probabilties for each word up to the EOS token + # aggregate the character probabilities of each word up to the EOS token probs = [ - float(preds_prob[i, : len(word)].clip(0, 1).mean()) if word else 0.0 for i, word in enumerate(word_values) + aggregate_confidence(preds_prob[i, : len(word)], self.confidence_aggregation) + for i, word in enumerate(word_values) ] return list(zip(word_values, probs)) diff --git a/doctr/models/zoo.py b/doctr/models/zoo.py index 5b1d999186..52c7690b17 100644 --- a/doctr/models/zoo.py +++ b/doctr/models/zoo.py @@ -5,6 +5,7 @@ from typing import Any +from ._utils import ConfidenceAggregation, _resolve_confidence_aggregation from .detection.zoo import detection_predictor from .kie_predictor import KIEPredictor from .layout.zoo import layout_predictor @@ -32,8 +33,13 @@ def _predictor( layout_arch: Any = "lw_detr_s", ignore_regions: list[str] | None = None, detect_tables: bool = False, + confidence_aggregation: ConfidenceAggregation | None = None, **kwargs, ) -> OCRPredictor: + if confidence_aggregation is not None: + # Fail before building the models + _resolve_confidence_aggregation(confidence_aggregation) + # Detection det_predictor = detection_predictor( det_arch, @@ -51,6 +57,7 @@ def _predictor( pretrained=pretrained, pretrained_backbone=pretrained_backbone, batch_size=reco_bs, + confidence_aggregation=confidence_aggregation, ) # Layout - required for table detection, so build it whenever layout or tables are requested @@ -113,6 +120,7 @@ def ocr_predictor( layout_arch: Any = "lw_detr_s", ignore_regions: list[str] | None = None, detect_tables: bool = False, + confidence_aggregation: ConfidenceAggregation | None = None, **kwargs: Any, ) -> OCRPredictor: """End-to-end OCR architecture using one model for localization, and another for text recognition. @@ -156,6 +164,11 @@ def ocr_predictor( structure model. Words falling inside a detected table are regrouped into a structured table (accessible via `page.tables`) and removed from the regular text output. This enables the layout model and slightly deteriorates the overall latency. + confidence_aggregation: aggregation method of the character probabilities of the recognition model into the + word confidence: "mean", "min", "max", "median", "geometric_mean", "harmonic_mean" or a Python callable + taking the 1D array of probabilities and returning a float. If None (default), the method of the + recognition model is kept ("min" for CRNN, VIPTR, SAR and MASTER, "mean" for ViTSTR and PARSeq). + A recognition model instance passed as `reco_arch` is modified. kwargs: keyword args of `OCRPredictor` Returns: @@ -177,6 +190,7 @@ def ocr_predictor( layout_arch=layout_arch, detect_tables=detect_tables, ignore_regions=ignore_regions, + confidence_aggregation=confidence_aggregation, **kwargs, ) @@ -197,8 +211,13 @@ def _kie_predictor( detect_layout: bool = False, layout_arch: Any = "lw_detr_s", ignore_regions: list[str] | None = None, + confidence_aggregation: ConfidenceAggregation | None = None, **kwargs, ) -> KIEPredictor: + if confidence_aggregation is not None: + # Fail before building the models + _resolve_confidence_aggregation(confidence_aggregation) + # Detection det_predictor = detection_predictor( det_arch, @@ -216,6 +235,7 @@ def _kie_predictor( pretrained=pretrained, pretrained_backbone=pretrained_backbone, batch_size=reco_bs, + confidence_aggregation=confidence_aggregation, ) # Layout - optional @@ -262,6 +282,7 @@ def kie_predictor( detect_layout: bool = False, layout_arch: Any = "lw_detr_s", ignore_regions: list[str] | None = None, + confidence_aggregation: ConfidenceAggregation | None = None, **kwargs: Any, ) -> KIEPredictor: """End-to-end KIE architecture using one model for localization, and another for text recognition. @@ -301,6 +322,11 @@ def kie_predictor( ignore_regions: optional list of layout class names to ignore during detection/recognition. If provided, the layout model will be used to locate the regions of the specified classes, and these regions will be masked out (filled with black) before passing the pages to the detection/recognition modules. + confidence_aggregation: aggregation method of the character probabilities of the recognition model into the + word confidence: "mean", "min", "max", "median", "geometric_mean", "harmonic_mean" or a Python callable + taking the 1D array of probabilities and returning a float. If None (default), the method of the + recognition model is kept ("min" for CRNN, VIPTR, SAR and MASTER, "mean" for ViTSTR and PARSeq). + A recognition model instance passed as `reco_arch` is modified. kwargs: keyword args of `OCRPredictor` Returns: @@ -321,5 +347,6 @@ def kie_predictor( detect_layout=detect_layout, layout_arch=layout_arch, ignore_regions=ignore_regions, + confidence_aggregation=confidence_aggregation, **kwargs, ) diff --git a/pyproject.toml b/pyproject.toml index d01a031d73..b50f62111c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ dependencies = [ "shapely>=1.6.0,<3.0.0", "langdetect>=1.0.9,<2.0.0", "rapidfuzz>=3.0.0,<4.0.0", - "huggingface-hub>=0.20.0,<2.0.0", + "huggingface-hub>=1.0.0,<3.0.0", "Pillow>=10.1.0", "anyascii>=0.3.2", "validators>=0.18.0", diff --git a/tests/common/test_cli.py b/tests/common/test_cli.py index 1ed0d77074..9ace213579 100644 --- a/tests/common/test_cli.py +++ b/tests/common/test_cli.py @@ -88,6 +88,7 @@ def test_parse_args_defaults(): assert args.device == "auto" assert args.bin_thresh is None assert args.box_thresh is None + assert args.confidence_aggregation is None assert args.assume_straight_pages is True assert args.straighten_pages is False assert args.preserve_aspect_ratio is True @@ -192,6 +193,8 @@ def test_parse_args_advanced_options(): "--paragraph_break", "0.1", "--keep_reading_order", + "--confidence_aggregation", + "geometric_mean", "--direction", "rtl", "--indent", @@ -211,6 +214,7 @@ def test_parse_args_advanced_options(): assert args.resolve_blocks is True assert args.paragraph_break == 0.1 assert args.keep_reading_order is True + assert args.confidence_aggregation == "geometric_mean" assert args.direction == "rtl" assert args.indent == 2 assert args.quiet is True @@ -233,6 +237,7 @@ def test_parse_args_version(capsys): ["--bin_thresh", "1.5"], ["--box_thresh", "-0.1"], ["--bin_thresh", "high"], + ["--confidence_aggregation", "average"], ], ) def test_parse_args_invalid_values(argv): @@ -416,6 +421,8 @@ def test_build_predictor_forwards_every_option(recorded_predictor): "4", "--reco_bs", "64", + "--confidence_aggregation", + "geometric_mean", "--detect_orientation", "--detect_language", "--detect_layout", @@ -450,6 +457,7 @@ def test_build_predictor_forwards_every_option(recorded_predictor): "detect_tables": True, "det_bs": 4, "reco_bs": 64, + "confidence_aggregation": "geometric_mean", "disable_page_orientation": True, "disable_crop_orientation": True, "preserve_original_coords": True, @@ -465,6 +473,7 @@ def test_build_predictor_forwards_every_option(recorded_predictor): assert model.table_predictor is not None assert model.det_predictor.pre_processor.batch_size == 4 assert model.reco_predictor.pre_processor.batch_size == 64 + assert model.reco_predictor.model.postprocessor.confidence_aggregation == "geometric_mean" builder = model.doc_builder assert (builder.resolve_lines, builder.resolve_blocks, builder.paragraph_break) == (False, True, 0.1) diff --git a/tests/common/test_models.py b/tests/common/test_models.py index 9cde5b69bf..c35b7b3a73 100644 --- a/tests/common/test_models.py +++ b/tests/common/test_models.py @@ -6,7 +6,13 @@ import requests from doctr.io import reader -from doctr.models._utils import estimate_orientation, get_language, invert_data_structure, mask_boxes +from doctr.models._utils import ( + aggregate_confidence, + estimate_orientation, + get_language, + invert_data_structure, + mask_boxes, +) from doctr.utils import geometry @@ -142,3 +148,40 @@ def test_convert_list_dict(): assert converted_dic == tar_dict assert converted_list == dic + + +@pytest.mark.parametrize( + "method, expected", + [ + ("mean", 0.77), + ("min", 0.12), + ("max", 0.99), + ("median", 0.985), + ("geometric_mean", 0.5827), + ("harmonic_mean", 0.3517), + (lambda scores: float(np.percentile(scores, 25)), 0.765), + ], +) +def test_aggregate_confidence(method, expected): + # One uncertain character among confident ones + scores = np.array([0.99, 0.98, 0.12, 0.99], dtype=np.float32) + assert aggregate_confidence(scores, method) == pytest.approx(expected, abs=1e-4) + assert aggregate_confidence(scores.tolist(), method) == pytest.approx(expected, abs=1e-4) + assert aggregate_confidence(scores.reshape(2, 2), method) == pytest.approx(expected, abs=1e-4) + assert aggregate_confidence([], method) == 0.0 + + +def test_aggregate_confidence_edge_cases(): + # A zero score cancels the geometric and harmonic means + assert aggregate_confidence([0.0, 0.9], "geometric_mean") == 0.0 + assert aggregate_confidence([0.0, 0.9], "harmonic_mean") == 0.0 + # The result is clipped to [0, 1] + assert aggregate_confidence([0.5], lambda scores: -1.0) == 0.0 + assert aggregate_confidence([0.5], lambda scores: 2.0) == 1.0 + # A NaN aggregation (e.g. the mean of an empty selection) counts as no confidence, a 1-element array is accepted + assert aggregate_confidence([0.4, 0.5], lambda scores: np.mean(scores[scores > 0.6])) == 0.0 + assert aggregate_confidence([0.4, 0.5], lambda scores: np.quantile(scores, [0.5])) == pytest.approx(0.45) + # Invalid methods, including values which are not hashable + for invalid in ["average", ["mean"], {"mean"}, None, 1]: + with pytest.raises(ValueError, match="Unknown confidence aggregation"): + aggregate_confidence([0.5], invalid) diff --git a/tests/common/test_models_recognition_predictor.py b/tests/common/test_models_recognition_predictor.py index 00aadb3cf8..9735082017 100644 --- a/tests/common/test_models_recognition_predictor.py +++ b/tests/common/test_models_recognition_predictor.py @@ -125,3 +125,13 @@ def test_invalid_split_overlap_ratio(split_overlap_ratio): target_ratio=4, split_overlap_ratio=split_overlap_ratio, ) + + +@pytest.mark.parametrize( + "confidence_aggregation, conf", + [("min", 0.1), ("mean", 0.5), ("geometric_mean", 0.3), (lambda confs: 0.42, 0.42)], +) +def test_remap_preds_confidence_aggregation(confidence_aggregation, conf): + # The empty part is ignored for the confidence + preds = [("hellowo", 0.9), ("", 0.0), ("loworld", 0.1)] + assert remap_preds(preds, [(0, 3, 0.5)], 0.5, confidence_aggregation) == [("helloworld", pytest.approx(conf))] diff --git a/tests/pytorch/test_models_recognition_pt.py b/tests/pytorch/test_models_recognition_pt.py index 06fafa0041..f7d4db023c 100644 --- a/tests/pytorch/test_models_recognition_pt.py +++ b/tests/pytorch/test_models_recognition_pt.py @@ -9,10 +9,12 @@ from doctr.io import DocumentFile from doctr.models import recognition +from doctr.models.preprocessor import PreProcessor from doctr.models.recognition.crnn.pytorch import CTCPostProcessor from doctr.models.recognition.master.pytorch import MASTERPostProcessor from doctr.models.recognition.parseq.pytorch import PARSeqPostProcessor from doctr.models.recognition.predictor import RecognitionPredictor +from doctr.models.recognition.predictor._utils import split_crops from doctr.models.recognition.sar.pytorch import SARPostProcessor from doctr.models.recognition.viptr.pytorch import VIPTRPostProcessor from doctr.models.recognition.vitstr.pytorch import ViTSTRPostProcessor @@ -84,7 +86,122 @@ def test_reco_postprocessors(post_processor, input_shape, mock_vocab): assert len(decoded) == input_shape[0] assert all(char in mock_vocab for word, _ in decoded for char in word) # Repr - assert repr(processor) == f"{post_processor.__name__}(vocab_size={len(mock_vocab)})" + default = "mean" if post_processor in (ViTSTRPostProcessor, PARSeqPostProcessor) else "min" + assert repr(processor) == ( + f"{post_processor.__name__}(vocab_size={len(mock_vocab)}, confidence_aggregation={default!r})" + ) + assert "confidence_aggregation=" in repr(post_processor(mock_vocab, confidence_aggregation=lambda p: 1.0)) + + +def _logits(probs: list[list[float]], num_classes: int) -> torch.Tensor: + probs_ = torch.zeros((1, len(probs), num_classes)) + probs_[0, :, : len(probs[0])] = torch.tensor(probs) + return probs_.clamp_min(1e-9).log() + + +@pytest.mark.parametrize( + "confidence_aggregation, ctc_conf, attention_conf", + [ + ("mean", 0.75, 0.7), + ("min", 0.6, 0.5), + ("geometric_mean", 0.54**0.5, 0.45**0.5), + (lambda probs: 1.0, 1.0, 1.0), + ], +) +@pytest.mark.parametrize( + "post_processor, num_classes", + [ + [CTCPostProcessor, 4], + [VIPTRPostProcessor, 4], + [SARPostProcessor, 4], + [ViTSTRPostProcessor, 5], + [MASTERPostProcessor, 6], + [PARSeqPostProcessor, 6], + ], +) +def test_reco_postprocessors_confidence_aggregation( + post_processor, num_classes, confidence_aggregation, ctc_conf, attention_conf +): + processor = post_processor("abc", confidence_aggregation=confidence_aggregation) + if post_processor in (CTCPostProcessor, VIPTRPostProcessor): + # "a a b b ": a character probability is the highest one within its run, blanks are ignored + probs = [[0.5, 0.2, 0.1, 0.2], [0.9, 0.05, 0.0, 0.05], [0.1, 0.1, 0.1, 0.7], [0.1, 0.4, 0.2, 0.3]] + probs += [[0.2, 0.6, 0.1, 0.1], [0.0, 0.0, 0.0, 1.0]] + word, conf = processor(_logits(probs, num_classes))[0] + assert (word, conf) == ("ab", pytest.approx(ctc_conf, abs=1e-5)) + else: + # "a b a": the probabilities after the token are ignored + probs = [[0.9, 0.05, 0.03, 0.02], [0.1, 0.5, 0.2, 0.2], [0.1, 0.1, 0.1, 0.7], [0.4, 0.2, 0.2, 0.2]] + word, conf = processor(_logits(probs, num_classes))[0] + assert (word, conf) == ("ab", pytest.approx(attention_conf, abs=1e-5)) + # Empty word + assert processor(_logits([[0.1, 0.1, 0.1, 0.7]] * 3, num_classes)) == [("", 0.0)] + for invalid in ["average", ["mean"], None]: + with pytest.raises(ValueError, match="Unknown confidence aggregation"): + post_processor("abc", confidence_aggregation=invalid) + + +class _MockRecoModel(torch.nn.Module): + """Recognition model returning a predefined confidence for each crop it receives""" + + def __init__(self, confidences: list[float], postprocessor=None) -> None: + super().__init__() + self.dummy = torch.nn.Parameter(torch.zeros(1)) + self.confidences = confidences + if postprocessor is not None: + self.postprocessor = postprocessor + + def forward(self, x: torch.Tensor, return_preds: bool = False, **kwargs): + return {"preds": [("ab", conf) for conf in self.confidences[: x.shape[0]]]} + + +@pytest.mark.parametrize( + "postprocessor", + [ + # The aggregation of the split parts is independent from the one of the character probabilities + CTCPostProcessor("abc", confidence_aggregation="max"), + PARSeqPostProcessor("abc"), + # A custom model does not have to provide a postprocessor with a confidence aggregation + None, + ], +) +def test_recognition_predictor_split_confidence_aggregation(postprocessor): + confidences = [0.9, 0.2, 0.7, 0.8, 0.6, 0.5] + predictor = RecognitionPredictor( + PreProcessor(output_size=(32, 128), batch_size=32, preserve_aspect_ratio=True), + _MockRecoModel(confidences, postprocessor), + ) + # A wide crop split into several parts + wide_crop = np.zeros((32, 32 * 20, 3), dtype=np.uint8) + num_parts = len(split_crops([wide_crop], predictor.critical_ar, predictor.target_ar, predictor.overlap_ratio)[0]) + assert 2 <= num_parts <= len(confidences) + confidences = confidences[:num_parts] + + # The lowest confidence of the parts by default + assert predictor.split_confidence_aggregation == "min" + assert predictor([wide_crop])[0][1] == pytest.approx(min(confidences)) + predictor.split_confidence_aggregation = "mean" + assert predictor([wide_crop])[0][1] == pytest.approx(np.mean(confidences)) + predictor.split_confidence_aggregation = lambda confs: float(confs.max()) + assert predictor([wide_crop])[0][1] == pytest.approx(max(confidences)) + # A crop which is not split keeps the confidence of the model + assert predictor([np.zeros((32, 128, 3), dtype=np.uint8)]) == [("ab", confidences[0])] + # An invalid method is rejected at the first call, even without a crop to split + predictor = RecognitionPredictor( + PreProcessor(output_size=(32, 128), batch_size=32, preserve_aspect_ratio=True), + _MockRecoModel(confidences, postprocessor), + ) + predictor.split_confidence_aggregation = "average" + with pytest.raises(ValueError, match="Unknown confidence aggregation"): + predictor([np.zeros((32, 128, 3), dtype=np.uint8)]) + + +@pytest.mark.parametrize( + "arch_name", ["crnn_mobilenet_v3_small", "sar_resnet31", "master", "vitstr_small", "parseq", "viptr_tiny"] +) +def test_recognition_models_confidence_aggregation(arch_name, mock_vocab): + model = recognition.__dict__[arch_name](vocab=mock_vocab, pretrained_backbone=False, confidence_aggregation="max") + assert model.postprocessor.confidence_aggregation == "max" @pytest.mark.parametrize( diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index 2654d69736..bd1a595fa5 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -301,6 +301,75 @@ def test_ocrpredictor_tables_factory(): assert predictor.table_predictor is None +def _custom_aggregation(scores): + return float(np.percentile(scores, 25)) + + +@pytest.mark.parametrize( + "confidence_aggregation, expected", + [ + # Default method of the recognition model + [None, "min"], + ["mean", "mean"], + ["geometric_mean", "geometric_mean"], + [_custom_aggregation, _custom_aggregation], + ], +) +def test_ocr_predictor_confidence_aggregation(confidence_aggregation, expected): + for factory in (models.ocr_predictor, models.kie_predictor): + predictor = factory( + "db_mobilenet_v3_large", + "crnn_mobilenet_v3_small", + pretrained=False, + pretrained_backbone=False, + confidence_aggregation=confidence_aggregation, + ) + assert predictor.reco_predictor.model.postprocessor.confidence_aggregation == expected + # The split crops keep their own aggregation + assert predictor.reco_predictor.split_confidence_aggregation == "min" + + +def test_ocr_predictor_confidence_aggregation_model_instances(mock_payslip): + reco_model = recognition.parseq(pretrained=True) + predictor = models.ocr_predictor( + detection.db_mobilenet_v3_large(pretrained=True), reco_model, confidence_aggregation="min" + ) + assert reco_model.postprocessor.confidence_aggregation == "min" + # None keeps the method of the model + assert models.ocr_predictor("db_mobilenet_v3_large", reco_model).reco_predictor.model is reco_model + assert reco_model.postprocessor.confidence_aggregation == "min" + + # End-to-end: the word confidences follow the aggregation method set on the recognition model + doc = DocumentFile.from_images(mock_payslip) + + def word_confidences(method): + reco_model.postprocessor.confidence_aggregation = method + out = predictor(doc) + return np.array([w.confidence for b in out.pages[0].blocks for line in b.lines for w in line.words]) + + min_confs, max_confs = word_confidences("min"), word_confidences("max") + mean_confs = word_confidences("mean") + assert min_confs.size > 0 + assert np.all((min_confs >= 0) & (max_confs <= 1)) + assert np.all(min_confs <= mean_confs + 1e-6) and np.all(mean_confs <= max_confs + 1e-6) + # The method is applied: at least one word has an uncertain character + assert np.any(min_confs < max_confs - 1e-3) + + +def test_recognition_predictor_confidence_aggregation(): + # Architecture names + reco_predictor = recognition_predictor("parseq", pretrained=False, confidence_aggregation="min") + assert reco_predictor.model.postprocessor.confidence_aggregation == "min" + # Model instances are modified + reco_model = recognition.crnn_mobilenet_v3_small(pretrained=False, pretrained_backbone=False) + assert recognition_predictor(reco_model, confidence_aggregation=np.median).model is reco_model + assert reco_model.postprocessor.confidence_aggregation is np.median + # None keeps the method of the model + assert recognition_predictor(reco_model).model.postprocessor.confidence_aggregation is np.median + with pytest.raises(ValueError, match="Unknown confidence aggregation"): + recognition_predictor("crnn_mobilenet_v3_small", pretrained=False, confidence_aggregation="average") + + @pytest.mark.parametrize( "detect_layout, detect_tables", [ From 9e1d642cea86f09f47fce8c85c0d8d37c64a47f7 Mon Sep 17 00:00:00 2001 From: felix Date: Fri, 9 Oct 2026 11:05:20 +0200 Subject: [PATCH 3/4] Add missing zoo modification --- doctr/models/recognition/zoo.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/doctr/models/recognition/zoo.py b/doctr/models/recognition/zoo.py index a28234c949..050bceb826 100644 --- a/doctr/models/recognition/zoo.py +++ b/doctr/models/recognition/zoo.py @@ -9,6 +9,7 @@ from doctr.models.utils import _CompiledModule from .. import recognition +from .._utils import ConfidenceAggregation, _resolve_confidence_aggregation from .predictor import RecognitionPredictor __all__ = ["recognition_predictor"] @@ -27,7 +28,15 @@ ] -def _predictor(arch: Any, pretrained: bool, **kwargs: Any) -> RecognitionPredictor: +def _predictor( + arch: Any, + pretrained: bool, + confidence_aggregation: ConfidenceAggregation | None = None, + **kwargs: Any, +) -> RecognitionPredictor: + if confidence_aggregation is not None: + _resolve_confidence_aggregation(confidence_aggregation) + if isinstance(arch, str): if arch not in ARCHS: raise ValueError(f"unknown architecture '{arch}'") @@ -51,6 +60,9 @@ def _predictor(arch: Any, pretrained: bool, **kwargs: Any) -> RecognitionPredict raise ValueError(f"unknown architecture: {type(arch)}") _model = arch + if confidence_aggregation is not None: + _model.postprocessor.confidence_aggregation = confidence_aggregation + kwargs.pop("pretrained_backbone", None) kwargs["mean"] = kwargs.get("mean", _model.cfg["mean"]) @@ -67,6 +79,7 @@ def recognition_predictor( pretrained: bool = False, symmetric_pad: bool = False, batch_size: int = 128, + confidence_aggregation: ConfidenceAggregation | None = None, **kwargs: Any, ) -> RecognitionPredictor: """Text recognition architecture. @@ -82,9 +95,19 @@ def recognition_predictor( pretrained: If True, returns a model pre-trained on our text recognition dataset symmetric_pad: if True, pad the image symmetrically instead of padding at the bottom-right batch_size: number of samples the model processes in parallel + confidence_aggregation: aggregation method of the character probabilities into the word confidence: + "mean", "min", "max", "median", "geometric_mean", "harmonic_mean" or a callable. If None, the one of + the model is kept ("min" for CRNN, VIPTR, SAR and MASTER, "mean" for ViTSTR and PARSeq by default) **kwargs: optional parameters to be passed to the architecture Returns: Recognition predictor """ - return _predictor(arch=arch, pretrained=pretrained, symmetric_pad=symmetric_pad, batch_size=batch_size, **kwargs) + return _predictor( + arch=arch, + pretrained=pretrained, + symmetric_pad=symmetric_pad, + batch_size=batch_size, + confidence_aggregation=confidence_aggregation, + **kwargs, + ) From e67d3ae512bf43c83cbe04d9caa04e8768e8d52d Mon Sep 17 00:00:00 2001 From: felix Date: Fri, 9 Oct 2026 11:11:12 +0200 Subject: [PATCH 4/4] Update --- doctr/models/_utils.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/doctr/models/_utils.py b/doctr/models/_utils.py index a0e72159f8..aa4961eaac 100644 --- a/doctr/models/_utils.py +++ b/doctr/models/_utils.py @@ -11,7 +11,6 @@ import cv2 import numpy as np -import torch from langdetect import LangDetectException, detect_langs from doctr.utils.geometry import rotate_image @@ -56,11 +55,6 @@ def _harmonic_mean(scores: np.ndarray) -> float: def _resolve_confidence_aggregation(method: ConfidenceAggregation) -> Callable[[np.ndarray], Any]: """Return the function of a confidence aggregation method, raise a `ValueError` if it is unknown""" - if isinstance(method, torch.nn.Module): - raise ValueError( - f"Unsupported confidence aggregation {type(method).__name__}: " - f"expected a Python callable or one of {list(_CONFIDENCE_AGGREGATIONS)}" - ) if callable(method): return method # the type check comes first: an unhashable value (e.g. a list) can not be looked up