metrics/eval_3D_GSO.py

import os
import numpy as np
import trimesh
import open3d as o3d
from metrics import chamfer, compute_iou
from tqdm import tqdm
# seed
np.random.seed(0)

def cartesian_to_spherical(xyz):
    ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
    xy = xyz[:, 0] ** 2 + xyz[:, 1] ** 2
    z = np.sqrt(xy + xyz[:, 2] ** 2)
    theta = np.arctan2(np.sqrt(xy), xyz[:, 2])  # for elevation angle defined from Z-axis down
    # ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
    azimuth = np.arctan2(xyz[:, 1], xyz[:, 0])
    return np.array([theta, azimuth, z])

def get_pose(target_RT):
    R, T = target_RT[:3, :3], target_RT[:, -1]
    T_target = -R.T @ T
    theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
    return theta_target, azimuth_target, z_target

def trimesh_to_open3d(src):
    dst = o3d.geometry.TriangleMesh()
    dst.vertices = o3d.utility.Vector3dVector(src.vertices)
    dst.triangles = o3d.utility.Vector3iVector(src.faces)
    vertex_colors = src.visual.vertex_colors[:, :3].astype(np.float32) / 255.0
    dst.vertex_colors = o3d.utility.Vector3dVector(vertex_colors)
    dst.compute_vertex_normals()

    return dst

def normalize_mesh(vertices):
    max_pt = np.max(vertices, 0)
    min_pt = np.min(vertices, 0)
    scale = 1 / np.max(max_pt - min_pt)
    vertices = vertices * scale

    max_pt = np.max(vertices, 0)
    min_pt = np.min(vertices, 0)
    center = (max_pt + min_pt) / 2
    vertices = vertices - center[None, :]
    return vertices

def capture_screenshots(mesh_rec_o3d, cam_param, render_param, img_name):
    vis = o3d.visualization.Visualizer()
    vis.create_window(width=512, height=512, visible=False)
    vis.add_geometry(mesh_rec_o3d)
    ctr = vis.get_view_control()
    parameters = o3d.io.read_pinhole_camera_parameters(cam_param)
    ctr.convert_from_pinhole_camera_parameters(parameters, allow_arbitrary=True)
    vis.get_render_option().load_from_json(render_param)  # rgb
    vis.poll_events()
    vis.update_renderer()
    vis.capture_screen_image(img_name, do_render=True)
    vis.destroy_window()
    del vis
    del ctr


def vis_3D_rec(GT_DIR, REC_DIR, method_name):
    N = 4096
    # get all folders
    obj_names = [f for f in os.listdir(GT_DIR) if os.path.isdir(os.path.join(GT_DIR, f))]

    CDs = []
    IoUs = []
    for obj_name in tqdm(obj_names):
        print(obj_name)
        gt_meshfile = os.path.join(GT_DIR, obj_name, "meshes", "model.obj")
        if "ours" in REC_DIR:
            condition_pose = np.load(os.path.join(GT_DIR, obj_name, "render_sync_36_single/model/000.npy"))
        else:
            condition_pose = np.load(os.path.join(GT_DIR, obj_name, "render_mvs_25/model/000.npy"))
            condition_pose = np.concatenate([condition_pose, np.array([[0, 0, 0, 1]])], axis=0)
        theta, azimu, radius = get_pose(condition_pose[:3, :])

        if "PointE" in REC_DIR:
            rec_pcfile = os.path.join(REC_DIR, obj_name, "pc.ply")
        if "RealFusion" in REC_DIR:
            rec_meshfile = os.path.join(REC_DIR, obj_name, "mesh/mesh.obj")
        elif "dreamgaussian" in REC_DIR:
            rec_meshfile = os.path.join(REC_DIR, obj_name+".obj")
        elif "Wonder3D" in REC_DIR:
            rec_meshfile = os.path.join(REC_DIR, "mesh-ortho-"+obj_name, "save/it3000-mc192.obj")
        else:
            rec_meshfile = os.path.join(REC_DIR, obj_name, "mesh.ply")



        mesh_gt = trimesh.load(gt_meshfile)
        mesh_gt_o3d = o3d.io.read_triangle_mesh(gt_meshfile, True)

        # trimesh load point cloud
        if "PointE" in REC_DIR:
            pc_rec = trimesh.load(rec_pcfile)

        if method_name == "GT":
            mesh_rec = trimesh.load(gt_meshfile)
            mesh_rec_o3d = o3d.io.read_triangle_mesh(gt_meshfile, True)
        else:
            mesh_rec = trimesh.load(rec_meshfile)
            mesh_rec_o3d = o3d.io.read_triangle_mesh(rec_meshfile, True)

        # normalize
        mesh_gt.vertices = normalize_mesh(mesh_gt.vertices)
        vertices_gt = np.asarray(mesh_gt_o3d.vertices)
        vertices_gt = normalize_mesh(vertices_gt)
        mesh_gt_o3d.vertices = o3d.utility.Vector3dVector(vertices_gt)


        if "PointE" in REC_DIR:
            pc_rec.vertices = normalize_mesh(pc_rec.vertices)

        # normalize
        mesh_rec.vertices = normalize_mesh(mesh_rec.vertices)
        vertices_rec = np.asarray(mesh_rec_o3d.vertices)
        vertices_rec = normalize_mesh(vertices_rec)
        mesh_rec_o3d.vertices = o3d.utility.Vector3dVector(vertices_rec)


        if "RealFusion" in REC_DIR or "Wonder3D_ours" in REC_DIR or "SyncDreamer" in REC_DIR:
            mesh_rec.vertices = trimesh.transformations.rotation_matrix(azimu[0], [0, 0, 1])[:3, :3].dot(
                mesh_rec.vertices.T).T
            # o3d
            R = mesh_rec_o3d.get_rotation_matrix_from_xyz(np.array([0., 0., azimu[0]]))
            mesh_rec_o3d.rotate(R, center=(0, 0, 0))
        elif "dreamgaussian" in REC_DIR:
            mesh_rec.vertices = trimesh.transformations.rotation_matrix(azimu[0]+np.pi/2, [0, 1, 0])[:3, :3].dot(
                mesh_rec.vertices.T).T
            # rotate 90 along x
            mesh_rec.vertices = trimesh.transformations.rotation_matrix(np.pi/2, [1, 0, 0])[:3, :3].dot(
                mesh_rec.vertices.T).T
            # o3d
            R = mesh_rec_o3d.get_rotation_matrix_from_xyz(np.array([0., azimu[0]+np.pi/2, 0.]))
            mesh_rec_o3d.rotate(R, center=(0, 0, 0))
            R = mesh_rec_o3d.get_rotation_matrix_from_xyz(np.array([np.pi/2, 0., 0.]))
            mesh_rec_o3d.rotate(R, center=(0, 0, 0))
        elif "one2345" in REC_DIR:
            # rotate along z axis by azimu degree
            # mesh_rec.apply_transform(trimesh.transformations.rotation_matrix(-azimu, [0, 0, 1]))
            azimu = np.rad2deg(azimu[0])
            azimu += 60 # https://github.com/One-2-3-45/One-2-3-45/issues/26
            # print("azimu", azimu)
            mesh_rec.vertices = trimesh.transformations.rotation_matrix(np.radians(azimu), [0, 0, 1])[:3, :3].dot(mesh_rec.vertices.T).T
            # # scale again
            # mesh_rec, rec_center, rec_scale = normalize_mesh(mesh_rec)
            # o3d
            R = mesh_rec_o3d.get_rotation_matrix_from_xyz(np.array([0., 0., np.radians(azimu)]))
            mesh_rec_o3d.rotate(R, center=(0, 0, 0))
            # # scale again
            # mesh_rec_o3d = mesh_rec_o3d.translate(-rec_center)
            # mesh_rec_o3d = mesh_rec_o3d.scale(1 / rec_scale, [0, 0, 0])
        elif "PointE" in REC_DIR or "ShapeE" in REC_DIR:
            # sample points from rec_pc
            if "PointE" in REC_DIR:
                rec_pc_tri = pc_rec
                rec_pc_tri.vertices = rec_pc_tri.vertices[np.random.choice(np.arange(len(pc_rec.vertices)), N)]
            else:
                rec_pc = trimesh.sample.sample_surface(mesh_rec, N)
                rec_pc_tri = trimesh.PointCloud(vertices=rec_pc[0])

            gt_pc = trimesh.sample.sample_surface(mesh_gt, N)
            gt_pc_tri = trimesh.PointCloud(vertices=gt_pc[0])
            # loop over all flips and 90 degrees rotations of rec_pc, pick the one with the smallest chamfer distance
            chamfer_dist_min = np.inf
            opt_axis = None
            opt_angle = None
            for axis in [[1, 0, 0], [0, 1, 0], [0, 0, 1]]:
                for angle in [0, 90, 180, 270]:
                    tmp_rec_pc_tri = rec_pc_tri.copy()
                    tmp_rec_pc_tri.vertices = trimesh.transformations.rotation_matrix(np.radians(angle), axis)[:3, :3].dot(tmp_rec_pc_tri.vertices.T).T
                    tmp_mesh_rec = mesh_rec.copy()
                    tmp_mesh_rec.vertices = trimesh.transformations.rotation_matrix(np.radians(angle), axis)[:3, :3].dot(tmp_mesh_rec.vertices.T).T
                    # compute chamfer distance
                    chamfer_dist = chamfer(gt_pc_tri.vertices, tmp_rec_pc_tri.vertices)
                    if chamfer_dist < chamfer_dist_min:
                        chamfer_dist_min = chamfer_dist
                        opt_axis = axis
                        opt_angle = angle

            chamfer_dist = chamfer_dist_min

            mesh_rec.vertices = trimesh.transformations.rotation_matrix(np.radians(opt_angle), opt_axis)[:3, :3].dot(mesh_rec.vertices.T).T
            # o3d
            if np.abs(opt_angle) > 1e-6:
                if opt_axis == [1, 0, 0]:
                    R = mesh_rec_o3d.get_rotation_matrix_from_xyz(np.array([np.radians(opt_angle), 0., 0.]))
                elif opt_axis == [0, 1, 0]:
                    R = mesh_rec_o3d.get_rotation_matrix_from_xyz(np.array([0., np.radians(opt_angle), 0.]))
                elif opt_axis == [0, 0, 1]:
                    R = mesh_rec_o3d.get_rotation_matrix_from_xyz(np.array([0., 0., np.radians(opt_angle)]))
                mesh_rec_o3d.rotate(R, center=(0, 0, 0))



        if "ours" in REC_DIR or "SyncDreamer" in REC_DIR:
            # invert the face
            mesh_rec.invert()
            # o3d Invert the mesh faces
            mesh_rec_o3d.triangles = o3d.utility.Vector3iVector(np.asarray(mesh_rec_o3d.triangles)[:, [0, 2, 1]])
            # Compute vertex normals to ensure correct orientation
            mesh_rec_o3d.compute_vertex_normals()




        # normalize
        mesh_rec.vertices = normalize_mesh(mesh_rec.vertices)
        vertices_rec = np.asarray(mesh_rec_o3d.vertices)
        vertices_rec = normalize_mesh(vertices_rec)
        mesh_rec_o3d.vertices = o3d.utility.Vector3dVector(vertices_rec)

        # print("mesh_gt_o3d ", np.asarray(mesh_gt_o3d.vertices).max(0), np.asarray(mesh_gt_o3d.vertices).min(0))
        # print("mesh_rec_o3d ", np.asarray(mesh_rec_o3d.vertices).max(0), np.asarray(mesh_rec_o3d.vertices).min(0))
        assert np.abs(np.asarray(mesh_gt_o3d.vertices)).max() <= 0.505
        assert np.abs(np.asarray(mesh_rec_o3d.vertices)).max() <= 0.505
        assert np.abs(np.asarray(mesh_gt.vertices)).max() <= 0.505
        assert np.abs(np.asarray(mesh_rec.vertices)).max() <= 0.505



        # compute chamfer distance
        chamfer_dist = chamfer(mesh_gt.vertices, mesh_rec.vertices)
        vol_iou = compute_iou(mesh_gt, mesh_rec)
        CDs.append(chamfer_dist)
        IoUs.append(vol_iou)

        # # todo save screenshots
        # mesh_axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0])
        # # draw bbox for gt and rec
        # bbox_gt = mesh_gt.bounding_box.bounds
        # bbox_rec = mesh_rec.bounding_box.bounds
        # bbox_gt_o3d = o3d.geometry.AxisAlignedBoundingBox(min_bound=bbox_gt[0], max_bound=bbox_gt[1])
        # bbox_rec_o3d = o3d.geometry.AxisAlignedBoundingBox(min_bound=bbox_rec[0], max_bound=bbox_rec[1])
        # # color red for gt, green for rec
        # bbox_gt_o3d.color = (1, 0, 0)
        # bbox_rec_o3d.color = (0, 1, 0)
        # # draw a bbox of unit cube [-1, 1]^3
        # bbox_unit_cube = o3d.geometry.AxisAlignedBoundingBox(min_bound=(-1, -1, -1), max_bound=(1, 1, 1))
        # bbox_unit_cube.color = (0, 0, 1)
        #
        # # o3d.visualization.draw_geometries(
        # #     [mesh_axis, mesh_gt_o3d, mesh_rec_o3d, bbox_gt_o3d, bbox_rec_o3d, bbox_unit_cube])
        #
        # # take a screenshot with circle view and save to file
        # # save screenshot to file
        # vis_output = os.path.join("screenshots", method_name)
        # os.makedirs(vis_output, exist_ok=True)
        # mesh_rec_o3d.compute_vertex_normals()
        #
        # # vis = o3d.visualization.Visualizer()
        # # vis.create_window(width=512, height=512)
        # # vis.add_geometry(mesh_rec_o3d)
        # # # show the window and save camera pose to json file
        # # vis.get_render_option().light_on = True
        # # vis.run()
        #
        # # rgb
        # for i in range(6):
        #     capture_screenshots(mesh_rec_o3d, f"ScreenCamera_{i}.json", "RenderOption_rgb.json", os.path.join(vis_output, obj_name + f"_{i}.png"))
        # # phong shading
        # for i in range(6):
        #     capture_screenshots(mesh_rec_o3d, f"ScreenCamera_{i}.json", "RenderOption_phong.json", os.path.join(vis_output, obj_name + f"_{i}_phong.png"))


        # todo 3D metrics
        # save metrics to a single file
        with open(os.path.join(REC_DIR, "metrics3D.txt"), "a") as f:
            # write metrics in one line with format: obj_name chamfer_dist volume_iou
            f.write(obj_name + " CD:" + str(chamfer_dist) + " IoU:" + str(vol_iou) + "\n")

    # average metrics and save to the file
    print("Average CD:", np.mean(CDs))
    print("Average IoU:", np.mean(IoUs))
    with open(os.path.join(REC_DIR, "metrics3D.txt"), "a") as f:
        f.write("Average CD:" + str(np.mean(CDs)) + " IoU:" + str(np.mean(IoUs)) + "\n")


### TODO
GT_DIR = "/home/xin/data/EscherNet/Data/GSO30/"
methods = {}
# methods["One2345-XL"] = ""
# methods["One2345"] = ""
# methods["PointE"] = ""
# methods["ShapeE"] = ""
# methods["DreamGaussian"] = ""
# methods["DreamGaussian-XL"] = ""
# methods["SyncDreamer"] = ""
methods["Ours_T1"] = "/GSO3D/ours_GSO_T1/NeuS/"
methods["Ours_T2"] = "/GSO3D/ours_GSO_T2/NeuS/"
methods["Ours_T3"] = "/GSO3D/ours_GSO_T3/NeuS/"
methods["Ours_T5"] = "/GSO3D/ours_GSO_T5/NeuS/"
methods["Ours_T10"] = "/GSO3D/ours_GSO_T10/NeuS"

for method_name in methods.keys():
    print("method_name: ", method_name)
    vis_3D_rec(GT_DIR, methods[method_name], method_name)