fix(inference): harden infer_argmax and add equivalence tests

python/rapidocr/inference_engine/base.py

@@ -94,9 +94,24 @@ class InferSession(abc.ABC):
         pass
 
     def infer_argmax(self, input_content: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
-        """argmax/max over the last axis of the output (rec logits)."""
+        """Return ``(argmax, max)`` along the *last* axis of the model output.
+
+        Used by the text recognizer (CTC decode), whose logits have shape
+        ``(batch, time, vocab)``: reducing on-device before the CPU copy avoids a
+        large device-to-host transfer. Subclasses running on an accelerator (e.g.
+        torch/CUDA) should override this to perform the reduction on-device.
+
+        Raises:
+            ValueError: if the model output has fewer than 2 dimensions.
+        """
         preds = self(input_content)
-        return preds.argmax(axis=2), preds.max(axis=2)
+        if preds.ndim < 2:
+            raise ValueError(
+                f"infer_argmax expects an output with at least 2 dimensions, "
+                f"got {preds.ndim} (shape={preds.shape}). It is intended for rec "
+                "logits of shape (batch, time, vocab)."
+            )
+        return preds.argmax(axis=-1), preds.max(axis=-1)
 
     @staticmethod
     def _verify_model(model_path: Union[str, Path, None]):

python/rapidocr/inference_engine/pytorch/main.py

@@ -25,8 +25,8 @@ class TorchInferSession(InferSession):
         return self._run(img, reduce=False)
 
     def infer_argmax(self, img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
-        """For rec: argmax/max of the (B, T, vocab) logits on the device, so only
-        (B, T) is copied to the CPU."""
+        """For the recognizer: argmax/max of the rec logits over the last axis
+        (vocab) on the device, so only (B, T) is copied to the CPU."""
         return self._run(img, reduce=True)
 
     def _run(self, img: np.ndarray, reduce: bool):
@@ -35,7 +35,7 @@ class TorchInferSession(InferSession):
             inp = inp.to(self.device)
             outputs = self.predictor(inp)
             if reduce:
-                prob, idx = outputs.max(dim=2)
+                prob, idx = outputs.max(dim=-1)
                 outputs = (idx.cpu().numpy(), prob.cpu().numpy())
             else:
                 outputs = outputs.cpu().numpy()

python/tests/test_infer_argmax.py

@@ -0,0 +1,100 @@
+# -*- encoding: utf-8 -*-
+# Unit tests for the rec-logits reduction introduced in PR #757.
+# These tests deliberately avoid downloading models or requiring a GPU.
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+root_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(root_dir))
+
+from rapidocr.ch_ppocr_rec.utils import CTCLabelDecode
+from rapidocr.inference_engine.base import InferSession
+from rapidocr.utils.utils import import_package
+
+
+class _ArrayReturningSession(InferSession):
+    """Minimal InferSession whose ``__call__`` returns a fixed 3D array."""
+
+    def __init__(self):
+        # bypass the abstract base __init__
+        class _Cfg:
+            pass
+
+        super().__init__(_Cfg())
+        self._fake = None
+
+    def setup(self, arr: np.ndarray) -> None:
+        self._fake = arr
+
+    def __call__(self, input_content):
+        return self._fake
+
+    def have_key(self, key: str = "character") -> bool:
+        return False
+
+    def get_character_list(self, key: str = "character"):
+        return []
+
+
+def test_infer_argmax_reduces_over_last_axis():
+    arr = np.random.RandomState(0).rand(2, 5, 10).astype(np.float32)
+    sess = _ArrayReturningSession()
+    sess.setup(arr)
+
+    idx, prob = sess.infer_argmax(np.zeros((1, 1), dtype=np.float32))
+
+    assert np.array_equal(idx, arr.argmax(axis=-1))
+    assert np.array_equal(prob, arr.max(axis=-1))
+
+
+def test_infer_argmax_rejects_under_two_dims():
+    arr = np.random.RandomState(1).rand(10).astype(np.float32)
+    sess = _ArrayReturningSession()
+    sess.setup(arr)
+
+    with pytest.raises(ValueError):
+        sess.infer_argmax(np.zeros((1,), dtype=np.float32))
+
+
+def test_ctc_label_decode_accepts_tuple_and_array():
+    character = list("abcdefghij")
+    decode = CTCLabelDecode(character=character)
+    logits = np.random.RandomState(2).rand(2, 6, len(character)).astype(np.float32)
+
+    out_from_array = decode(logits)
+    out_from_tuple = decode((logits.argmax(axis=2), logits.max(axis=2)))
+
+    assert out_from_array == out_from_tuple
+
+
+@pytest.mark.skipif(not import_package("torch"), reason="torch is not installed")
+def test_torch_infer_argmax_matches_base_reduction():
+    from rapidocr.inference_engine.pytorch.main import TorchInferSession
+
+    character = list("abcdefghij")
+
+    class _FakeTorchSession(TorchInferSession):
+        def __init__(self):
+            import torch
+
+            self.device = torch.device("cpu")
+
+            def predictor(x):
+                # Deterministic given x so __call__ and infer_argmax agree.
+                return torch.full(
+                    (x.shape[0], 25, len(character)), float(x.mean().item())
+                )
+
+            self.predictor = predictor
+
+    sess = _FakeTorchSession()
+    x = np.random.RandomState(3).rand(2, 3, 32, 100).astype(np.float32)
+
+    full = sess(x)  # reduce=False -> (B, T, vocab) ndarray
+    idx, prob = sess.infer_argmax(x)  # reduce=True -> (B, T) tuple
+
+    assert np.array_equal(idx, full.argmax(axis=-1))
+    assert np.array_equal(prob, full.max(axis=-1))