pifuhd/lib/evaluator.py

# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import trimesh
import trimesh.proximity
import trimesh.sample
import numpy as np
import math
import os
from PIL import Image

import argparse

def euler_to_rot_mat(r_x, r_y, r_z):
    R_x = np.array([[1, 0, 0],
                    [0, math.cos(r_x), -math.sin(r_x)],
                    [0, math.sin(r_x), math.cos(r_x)]
                    ])

    R_y = np.array([[math.cos(r_y), 0, math.sin(r_y)],
                    [0, 1, 0],
                    [-math.sin(r_y), 0, math.cos(r_y)]
                    ])

    R_z = np.array([[math.cos(r_z), -math.sin(r_z), 0],
                    [math.sin(r_z), math.cos(r_z), 0],
                    [0, 0, 1]
                    ])

    R = np.dot(R_z, np.dot(R_y, R_x))

    return R


class MeshEvaluator:
    _normal_render = None

    @staticmethod
    def init_gl():
        from .render.gl.normal_render import NormalRender
        MeshEvaluator._normal_render = NormalRender(width=512, height=512)

    def __init__(self):
        pass

    def set_mesh(self, src_path, tgt_path, scale_factor=1.0, offset=0):
        self.src_mesh = trimesh.load(src_path)
        self.tgt_mesh = trimesh.load(tgt_path)

        self.scale_factor = scale_factor
        self.offset = offset


    def get_chamfer_dist(self, num_samples=10000):
        # Chamfer
        src_surf_pts, _ = trimesh.sample.sample_surface(self.src_mesh, num_samples)
        tgt_surf_pts, _ = trimesh.sample.sample_surface(self.tgt_mesh, num_samples)

        _, src_tgt_dist, _ = trimesh.proximity.closest_point(self.tgt_mesh, src_surf_pts)
        _, tgt_src_dist, _ = trimesh.proximity.closest_point(self.src_mesh, tgt_surf_pts)

        src_tgt_dist[np.isnan(src_tgt_dist)] = 0
        tgt_src_dist[np.isnan(tgt_src_dist)] = 0

        src_tgt_dist = src_tgt_dist.mean()
        tgt_src_dist = tgt_src_dist.mean()

        chamfer_dist = (src_tgt_dist + tgt_src_dist) / 2

        return chamfer_dist

    def get_surface_dist(self, num_samples=10000):
        # P2S
        src_surf_pts, _ = trimesh.sample.sample_surface(self.src_mesh, num_samples)

        _, src_tgt_dist, _ = trimesh.proximity.closest_point(self.tgt_mesh, src_surf_pts)

        src_tgt_dist[np.isnan(src_tgt_dist)] = 0

        src_tgt_dist = src_tgt_dist.mean()

        return src_tgt_dist

    def _render_normal(self, mesh, deg):
        view_mat = np.identity(4)
        view_mat[:3, :3] *= 2 / 256
        rz = deg / 180. * np.pi
        model_mat = np.identity(4)
        model_mat[:3, :3] = 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)
        self._normal_render.set_normal_mesh(self.scale_factor*mesh.vertices, mesh.faces, mesh.vertex_normals, mesh.faces)
        self._normal_render.draw()
        normal_img = self._normal_render.get_color()
        return normal_img

    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).mean() * 3

        return error, src_normal, tgt_normal

    def get_reproj_normal_error(self, frontal=True, back=True, left=True, right=True, save_demo_img=None):
        # 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, src_normal, tgt_normal = self._get_reproj_normal_error(0)
            total_error += error
            demo_list.append(np.concatenate([src_normal, tgt_normal], axis=0))
        if back:
            side_cnt += 1
            error, src_normal, tgt_normal = self._get_reproj_normal_error(180)
            total_error += error
            demo_list.append(np.concatenate([src_normal, tgt_normal], axis=0))
        if left:
            side_cnt += 1
            error, src_normal, tgt_normal = self._get_reproj_normal_error(90)
            total_error += error
            demo_list.append(np.concatenate([src_normal, tgt_normal], axis=0))
        if right:
            side_cnt += 1
            error, src_normal, tgt_normal = self._get_reproj_normal_error(270)
            total_error += error
            demo_list.append(np.concatenate([src_normal, tgt_normal], 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)
        return total_error / side_cnt

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-r', '--root', type=str, required=True)
    parser.add_argument('-t', '--tar_path', type=str, required=True)
    args = parser.parse_args()

    evaluator = MeshEvaluator()
    evaluator.init_gl()

    def run(root, exp_name, tar_path):
        src_path = os.path.join(root, exp_name, 'recon')
        rp_path = os.path.join(tar_path, 'RP', 'GEO', 'OBJ')
        bf_path = os.path.join(tar_path, 'BUFF', 'GEO', 'PLY')

        buff_files = [f for f in os.listdir(bf_path) if '.ply' in f]

        src_names = ['0_0_00.obj', '90_0_00.obj', '180_0_00.obj', '270_0_00.obj']

        total_vals = []
        items = []
        for file in buff_files:
            tar_name = os.path.join(bf_path, file)
            name = tar_name.split('/')[-1][:-4]

            for src in src_names:
                src_name = os.path.join(src_path, 'result_%s_%s' % (name, src))
                if not os.path.exists(src_name):
                    continue
                evaluator.set_mesh(src_name, tar_name, 0.13, -40)

                vals = []
                vals.append(0.1 * evaluator.get_chamfer_dist())
                vals.append(0.1 * evaluator.get_surface_dist())
                vals.append(4.0 * evaluator.get_reproj_normal_error(save_demo_img=os.path.join(src_path, '%s_%s.png' % (name, src[:-4]))))

                item = {
                    'name': '%s_%s' % (name, src),
                    'vals': vals
                }

                total_vals.append(vals)
                items.append(item)

        vals = np.array(total_vals).mean(0)
        buf_val = vals

        np.save(os.path.join(root, exp_name, 'buff-item.npy'), np.array(items))
        np.save(os.path.join(root, exp_name, 'buff-vals.npy'), total_vals)

        rp_files = [f for f in os.listdir(rp_path) if '.obj' in f]

        total_vals = []
        items = []
        for file in rp_files:
            tar_name = os.path.join(rp_path, file)
            name = tar_name.split('/')[-1][:-9]

            for src in src_names:
                src_name = os.path.join(src_path, 'result_%s_%s' % (name, src))
                if not os.path.exists(src_name):
                    continue

                evaluator.set_mesh(src_name, tar_name, 1.3, -120)

                vals = []
                vals.append(evaluator.get_chamfer_dist())
                vals.append(evaluator.get_surface_dist())
                vals.append(4.0 * evaluator.get_reproj_normal_error(save_demo_img=os.path.join(src_path, '%s_%s.png' % (name, src[:-4]))))

                item = {
                    'name': '%s_%s' % (name, src),
                    'vals': vals
                }

                total_vals.append(vals)
                items.append(item)

        np.save(os.path.join(root, exp_name, 'rp-item.npy'), np.array(items))
        np.save(os.path.join(root, exp_name, 'rp-vals.npy'), total_vals)

        vals = np.array(total_vals).mean(0)
        print('BUFF - chamfer: %.4f  p2s: %.4f  nml: %.4f' % (buf_val[0], buf_val[1], buf_val[2]))
        print('RP - chamfer: %.4f  p2s: %.4f  nml: %.4f' % (vals[0], vals[1], vals[2]))

    exp_list = ['pifuhd_final']

    root = args.root
    tar_path = args.tar_path

    for exp in exp_list:
        run(root, exp, tar_path)