mmpose-0.29.0/tests/test_backbones/test_cpm.py

# 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])