File size: 4,825 Bytes
938e515 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
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
# non-strict loading to finetune on different meshes
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}")
|