mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 05:20:49 +08:00
101 lines
3.0 KiB
Python
101 lines
3.0 KiB
Python
# -*- 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))
|