|
|
|
|
| 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}")
|
|
|
