Spaces:
Build error
Build error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from unittest import TestCase | |
| import torch | |
| from mmpose.models.backbones import CPM | |
| from mmpose.models.backbones.cpm import CpmBlock | |
| class TestCPM(TestCase): | |
| def test_cpm_block(self): | |
| with self.assertRaises(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) | |
| self.assertEqual(feat.shape, torch.Size([1, 3, 10, 10])) | |
| def test_cpm_backbone(self): | |
| with self.assertRaises(AssertionError): | |
| # CPM's num_stacks should larger than 0 | |
| CPM(in_channels=3, out_channels=17, num_stages=-1) | |
| with self.assertRaises(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) | |
| self.assertEqual(len(feat), 1) | |
| self.assertEqual(feat[0].shape, torch.Size([1, 17, 32, 24])) | |
| imgs = torch.randn(1, 3, 384, 288) | |
| feat = model(imgs) | |
| self.assertEqual(len(feat), 1) | |
| self.assertEqual(feat[0].shape, torch.Size([1, 17, 48, 36])) | |
| imgs = torch.randn(1, 3, 368, 368) | |
| feat = model(imgs) | |
| self.assertEqual(len(feat), 1) | |
| self.assertEqual(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) | |
| self.assertEqual(len(feat), 2) | |
| self.assertEqual(feat[0].shape, torch.Size([1, 17, 46, 46])) | |
| self.assertEqual(feat[1].shape, torch.Size([1, 17, 46, 46])) | |