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