diff --git a/doctr/datasets/utils.py b/doctr/datasets/utils.py index 77b0aceb4c..9b7d3f711b 100644 --- a/doctr/datasets/utils.py +++ b/doctr/datasets/utils.py @@ -111,10 +111,11 @@ def decode_sequence( """ if not isinstance(input_seq, (Sequence, np.ndarray)): raise TypeError("Invalid sequence type") - if isinstance(input_seq, np.ndarray) and ( - not np.issubdtype(input_seq.dtype, np.integer) or (input_seq.size > 0 and input_seq.max() >= len(mapping)) - ): - raise AssertionError("Input must be an array of int, with max less than mapping size") + if isinstance(input_seq, np.ndarray) and not np.issubdtype(input_seq.dtype, np.integer): + raise AssertionError("Input must be an array of int, with values between 0 and the mapping size") + # A negative index would silently wrap around to the end of the mapping (e.g. an EOS of -1 read as a char) + if len(input_seq) > 0 and (min(input_seq) < 0 or max(input_seq) >= len(mapping)): + raise AssertionError("Input must be an array of int, with values between 0 and the mapping size") return "".join(map(mapping.__getitem__, input_seq)) diff --git a/tests/common/test_datasets_utils.py b/tests/common/test_datasets_utils.py index 9343c48908..d48518e1ba 100644 --- a/tests/common/test_datasets_utils.py +++ b/tests/common/test_datasets_utils.py @@ -61,6 +61,13 @@ def test_decode_sequence(): utils.decode_sequence(np.array([2, 10]), mapping) with pytest.raises(AssertionError): utils.decode_sequence(np.array([2, 4.5]), mapping) + # negative indices (such as the default EOS of -1) are not wrapped around to the last chars of the mapping + with pytest.raises(AssertionError): + utils.decode_sequence(np.array([2, -1]), mapping) + with pytest.raises(AssertionError): + utils.decode_sequence([2, -1], mapping) + with pytest.raises(AssertionError): + utils.decode_sequence([2, 10], mapping) assert utils.decode_sequence([3, 4, 3, 4], mapping) == "dede" assert utils.decode_sequence(np.array([3, 4, 3, 4], dtype=np.int64), mapping) == "dede"