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)