[models] Configurable word confidence aggregation for the recognition models

.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 ]

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*

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.

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")

doctr/models/_utils.py

@@ -4,9 +4,10 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> 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
@@ -15,6 +16,8 @@
 from doctr.utils.geometry import rotate_image
 
 __all__ = [
+    "ConfidenceAggregation",
+    "aggregate_confidence",
     "estimate_orientation",
     "get_language",
     "invert_data_structure",
@@ -23,6 +26,84 @@
     "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 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.

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) + ["<eos>"]
 
     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)}"
+        )

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
         <https://github.com/githubharald/CTCDecoder>`_.
@@ -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

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) + ["<eos>"] + ["<sos>"] + ["<pad>"]