Spaces:
Build error
Build error
| # 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]) | |