|
|
|
|
|
|
|
|
|
from typing import Optional
|
|
from torch import nn
|
|
|
|
from detectron2.config import CfgNode
|
|
|
|
from .cse.embedder import Embedder
|
|
from .filter import DensePoseDataFilter
|
|
|
|
|
|
def build_densepose_predictor(cfg: CfgNode, input_channels: int):
|
|
"""
|
|
Create an instance of DensePose predictor based on configuration options.
|
|
|
|
Args:
|
|
cfg (CfgNode): configuration options
|
|
input_channels (int): input tensor size along the channel dimension
|
|
Return:
|
|
An instance of DensePose predictor
|
|
"""
|
|
from .predictors import DENSEPOSE_PREDICTOR_REGISTRY
|
|
|
|
predictor_name = cfg.MODEL.ROI_DENSEPOSE_HEAD.PREDICTOR_NAME
|
|
return DENSEPOSE_PREDICTOR_REGISTRY.get(predictor_name)(cfg, input_channels)
|
|
|
|
|
|
def build_densepose_data_filter(cfg: CfgNode):
|
|
"""
|
|
Build DensePose data filter which selects data for training
|
|
|
|
Args:
|
|
cfg (CfgNode): configuration options
|
|
|
|
Return:
|
|
Callable: list(Tensor), list(Instances) -> list(Tensor), list(Instances)
|
|
An instance of DensePose filter, which takes feature tensors and proposals
|
|
as an input and returns filtered features and proposals
|
|
"""
|
|
dp_filter = DensePoseDataFilter(cfg)
|
|
return dp_filter
|
|
|
|
|
|
def build_densepose_head(cfg: CfgNode, input_channels: int):
|
|
"""
|
|
Build DensePose head based on configurations options
|
|
|
|
Args:
|
|
cfg (CfgNode): configuration options
|
|
input_channels (int): input tensor size along the channel dimension
|
|
Return:
|
|
An instance of DensePose head
|
|
"""
|
|
from .roi_heads.registry import ROI_DENSEPOSE_HEAD_REGISTRY
|
|
|
|
head_name = cfg.MODEL.ROI_DENSEPOSE_HEAD.NAME
|
|
return ROI_DENSEPOSE_HEAD_REGISTRY.get(head_name)(cfg, input_channels)
|
|
|
|
|
|
def build_densepose_losses(cfg: CfgNode):
|
|
"""
|
|
Build DensePose loss based on configurations options
|
|
|
|
Args:
|
|
cfg (CfgNode): configuration options
|
|
Return:
|
|
An instance of DensePose loss
|
|
"""
|
|
from .losses import DENSEPOSE_LOSS_REGISTRY
|
|
|
|
loss_name = cfg.MODEL.ROI_DENSEPOSE_HEAD.LOSS_NAME
|
|
return DENSEPOSE_LOSS_REGISTRY.get(loss_name)(cfg)
|
|
|
|
|
|
def build_densepose_embedder(cfg: CfgNode) -> Optional[nn.Module]:
|
|
"""
|
|
Build embedder used to embed mesh vertices into an embedding space.
|
|
Embedder contains sub-embedders, one for each mesh ID.
|
|
|
|
Args:
|
|
cfg (cfgNode): configuration options
|
|
Return:
|
|
Embedding module
|
|
"""
|
|
if cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS:
|
|
return Embedder(cfg)
|
|
return None
|
|
|
