def __init__(
self,
cfg,
texture_atlases_dict,
device="cuda",
default_class=0,
**kwargs,
):
self.embedder = build_densepose_embedder(cfg)
self.texture_image_dict = {}
self.alpha_dict = {}
for mesh_name in texture_atlases_dict.keys():
if texture_atlases_dict[mesh_name].shape[
-1] == 4: # Image with alpha channel
self.alpha_dict[mesh_name] = texture_atlases_dict[
mesh_name][:, :, -1] / 255.0
self.texture_image_dict[mesh_name] = texture_atlases_dict[
mesh_name][:, :, :3]
else:
self.alpha_dict[mesh_name] = texture_atlases_dict[
mesh_name].sum(axis=-1) > 0
self.texture_image_dict[mesh_name] = texture_atlases_dict[
mesh_name]
self.device = torch.device(device)
self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
self.default_class = default_class
self.mesh_vertex_embeddings = {
mesh_name: self.embedder(mesh_name).to(self.device)
for mesh_name in self.class_to_mesh_name.values()
}
def __init__(
self,
cfg,
inplace=True,
cmap=cv2.COLORMAP_JET,
alpha=0.7,
device="cuda",
default_class=0,
class_to_mesh_name=DEFAULT_CLASS_TO_MESH_NAME,
**kwargs,
):
self.mask_visualizer = MatrixVisualizer(inplace=inplace,
cmap=cmap,
val_scale=1.0,
alpha=alpha)
self.class_to_mesh_name = class_to_mesh_name
self.embedder = build_densepose_embedder(cfg)
self.device = torch.device(device)
self.default_class = default_class
self.embed_map_rescaled = get_smpl_euclidean_vertex_embedding()[:, 0]
self.embed_map_rescaled -= self.embed_map_rescaled.min()
self.embed_map_rescaled /= self.embed_map_rescaled.max()
self.mesh_vertex_embeddings = {
mesh_name: self.embedder(mesh_name).to(self.device)
for mesh_name in self.class_to_mesh_name.values()
}
def visualize(
self,
image_bgr: Image,
outputs_boxes_xywh_classes: Tuple[
Optional[DensePoseEmbeddingPredictorOutput], Optional[Boxes],
Optional[List[int]]],
) -> Image:
if outputs_boxes_xywh_classes[0] is None:
return image_bgr
embedder = build_densepose_embedder(self.cfg)
embed_map_rescaled = get_smpl_euclidean_vertex_embedding()[:, 0]
embed_map_rescaled -= embed_map_rescaled.min()
embed_map_rescaled /= embed_map_rescaled.max()
S, E, N, bboxes_xywh, pred_classes = self.extract_and_check_outputs_and_boxes(
outputs_boxes_xywh_classes)
mesh_vertex_embeddings = {
p: embedder(self.class_to_mesh_name[p]).to(self.device)
for p in np.unique(pred_classes)
}
for n in range(N):
x, y, w, h = bboxes_xywh[n].int().tolist()
closest_vertices, mask = self.get_closest_vertices_mask_from_ES(
E[[n]], S[[n]], h, w, mesh_vertex_embeddings[pred_classes[n]])
vis = (embed_map_rescaled[closest_vertices].clip(0, 1) *
255.0).cpu().numpy()
mask_numpy = mask.cpu().numpy().astype(dtype=np.uint8)
image_bgr = self.mask_visualizer.visualize(image_bgr, mask_numpy,
vis, [x, y, w, h])
return image_bgr
def build_inference_based_loaders(
cfg: CfgNode, model: torch.nn.Module
) -> Tuple[List[InferenceBasedLoader], List[float]]:
loaders = []
ratios = []
embedder = build_densepose_embedder(cfg).to(device=model.device) # pyre-ignore[16]
for dataset_spec in cfg.BOOTSTRAP_DATASETS:
dataset_cfg = get_bootstrap_dataset_config().clone()
dataset_cfg.merge_from_other_cfg(CfgNode(dataset_spec))
loader = build_inference_based_loader(cfg, dataset_cfg, model, embedder)
loaders.append(loader)
ratios.append(dataset_cfg.RATIO)
return loaders, ratios
def __init__(
self,
cfg,
inplace=True,
cmap=cv2.COLORMAP_JET,
alpha=0.7,
device="cuda",
default_class=0,
**kwargs,
):
self.mask_visualizer = MatrixVisualizer(
inplace=inplace, cmap=cmap, val_scale=1.0, alpha=alpha
)
self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
self.embedder = build_densepose_embedder(cfg)
self.device = torch.device(device)
self.default_class = default_class
self.mesh_vertex_embeddings = {
mesh_name: self.embedder(mesh_name).to(self.device)
for mesh_name in self.class_to_mesh_name.values()
if self.embedder.has_embeddings(mesh_name)
}
