mirror of
https://github.com/open-mmlab/mmocr.git
synced 2025-06-03 21:54:47 +08:00
32 lines
865 B
Python
32 lines
865 B
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import pytest
|
|
import torch
|
|
|
|
from mmocr.models.textdet.necks import FPN_UNet
|
|
|
|
|
|
def test_fpn_unet_neck():
|
|
s = 64
|
|
feat_sizes = [s // 2**i for i in range(4)] # [64, 32, 16, 8]
|
|
in_channels = [8, 16, 32, 64]
|
|
out_channels = 4
|
|
|
|
# len(in_channcels) is not equal to 4
|
|
with pytest.raises(AssertionError):
|
|
FPN_UNet(in_channels + [128], out_channels)
|
|
|
|
# `out_channels` is not int type
|
|
with pytest.raises(AssertionError):
|
|
FPN_UNet(in_channels, [2, 4])
|
|
|
|
feats = [
|
|
torch.rand(1, in_channels[i], feat_sizes[i], feat_sizes[i])
|
|
for i in range(len(in_channels))
|
|
]
|
|
|
|
fpn_unet_neck = FPN_UNet(in_channels, out_channels)
|
|
fpn_unet_neck.init_weights()
|
|
|
|
out_neck = fpn_unet_neck(feats)
|
|
assert out_neck.shape == torch.Size([1, out_channels, s * 4, s * 4])
|