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