|
|
|
|
|
import pytest |
|
|
import torch |
|
|
|
|
|
from mmocr.models.textrecog.encoders import (ABIVisionModel, BaseEncoder, |
|
|
NRTREncoder, SAREncoder, |
|
|
SatrnEncoder, TransformerEncoder) |
|
|
|
|
|
|
|
|
def test_sar_encoder(): |
|
|
with pytest.raises(AssertionError): |
|
|
SAREncoder(enc_bi_rnn='bi') |
|
|
with pytest.raises(AssertionError): |
|
|
SAREncoder(enc_do_rnn=2) |
|
|
with pytest.raises(AssertionError): |
|
|
SAREncoder(enc_gru='gru') |
|
|
with pytest.raises(AssertionError): |
|
|
SAREncoder(d_model=512.5) |
|
|
with pytest.raises(AssertionError): |
|
|
SAREncoder(d_enc=200.5) |
|
|
with pytest.raises(AssertionError): |
|
|
SAREncoder(mask='mask') |
|
|
|
|
|
encoder = SAREncoder() |
|
|
encoder.init_weights() |
|
|
encoder.train() |
|
|
|
|
|
feat = torch.randn(1, 512, 4, 40) |
|
|
img_metas = [{'valid_ratio': 1.0}] |
|
|
with pytest.raises(AssertionError): |
|
|
encoder(feat, img_metas * 2) |
|
|
out_enc = encoder(feat, img_metas) |
|
|
|
|
|
assert out_enc.shape == torch.Size([1, 512]) |
|
|
|
|
|
|
|
|
def test_nrtr_encoder(): |
|
|
tf_encoder = NRTREncoder() |
|
|
tf_encoder.init_weights() |
|
|
tf_encoder.train() |
|
|
|
|
|
feat = torch.randn(1, 512, 1, 25) |
|
|
out_enc = tf_encoder(feat) |
|
|
print('hello', out_enc.size()) |
|
|
assert out_enc.shape == torch.Size([1, 25, 512]) |
|
|
|
|
|
|
|
|
def test_satrn_encoder(): |
|
|
satrn_encoder = SatrnEncoder() |
|
|
satrn_encoder.init_weights() |
|
|
satrn_encoder.train() |
|
|
|
|
|
feat = torch.randn(1, 512, 8, 25) |
|
|
out_enc = satrn_encoder(feat) |
|
|
assert out_enc.shape == torch.Size([1, 200, 512]) |
|
|
|
|
|
|
|
|
def test_base_encoder(): |
|
|
encoder = BaseEncoder() |
|
|
encoder.init_weights() |
|
|
encoder.train() |
|
|
|
|
|
feat = torch.randn(1, 256, 4, 40) |
|
|
out_enc = encoder(feat) |
|
|
assert out_enc.shape == torch.Size([1, 256, 4, 40]) |
|
|
|
|
|
|
|
|
def test_transformer_encoder(): |
|
|
model = TransformerEncoder() |
|
|
x = torch.randn(10, 512, 8, 32) |
|
|
assert model(x).shape == torch.Size([10, 512, 8, 32]) |
|
|
|
|
|
|
|
|
def test_abi_vision_model(): |
|
|
model = ABIVisionModel( |
|
|
decoder=dict(type='ABIVisionDecoder', max_seq_len=10, use_result=None)) |
|
|
x = torch.randn(1, 512, 8, 32) |
|
|
result = model(x) |
|
|
assert result['feature'].shape == torch.Size([1, 10, 512]) |
|
|
assert result['logits'].shape == torch.Size([1, 10, 90]) |
|
|
assert result['attn_scores'].shape == torch.Size([1, 10, 8, 32]) |
|
|
|
