|
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
from detectron2.config import CfgNode
|
|
from detectron2.layers import Conv2d
|
|
|
|
from ..utils import initialize_module_params
|
|
from .registry import ROI_DENSEPOSE_HEAD_REGISTRY
|
|
|
|
|
|
@ROI_DENSEPOSE_HEAD_REGISTRY.register()
|
|
class DensePoseV1ConvXHead(nn.Module):
|
|
"""
|
|
Fully convolutional DensePose head.
|
|
"""
|
|
|
|
def __init__(self, cfg: CfgNode, input_channels: int):
|
|
"""
|
|
Initialize DensePose fully convolutional head
|
|
|
|
Args:
|
|
cfg (CfgNode): configuration options
|
|
input_channels (int): number of input channels
|
|
"""
|
|
super(DensePoseV1ConvXHead, self).__init__()
|
|
|
|
hidden_dim = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM
|
|
kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL
|
|
self.n_stacked_convs = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS
|
|
|
|
pad_size = kernel_size // 2
|
|
n_channels = input_channels
|
|
for i in range(self.n_stacked_convs):
|
|
layer = Conv2d(n_channels, hidden_dim, kernel_size, stride=1, padding=pad_size)
|
|
layer_name = self._get_layer_name(i)
|
|
self.add_module(layer_name, layer)
|
|
n_channels = hidden_dim
|
|
self.n_out_channels = n_channels
|
|
initialize_module_params(self)
|
|
|
|
def forward(self, features: torch.Tensor):
|
|
"""
|
|
Apply DensePose fully convolutional head to the input features
|
|
|
|
Args:
|
|
features (tensor): input features
|
|
Result:
|
|
A tensor of DensePose head outputs
|
|
"""
|
|
x = features
|
|
output = x
|
|
for i in range(self.n_stacked_convs):
|
|
layer_name = self._get_layer_name(i)
|
|
x = getattr(self, layer_name)(x)
|
|
x = F.relu(x)
|
|
output = x
|
|
return output
|
|
|
|
def _get_layer_name(self, i: int):
|
|
layer_name = "body_conv_fcn{}".format(i + 1)
|
|
return layer_name
|
|
|
