| |
|
|
| import logging |
| import numpy as np |
| import pickle |
| from enum import Enum |
| from typing import Optional |
| import torch |
| from torch import nn |
|
|
| from detectron2.config import CfgNode |
| from detectron2.utils.file_io import PathManager |
|
|
| from .vertex_direct_embedder import VertexDirectEmbedder |
| from .vertex_feature_embedder import VertexFeatureEmbedder |
|
|
|
|
| class EmbedderType(Enum): |
| """ |
| Embedder type which defines how vertices are mapped into the embedding space: |
| - "vertex_direct": direct vertex embedding |
| - "vertex_feature": embedding vertex features |
| """ |
|
|
| VERTEX_DIRECT = "vertex_direct" |
| VERTEX_FEATURE = "vertex_feature" |
|
|
|
|
| def create_embedder(embedder_spec: CfgNode, embedder_dim: int) -> nn.Module: |
| """ |
| Create an embedder based on the provided configuration |
| |
| Args: |
| embedder_spec (CfgNode): embedder configuration |
| embedder_dim (int): embedding space dimensionality |
| Return: |
| An embedder instance for the specified configuration |
| Raises ValueError, in case of unexpected embedder type |
| """ |
| embedder_type = EmbedderType(embedder_spec.TYPE) |
| if embedder_type == EmbedderType.VERTEX_DIRECT: |
| embedder = VertexDirectEmbedder( |
| num_vertices=embedder_spec.NUM_VERTICES, |
| embed_dim=embedder_dim, |
| ) |
| if embedder_spec.INIT_FILE != "": |
| embedder.load(embedder_spec.INIT_FILE) |
| elif embedder_type == EmbedderType.VERTEX_FEATURE: |
| embedder = VertexFeatureEmbedder( |
| num_vertices=embedder_spec.NUM_VERTICES, |
| feature_dim=embedder_spec.FEATURE_DIM, |
| embed_dim=embedder_dim, |
| train_features=embedder_spec.FEATURES_TRAINABLE, |
| ) |
| if embedder_spec.INIT_FILE != "": |
| embedder.load(embedder_spec.INIT_FILE) |
| else: |
| raise ValueError(f"Unexpected embedder type {embedder_type}") |
|
|
| if not embedder_spec.IS_TRAINABLE: |
| embedder.requires_grad_(False) |
|
|
| return embedder |
|
|
|
|
| class Embedder(nn.Module): |
| """ |
| Embedder module that serves as a container for embedders to use with different |
| meshes. Extends Module to automatically save / load state dict. |
| """ |
|
|
| DEFAULT_MODEL_CHECKPOINT_PREFIX = "roi_heads.embedder." |
|
|
| def __init__(self, cfg: CfgNode): |
| """ |
| Initialize mesh embedders. An embedder for mesh `i` is stored in a submodule |
| "embedder_{i}". |
| |
| Args: |
| cfg (CfgNode): configuration options |
| """ |
| super(Embedder, self).__init__() |
| self.mesh_names = set() |
| embedder_dim = cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE |
| logger = logging.getLogger(__name__) |
| for mesh_name, embedder_spec in cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS.items(): |
| logger.info(f"Adding embedder embedder_{mesh_name} with spec {embedder_spec}") |
| self.add_module(f"embedder_{mesh_name}", create_embedder(embedder_spec, embedder_dim)) |
| self.mesh_names.add(mesh_name) |
| if cfg.MODEL.WEIGHTS != "": |
| self.load_from_model_checkpoint(cfg.MODEL.WEIGHTS) |
|
|
| def load_from_model_checkpoint(self, fpath: str, prefix: Optional[str] = None): |
| if prefix is None: |
| prefix = Embedder.DEFAULT_MODEL_CHECKPOINT_PREFIX |
| state_dict = None |
| if fpath.endswith(".pkl"): |
| with PathManager.open(fpath, "rb") as hFile: |
| state_dict = pickle.load(hFile, encoding="latin1") |
| else: |
| with PathManager.open(fpath, "rb") as hFile: |
| state_dict = torch.load(hFile, map_location=torch.device("cpu")) |
| if state_dict is not None and "model" in state_dict: |
| state_dict_local = {} |
| for key in state_dict["model"]: |
| if key.startswith(prefix): |
| v_key = state_dict["model"][key] |
| if isinstance(v_key, np.ndarray): |
| v_key = torch.from_numpy(v_key) |
| state_dict_local[key[len(prefix) :]] = v_key |
| |
| self.load_state_dict(state_dict_local, strict=False) |
|
|
| def forward(self, mesh_name: str) -> torch.Tensor: |
| """ |
| Produce vertex embeddings for the specific mesh; vertex embeddings are |
| a tensor of shape [N, D] where: |
| N = number of vertices |
| D = number of dimensions in the embedding space |
| Args: |
| mesh_name (str): name of a mesh for which to obtain vertex embeddings |
| Return: |
| Vertex embeddings, a tensor of shape [N, D] |
| """ |
| return getattr(self, f"embedder_{mesh_name}")() |
|
|
| def has_embeddings(self, mesh_name: str) -> bool: |
| return hasattr(self, f"embedder_{mesh_name}") |
|
|
