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