Spaces:
Runtime error
Runtime error
File size: 525 Bytes
2366e36 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
from mmocr.models.textrecog.necks import FPNOCR
def test_fpn_ocr():
in_s1 = torch.rand(1, 128, 32, 256)
in_s2 = torch.rand(1, 256, 16, 128)
in_s3 = torch.rand(1, 512, 8, 64)
in_s4 = torch.rand(1, 512, 4, 32)
fpn_ocr = FPNOCR(in_channels=[128, 256, 512, 512], out_channels=256)
fpn_ocr.init_weights()
fpn_ocr.train()
out_neck = fpn_ocr((in_s1, in_s2, in_s3, in_s4))
assert out_neck[0].shape == torch.Size([1, 256, 32, 256])
|