|
|
|
|
| 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
|
|
|
