Spaces:
Runtime error
Runtime error
File size: 2,378 Bytes
2366e36 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
# Copyright (c) OpenMMLab. All rights reserved.
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])
|