# -*- coding: utf-8 -*- # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is # holder of all proprietary rights on this computer program. # You can only use this computer program if you have closed # a license agreement with MPG or you get the right to use the computer # program from someone who is authorized to grant you that right. # Any use of the computer program without a valid license is prohibited and # liable to prosecution. # # Copyright©2019 Max-Planck-Gesellschaft zur Förderung # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute # for Intelligent Systems. All rights reserved. # # Contact: ps-license@tuebingen.mpg.de from lib.renderer.gl.normal_render import NormalRender from lib.dataset.mesh_util import projection from lib.common.render import Render from PIL import Image import numpy as np import torch from torch import nn import trimesh import os.path as osp from PIL import Image class Evaluator: _normal_render = None @staticmethod def init_gl(): Evaluator._normal_render = NormalRender(width=512, height=512) def __init__(self, device): self.device = device self.render = Render(size=512, device=self.device) self.error_term = nn.MSELoss() self.offset = 0.0 self.scale_factor = None def set_mesh(self, result_dict, scale_factor=1.0, offset=0.0): for key in result_dict.keys(): if torch.is_tensor(result_dict[key]): result_dict[key] = result_dict[key].detach().cpu().numpy() for k, v in result_dict.items(): setattr(self, k, v) self.scale_factor = scale_factor self.offset = offset def _render_normal(self, mesh, deg, norms=None): view_mat = np.identity(4) rz = deg / 180.0 * np.pi model_mat = np.identity(4) model_mat[:3, :3] = self._normal_render.euler_to_rot_mat(0, rz, 0) model_mat[1, 3] = self.offset view_mat[2, 2] *= -1 self._normal_render.set_matrices(view_mat, model_mat) if norms is None: norms = mesh.vertex_normals self._normal_render.set_normal_mesh(self.scale_factor * mesh.vertices, mesh.faces, norms, mesh.faces) self._normal_render.draw() normal_img = self._normal_render.get_color() return normal_img def render_mesh_list(self, mesh_lst): self.offset = 0.0 self.scale_factor = 1.0 full_list = [] for mesh in mesh_lst: row_lst = [] for deg in np.arange(0, 360, 90): normal = self._render_normal(mesh, deg) row_lst.append(normal) full_list.append(np.concatenate(row_lst, axis=1)) res_array = np.concatenate(full_list, axis=0) return res_array def _get_reproj_normal_error(self, deg): tgt_normal = self._render_normal(self.tgt_mesh, deg) src_normal = self._render_normal(self.src_mesh, deg) error = (((src_normal[:, :, :3] - tgt_normal[:, :, :3])**2).sum(axis=2).mean(axis=(0, 1))) return error, [src_normal, tgt_normal] def render_normal(self, verts, faces): verts = verts[0].detach().cpu().numpy() faces = faces[0].detach().cpu().numpy() mesh_F = trimesh.Trimesh(verts * np.array([1.0, -1.0, 1.0]), faces) mesh_B = trimesh.Trimesh(verts * np.array([1.0, -1.0, -1.0]), faces) self.scale_factor = 1.0 normal_F = self._render_normal(mesh_F, 0) normal_B = self._render_normal(mesh_B, 0, norms=mesh_B.vertex_normals * np.array([-1.0, -1.0, 1.0])) mask = normal_F[:, :, 3:4] normal_F = (torch.as_tensor(2.0 * (normal_F - 0.5) * mask).permute( 2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device)) normal_B = (torch.as_tensor(2.0 * (normal_B - 0.5) * mask).permute( 2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device)) return {"T_normal_F": normal_F, "T_normal_B": normal_B} def calculate_normal_consist( self, frontal=True, back=True, left=True, right=True, save_demo_img=None, return_demo=False, ): # reproj error # if save_demo_img is not None, save a visualization at the given path (etc, "./test.png") if self._normal_render is None: print( "In order to use normal render, " "you have to call init_gl() before initialing any evaluator objects." ) return -1 side_cnt = 0 total_error = 0 demo_list = [] if frontal: side_cnt += 1 error, normal_lst = self._get_reproj_normal_error(0) total_error += error demo_list.append(np.concatenate(normal_lst, axis=0)) if back: side_cnt += 1 error, normal_lst = self._get_reproj_normal_error(180) total_error += error demo_list.append(np.concatenate(normal_lst, axis=0)) if left: side_cnt += 1 error, normal_lst = self._get_reproj_normal_error(90) total_error += error demo_list.append(np.concatenate(normal_lst, axis=0)) if right: side_cnt += 1 error, normal_lst = self._get_reproj_normal_error(270) total_error += error demo_list.append(np.concatenate(normal_lst, axis=0)) if save_demo_img is not None: res_array = np.concatenate(demo_list, axis=1) res_img = Image.fromarray((res_array * 255).astype(np.uint8)) res_img.save(save_demo_img) if return_demo: res_array = np.concatenate(demo_list, axis=1) return res_array else: return total_error def space_transfer(self): # convert from GT to SDF self.verts_pr -= self.recon_size / 2.0 self.verts_pr /= self.recon_size / 2.0 self.verts_gt = projection(self.verts_gt, self.calib) self.verts_gt[:, 1] *= -1 self.tgt_mesh = trimesh.Trimesh(self.verts_gt, self.faces_gt) self.src_mesh = trimesh.Trimesh(self.verts_pr, self.faces_pr) # (self.tgt_mesh+self.src_mesh).show() def export_mesh(self, dir, name): self.tgt_mesh.visual.vertex_colors = np.array([255, 0, 0]) self.src_mesh.visual.vertex_colors = np.array([0, 255, 0]) (self.tgt_mesh + self.src_mesh).export( osp.join(dir, f"{name}_gt_pr.obj")) def calculate_chamfer_p2s(self, sampled_points=1000): """calculate the geometry metrics [chamfer, p2s, chamfer_H, p2s_H] Args: verts_gt (torch.cuda.tensor): [N, 3] faces_gt (torch.cuda.tensor): [M, 3] verts_pr (torch.cuda.tensor): [N', 3] faces_pr (torch.cuda.tensor): [M', 3] sampled_points (int, optional): use smaller number for faster testing. Defaults to 1000. Returns: tuple: chamfer, p2s, chamfer_H, p2s_H """ gt_surface_pts, _ = trimesh.sample.sample_surface_even( self.tgt_mesh, sampled_points) pred_surface_pts, _ = trimesh.sample.sample_surface_even( self.src_mesh, sampled_points) _, dist_pred_gt, _ = trimesh.proximity.closest_point( self.src_mesh, gt_surface_pts) _, dist_gt_pred, _ = trimesh.proximity.closest_point( self.tgt_mesh, pred_surface_pts) dist_pred_gt[np.isnan(dist_pred_gt)] = 0 dist_gt_pred[np.isnan(dist_gt_pred)] = 0 chamfer_dist = 0.5 * (dist_pred_gt.mean() + dist_gt_pred.mean()).item() * 100 p2s_dist = dist_pred_gt.mean().item() * 100 return chamfer_dist, p2s_dist def calc_acc(self, output, target, thres=0.5, use_sdf=False): # # remove the surface points with thres # non_surf_ids = (target != thres) # output = output[non_surf_ids] # target = target[non_surf_ids] with torch.no_grad(): output = output.masked_fill(output < thres, 0.0) output = output.masked_fill(output > thres, 1.0) if use_sdf: target = target.masked_fill(target < thres, 0.0) target = target.masked_fill(target > thres, 1.0) acc = output.eq(target).float().mean() # iou, precison, recall output = output > thres target = target > thres union = output | target inter = output & target _max = torch.tensor(1.0).to(output.device) union = max(union.sum().float(), _max) true_pos = max(inter.sum().float(), _max) vol_pred = max(output.sum().float(), _max) vol_gt = max(target.sum().float(), _max) return acc, true_pos / union, true_pos / vol_pred, true_pos / vol_gt