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