Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,807 Bytes
938e515 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
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
# use the provided mesh names if not None and not an empty list
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
|