|
|
|
|
|
import torch |
|
from torch import nn |
|
|
|
from detectron2.config import CfgNode |
|
from detectron2.layers import ConvTranspose2d, interpolate |
|
|
|
from ...structures import DensePoseEmbeddingPredictorOutput |
|
from ..utils import initialize_module_params |
|
from .registry import DENSEPOSE_PREDICTOR_REGISTRY |
|
|
|
|
|
@DENSEPOSE_PREDICTOR_REGISTRY.register() |
|
class DensePoseEmbeddingPredictor(nn.Module): |
|
""" |
|
Last layers of a DensePose model that take DensePose head outputs as an input |
|
and produce model outputs for continuous surface embeddings (CSE). |
|
""" |
|
|
|
def __init__(self, cfg: CfgNode, input_channels: int): |
|
""" |
|
Initialize predictor using configuration options |
|
|
|
Args: |
|
cfg (CfgNode): configuration options |
|
input_channels (int): input tensor size along the channel dimension |
|
""" |
|
super().__init__() |
|
dim_in = input_channels |
|
n_segm_chan = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS |
|
embed_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE |
|
kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL |
|
|
|
self.coarse_segm_lowres = ConvTranspose2d( |
|
dim_in, n_segm_chan, kernel_size, stride=2, padding=int(kernel_size / 2 - 1) |
|
) |
|
|
|
self.embed_lowres = ConvTranspose2d( |
|
dim_in, embed_size, kernel_size, stride=2, padding=int(kernel_size / 2 - 1) |
|
) |
|
self.scale_factor = cfg.MODEL.ROI_DENSEPOSE_HEAD.UP_SCALE |
|
initialize_module_params(self) |
|
|
|
def interp2d(self, tensor_nchw: torch.Tensor): |
|
""" |
|
Bilinear interpolation method to be used for upscaling |
|
|
|
Args: |
|
tensor_nchw (tensor): tensor of shape (N, C, H, W) |
|
Return: |
|
tensor of shape (N, C, Hout, Wout), where Hout and Wout are computed |
|
by applying the scale factor to H and W |
|
""" |
|
return interpolate( |
|
tensor_nchw, scale_factor=self.scale_factor, mode="bilinear", align_corners=False |
|
) |
|
|
|
def forward(self, head_outputs): |
|
""" |
|
Perform forward step on DensePose head outputs |
|
|
|
Args: |
|
head_outputs (tensor): DensePose head outputs, tensor of shape [N, D, H, W] |
|
""" |
|
embed_lowres = self.embed_lowres(head_outputs) |
|
coarse_segm_lowres = self.coarse_segm_lowres(head_outputs) |
|
embed = self.interp2d(embed_lowres) |
|
coarse_segm = self.interp2d(coarse_segm_lowres) |
|
return DensePoseEmbeddingPredictorOutput(embedding=embed, coarse_segm=coarse_segm) |
|
|
