File size: 1,762 Bytes
3bbb319 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
import pytest
import torch
from mmpose.models import CPM
from mmpose.models.backbones.cpm import CpmBlock
def test_cpm_block():
with pytest.raises(AssertionError):
# len(channels) == len(kernels)
CpmBlock(
3, channels=[3, 3, 3], kernels=[
1,
])
# Test CPM Block
model = CpmBlock(3, channels=[3, 3, 3], kernels=[1, 1, 1])
model.train()
imgs = torch.randn(1, 3, 10, 10)
feat = model(imgs)
assert feat.shape == torch.Size([1, 3, 10, 10])
def test_cpm_backbone():
with pytest.raises(AssertionError):
# CPM's num_stacks should larger than 0
CPM(in_channels=3, out_channels=17, num_stages=-1)
with pytest.raises(AssertionError):
# CPM's in_channels should be 3
CPM(in_channels=2, out_channels=17)
# Test CPM
model = CPM(in_channels=3, out_channels=17, num_stages=1)
model.init_weights()
model.train()
imgs = torch.randn(1, 3, 256, 192)
feat = model(imgs)
assert len(feat) == 1
assert feat[0].shape == torch.Size([1, 17, 32, 24])
imgs = torch.randn(1, 3, 384, 288)
feat = model(imgs)
assert len(feat) == 1
assert feat[0].shape == torch.Size([1, 17, 48, 36])
imgs = torch.randn(1, 3, 368, 368)
feat = model(imgs)
assert len(feat) == 1
assert feat[0].shape == torch.Size([1, 17, 46, 46])
# Test CPM multi-stages
model = CPM(in_channels=3, out_channels=17, num_stages=2)
model.init_weights()
model.train()
imgs = torch.randn(1, 3, 368, 368)
feat = model(imgs)
assert len(feat) == 2
assert feat[0].shape == torch.Size([1, 17, 46, 46])
assert feat[1].shape == torch.Size([1, 17, 46, 46])
|