mirror of https://github.com/open-mmlab/mmocr.git
17 lines
418 B
Python
17 lines
418 B
Python
import pytest
|
|
import torch
|
|
|
|
from mmocr.models.textrecog import SegHead
|
|
|
|
|
|
def test_seg_head():
|
|
with pytest.raises(AssertionError):
|
|
SegHead(num_classes='100')
|
|
with pytest.raises(AssertionError):
|
|
SegHead(num_classes=-1)
|
|
|
|
seg_head = SegHead(num_classes=37)
|
|
out_neck = (torch.rand(1, 128, 32, 32), )
|
|
out_head = seg_head(out_neck)
|
|
assert out_head.shape == torch.Size([1, 37, 32, 32])
|