lib/dataset/Evaluator.py

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