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