|
|
|
import copy |
|
|
|
import torch |
|
import torch.nn as nn |
|
from mmcv.cnn import ConvModule, constant_init, normal_init |
|
from torch.nn.modules.batchnorm import _BatchNorm |
|
|
|
from mmpose.utils import get_root_logger |
|
from ..builder import BACKBONES |
|
from .base_backbone import BaseBackbone |
|
from .utils import load_checkpoint |
|
|
|
|
|
class CpmBlock(nn.Module): |
|
"""CpmBlock for Convolutional Pose Machine. |
|
|
|
Args: |
|
in_channels (int): Input channels of this block. |
|
channels (list): Output channels of each conv module. |
|
kernels (list): Kernel sizes of each conv module. |
|
""" |
|
|
|
def __init__(self, |
|
in_channels, |
|
channels=(128, 128, 128), |
|
kernels=(11, 11, 11), |
|
norm_cfg=None): |
|
super().__init__() |
|
|
|
assert len(channels) == len(kernels) |
|
layers = [] |
|
for i in range(len(channels)): |
|
if i == 0: |
|
input_channels = in_channels |
|
else: |
|
input_channels = channels[i - 1] |
|
layers.append( |
|
ConvModule( |
|
input_channels, |
|
channels[i], |
|
kernels[i], |
|
padding=(kernels[i] - 1) // 2, |
|
norm_cfg=norm_cfg)) |
|
self.model = nn.Sequential(*layers) |
|
|
|
def forward(self, x): |
|
"""Model forward function.""" |
|
out = self.model(x) |
|
return out |
|
|
|
|
|
@BACKBONES.register_module() |
|
class CPM(BaseBackbone): |
|
"""CPM backbone. |
|
|
|
Convolutional Pose Machines. |
|
More details can be found in the `paper |
|
<https://arxiv.org/abs/1602.00134>`__ . |
|
|
|
Args: |
|
in_channels (int): The input channels of the CPM. |
|
out_channels (int): The output channels of the CPM. |
|
feat_channels (int): Feature channel of each CPM stage. |
|
middle_channels (int): Feature channel of conv after the middle stage. |
|
num_stages (int): Number of stages. |
|
norm_cfg (dict): Dictionary to construct and config norm layer. |
|
|
|
Example: |
|
>>> from mmpose.models import CPM |
|
>>> import torch |
|
>>> self = CPM(3, 17) |
|
>>> self.eval() |
|
>>> inputs = torch.rand(1, 3, 368, 368) |
|
>>> level_outputs = self.forward(inputs) |
|
>>> for level_output in level_outputs: |
|
... print(tuple(level_output.shape)) |
|
(1, 17, 46, 46) |
|
(1, 17, 46, 46) |
|
(1, 17, 46, 46) |
|
(1, 17, 46, 46) |
|
(1, 17, 46, 46) |
|
(1, 17, 46, 46) |
|
""" |
|
|
|
def __init__(self, |
|
in_channels, |
|
out_channels, |
|
feat_channels=128, |
|
middle_channels=32, |
|
num_stages=6, |
|
norm_cfg=dict(type='BN', requires_grad=True)): |
|
|
|
norm_cfg = copy.deepcopy(norm_cfg) |
|
super().__init__() |
|
|
|
assert in_channels == 3 |
|
|
|
self.num_stages = num_stages |
|
assert self.num_stages >= 1 |
|
|
|
self.stem = nn.Sequential( |
|
ConvModule(in_channels, 128, 9, padding=4, norm_cfg=norm_cfg), |
|
nn.MaxPool2d(kernel_size=3, stride=2, padding=1), |
|
ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg), |
|
nn.MaxPool2d(kernel_size=3, stride=2, padding=1), |
|
ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg), |
|
nn.MaxPool2d(kernel_size=3, stride=2, padding=1), |
|
ConvModule(128, 32, 5, padding=2, norm_cfg=norm_cfg), |
|
ConvModule(32, 512, 9, padding=4, norm_cfg=norm_cfg), |
|
ConvModule(512, 512, 1, padding=0, norm_cfg=norm_cfg), |
|
ConvModule(512, out_channels, 1, padding=0, act_cfg=None)) |
|
|
|
self.middle = nn.Sequential( |
|
ConvModule(in_channels, 128, 9, padding=4, norm_cfg=norm_cfg), |
|
nn.MaxPool2d(kernel_size=3, stride=2, padding=1), |
|
ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg), |
|
nn.MaxPool2d(kernel_size=3, stride=2, padding=1), |
|
ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg), |
|
nn.MaxPool2d(kernel_size=3, stride=2, padding=1)) |
|
|
|
self.cpm_stages = nn.ModuleList([ |
|
CpmBlock( |
|
middle_channels + out_channels, |
|
channels=[feat_channels, feat_channels, feat_channels], |
|
kernels=[11, 11, 11], |
|
norm_cfg=norm_cfg) for _ in range(num_stages - 1) |
|
]) |
|
|
|
self.middle_conv = nn.ModuleList([ |
|
nn.Sequential( |
|
ConvModule( |
|
128, middle_channels, 5, padding=2, norm_cfg=norm_cfg)) |
|
for _ in range(num_stages - 1) |
|
]) |
|
|
|
self.out_convs = nn.ModuleList([ |
|
nn.Sequential( |
|
ConvModule( |
|
feat_channels, |
|
feat_channels, |
|
1, |
|
padding=0, |
|
norm_cfg=norm_cfg), |
|
ConvModule(feat_channels, out_channels, 1, act_cfg=None)) |
|
for _ in range(num_stages - 1) |
|
]) |
|
|
|
def init_weights(self, pretrained=None): |
|
"""Initialize the weights in backbone. |
|
|
|
Args: |
|
pretrained (str, optional): Path to pre-trained weights. |
|
Defaults to None. |
|
""" |
|
if isinstance(pretrained, str): |
|
logger = get_root_logger() |
|
load_checkpoint(self, pretrained, strict=False, logger=logger) |
|
elif pretrained is None: |
|
for m in self.modules(): |
|
if isinstance(m, nn.Conv2d): |
|
normal_init(m, std=0.001) |
|
elif isinstance(m, (_BatchNorm, nn.GroupNorm)): |
|
constant_init(m, 1) |
|
else: |
|
raise TypeError('pretrained must be a str or None') |
|
|
|
def forward(self, x): |
|
"""Model forward function.""" |
|
stage1_out = self.stem(x) |
|
middle_out = self.middle(x) |
|
out_feats = [] |
|
|
|
out_feats.append(stage1_out) |
|
|
|
for ind in range(self.num_stages - 1): |
|
single_stage = self.cpm_stages[ind] |
|
out_conv = self.out_convs[ind] |
|
|
|
inp_feat = torch.cat( |
|
[out_feats[-1], self.middle_conv[ind](middle_out)], 1) |
|
cpm_feat = single_stage(inp_feat) |
|
out_feat = out_conv(cpm_feat) |
|
out_feats.append(out_feat) |
|
|
|
return out_feats |
|
|
