| | |
| |
|
| | 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 |
| |
|
