|
|
|
|
| import json
|
| import logging
|
| from typing import List, Optional
|
| import torch
|
| from torch import nn
|
|
|
| from detectron2.utils.file_io import PathManager
|
|
|
| from densepose.structures.mesh import create_mesh
|
|
|
|
|
| class MeshAlignmentEvaluator:
|
| """
|
| Class for evaluation of 3D mesh alignment based on the learned vertex embeddings
|
| """
|
|
|
| def __init__(self, embedder: nn.Module, mesh_names: Optional[List[str]]):
|
| self.embedder = embedder
|
|
|
| self.mesh_names = mesh_names if mesh_names else embedder.mesh_names
|
| self.logger = logging.getLogger(__name__)
|
| with PathManager.open(
|
| "https://dl.fbaipublicfiles.com/densepose/data/cse/mesh_keyvertices_v0.json", "r"
|
| ) as f:
|
| self.mesh_keyvertices = json.load(f)
|
|
|
| def evaluate(self):
|
| ge_per_mesh = {}
|
| gps_per_mesh = {}
|
| for mesh_name_1 in self.mesh_names:
|
| avg_errors = []
|
| avg_gps = []
|
| embeddings_1 = self.embedder(mesh_name_1)
|
| keyvertices_1 = self.mesh_keyvertices[mesh_name_1]
|
| keyvertex_names_1 = list(keyvertices_1.keys())
|
| keyvertex_indices_1 = [keyvertices_1[name] for name in keyvertex_names_1]
|
| for mesh_name_2 in self.mesh_names:
|
| if mesh_name_1 == mesh_name_2:
|
| continue
|
| embeddings_2 = self.embedder(mesh_name_2)
|
| keyvertices_2 = self.mesh_keyvertices[mesh_name_2]
|
| sim_matrix_12 = embeddings_1[keyvertex_indices_1].mm(embeddings_2.T)
|
| vertices_2_matching_keyvertices_1 = sim_matrix_12.argmax(axis=1)
|
| mesh_2 = create_mesh(mesh_name_2, embeddings_2.device)
|
| geodists = mesh_2.geodists[
|
| vertices_2_matching_keyvertices_1,
|
| [keyvertices_2[name] for name in keyvertex_names_1],
|
| ]
|
| Current_Mean_Distances = 0.255
|
| gps = (-(geodists**2) / (2 * (Current_Mean_Distances**2))).exp()
|
| avg_errors.append(geodists.mean().item())
|
| avg_gps.append(gps.mean().item())
|
|
|
| ge_mean = torch.as_tensor(avg_errors).mean().item()
|
| gps_mean = torch.as_tensor(avg_gps).mean().item()
|
| ge_per_mesh[mesh_name_1] = ge_mean
|
| gps_per_mesh[mesh_name_1] = gps_mean
|
| ge_mean_global = torch.as_tensor(list(ge_per_mesh.values())).mean().item()
|
| gps_mean_global = torch.as_tensor(list(gps_per_mesh.values())).mean().item()
|
| per_mesh_metrics = {
|
| "GE": ge_per_mesh,
|
| "GPS": gps_per_mesh,
|
| }
|
| return ge_mean_global, gps_mean_global, per_mesh_metrics
|
|
|
