[models] Configurable word confidence aggregation for the recognition models

[felixdittrich92]

Summary

The character probabilities of a recognition model can now be aggregated into the word confidence in different ways:

• "mean", "min", "max", "median", "geometric_mean" or "harmonic_mean";
• a Python callable that takes the 1D array of probabilities and returns a float.

The method can be set in four places:

• on the model or its postprocessor (confidence_aggregation=...);
• through recognition_predictor, ocr_predictor and kie_predictor (confidence_aggregation=...);
• with the CLI (--confidence_aggregation);
• afterwards, with predictor.reco_predictor.model.postprocessor.confidence_aggregation = ....

The defaults keep the method each architecture used before: "min" for CRNN, VIPTR, SAR and MASTER, and "mean" for ViTSTR and PARSeq.

predictor = ocr_predictor(pretrained=True, confidence_aggregation="geometric_mean")

Behaviour change of the default confidences

Only the probabilities of the predicted characters are aggregated now:

• CTC (CRNN, VIPTR): blank frames are ignored, and each character gets the highest probability among its frames. Before, the "min" ran over every frame, blank frames included.
• SAR, MASTER: aggregation stops at the end-of-sequence token. Before, it also covered the EOS position and the garbage positions after it.
• Empty prediction: confidence 0.

Decoded words are unchanged. Word confidences only go up, except for empty predictions. ViTSTR and PARSeq are bit-identical to main.

Quality impact on real documents. Setup:

• 565 word crops from three document photos (a typed German specification sheet and an invoice, photographed straight and rotated);
• every pretrained recognition checkpoint;
• the same crops on main and on this branch;
• hand-written ground truth;
• paired bootstrap 95% confidence intervals.

AUROC measures how well the confidence separates correct words from misread ones (higher is better):

model	main	this PR (default)	Δ AUROC [95% CI]
crnn_vgg16_bn	0.727	0.841	+0.115 [+0.073, +0.158]
crnn_mobilenet_v3_small	0.836	0.882	+0.046 [+0.019, +0.074]
crnn_mobilenet_v3_large	0.800	0.863	+0.064 [+0.033, +0.094]
viptr_tiny	0.777	0.866	+0.088 [+0.058, +0.120]
sar_resnet31	0.602	0.852	+0.250 [+0.205, +0.293]
master	0.766	0.858	+0.092 [+0.059, +0.129]
vitstr_small / vitstr_base / parseq	unchanged	unchanged	0

Averaged over the 9 models, per method:

	min	mean	geometric_mean	harmonic_mean	median	max
AUROC (higher is better)	0.855	0.854	0.854	0.855	0.742	0.652
ECE (lower is better)	0.200	0.264	0.261	0.257	0.275	0.281

• min, mean, geometric_mean and harmonic_mean rank errors about equally well, and min is the best calibrated of them.
• max and median are poor signals, and the docs warn about them.

Alternatives that were evaluated and rejected:

alternative	result
Keep the blank frames for CTC	Significantly worse on all 4 CTC models (−0.03 to −0.10 AUROC)
Use min or mean within a CTC run instead of max	Worse (−0.02 to −0.05)
Include the EOS probability	Model-dependent: +0.04 PARSeq, −0.09 SAR

Thresholds tuned on the confidences of previous versions may need to be adjusted. The docs contain a note about it.

Changes

• doctr/models/_utils.py: aggregate_confidence and the method resolution.
  • Validation: unknown names, unhashable values and torch.nn.Module raise a ValueError.
  • The result is clipped to [0, 1]; NaN and empty input give 0.
• Recognition postprocessors (CRNN, VIPTR, SAR, MASTER, ViTSTR, PARSeq): use the method.
  • CTC best-path decoding is vectorised with numpy; decoded strings are identical to the former groupby implementation.
  • MASTER finds the EOS from the predicted token indices, because <sos> and <pad> decode as multi-character strings.
• Model constructors: confidence_aggregation is the last argument (after cfg), so positional calls of previous versions are unchanged.
• remap_preds / RecognitionPredictor.split_confidence_aggregation ("min" by default, as before): aggregates the confidences of the parts of a split wide crop.
• Entry points: recognition_predictor, ocr_predictor, kie_predictor and the CLI. A model instance passed to them is modified. Invalid methods are rejected before any model is built.
• Docs: using_models.rst (new section with the behaviour-change note) and using_cli.rst.

Tests

• aggregate_confidence: every method, the edge cases and the validation.
• Hand-computed CTC and attention postprocessor values: blanks ignored, EOS handling, MASTER <pad>, empty words.
• bf16 / fp16 dtypes, split-crop aggregation, the zoo / CLI routing, and an end-to-end check that the method changes the word confidences.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.55%. Comparing base (fac9585) to head (e67d3ae).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2162   +/-   ##
=======================================
  Coverage   97.54%   97.55%           
=======================================
  Files         170      170           
  Lines       10665    10737   +72     
=======================================
+ Hits        10403    10474   +71     
- Misses        262      263    +1     

Flag	Coverage Δ
unittests	97.55% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.