mmocr/tests/test_models/test_ocr_backbone.py

73 lines
1.7 KiB
Python
Raw Normal View History

# Copyright (c) OpenMMLab. All rights reserved.
2021-04-02 23:54:57 +08:00
import pytest
import torch
from mmocr.models.textrecog.backbones import (ResNet31OCR, ResNetABI,
ShallowCNN, VeryDeepVgg)
2021-04-02 23:54:57 +08:00
def test_resnet31_ocr_backbone():
"""Test resnet backbone."""
with pytest.raises(AssertionError):
ResNet31OCR(2.5)
with pytest.raises(AssertionError):
ResNet31OCR(3, layers=5)
with pytest.raises(AssertionError):
ResNet31OCR(3, channels=5)
# Test ResNet18 forward
model = ResNet31OCR()
model.init_weights()
model.train()
imgs = torch.randn(1, 3, 32, 160)
feat = model(imgs)
assert feat.shape == torch.Size([1, 512, 4, 40])
def test_vgg_deep_vgg_ocr_backbone():
model = VeryDeepVgg()
model.init_weights()
model.train()
imgs = torch.randn(1, 3, 32, 160)
feats = model(imgs)
assert feats.shape == torch.Size([1, 512, 1, 41])
[Feature] Add Satrn (#405) * Add SATRN * Create satrn_small_academic.py * Update README.md * change config name * Update mmocr/models/textrecog/backbones/shallow_cnn.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update configs/textrecog/satrn/satrn_academic.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update configs/textrecog/satrn/satrn_small.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update shallow_cnn.py * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update test_ocr_encoder.py * change keep_aspect_ratio=False * Update transformer_layer.py * Update configs/textrecog/satrn/satrn_small.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update configs/textrecog/satrn/satrn_academic.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update mmocr/models/textrecog/layers/transformer_layer.py Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update transformer_layer.py * Apply suggestions from code review Co-authored-by: Tong Gao <gaotongxiao@gmail.com> * Update transformer_layer.py * update satrn readme * add satrn to ocr.py * add satrn_sm and fix configs * add a test for config * add copyright info * use mmocr registry Co-authored-by: Tong Gao <gaotongxiao@gmail.com>
2021-08-19 22:02:58 +08:00
def test_shallow_cnn_ocr_backbone():
model = ShallowCNN()
model.init_weights()
model.train()
imgs = torch.randn(1, 1, 32, 100)
feat = model(imgs)
assert feat.shape == torch.Size([1, 512, 8, 25])
def test_resnet_abi():
"""Test resnet backbone."""
with pytest.raises(AssertionError):
ResNetABI(2.5)
with pytest.raises(AssertionError):
ResNetABI(3, arch_settings=5)
with pytest.raises(AssertionError):
ResNetABI(3, stem_channels=None)
with pytest.raises(AssertionError):
ResNetABI(arch_settings=[3, 4, 6, 6], strides=[1, 2, 1, 2, 1])
# Test forwarding
model = ResNetABI()
model.train()
imgs = torch.randn(1, 3, 32, 160)
feat = model(imgs)
assert feat.shape == torch.Size([1, 512, 8, 40])