65 lines
2.8 KiB
Python
65 lines
2.8 KiB
Python
import numpy as np
|
|
import open3d as o3d
|
|
from pts import PtsUtil
|
|
from scipy.spatial import cKDTree
|
|
|
|
class ReconstructionUtil:
|
|
@staticmethod
|
|
def reconstruct_with_pts(pts):
|
|
pcd = o3d.geometry.PointCloud()
|
|
pcd.points = o3d.utility.Vector3dVector(pts)
|
|
pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
|
|
mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=9)
|
|
densities = np.asarray(densities)
|
|
vertices_to_remove = densities < np.quantile(densities, 0.2)
|
|
mesh.remove_vertices_by_mask(vertices_to_remove)
|
|
return mesh
|
|
|
|
@staticmethod
|
|
def filter_points(points, points_normals, cam_pose, voxel_size=0.005, theta=45):
|
|
sampled_points = PtsUtil.voxel_downsample_point_cloud(points, voxel_size)
|
|
#sampled_points = points
|
|
kdtree = cKDTree(points_normals[:,:3])
|
|
_, indices = kdtree.query(sampled_points)
|
|
nearest_points = points_normals[indices]
|
|
|
|
normals = nearest_points[:, 3:]
|
|
camera_axis = -cam_pose[:3, 2]
|
|
normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
|
|
cos_theta = np.dot(normals_normalized, camera_axis)
|
|
theta_rad = np.deg2rad(theta)
|
|
filtered_sampled_points= sampled_points[cos_theta > np.cos(theta_rad)]
|
|
|
|
return filtered_sampled_points[:, :3]
|
|
|
|
@staticmethod
|
|
def get_new_added_points(old_combined_pts, new_pts, threshold=0.005):
|
|
if old_combined_pts.size == 0:
|
|
return new_pts
|
|
if new_pts.size == 0:
|
|
return np.array([])
|
|
|
|
tree = cKDTree(old_combined_pts)
|
|
distances, _ = tree.query(new_pts, k=1)
|
|
new_added_points = new_pts[distances > threshold]
|
|
return new_added_points
|
|
|
|
if __name__ == "__main__":
|
|
import os
|
|
root = "/media/hofee/data/project/python/nbv_reconstruction/nbv_rec_visualize/data/sample/"
|
|
name = "google_scan-box_0106"
|
|
model_path = os.path.join(root, name, "sampled_model_points.txt")
|
|
rec_path = os.path.join(root, name, "best_reconstructed_pts.txt")
|
|
model_pts = np.loadtxt(model_path)
|
|
rec_pts = np.loadtxt(rec_path)
|
|
import time
|
|
start = time.time()
|
|
model_mesh = ReconstructionUtil.reconstruct_with_pts(model_pts)
|
|
end = time.time()
|
|
print(f"Time taken to reconstruct model: {end-start}")
|
|
rec_mesh = ReconstructionUtil.reconstruct_with_pts(rec_pts)
|
|
output_dir = "/media/hofee/data/project/python/nbv_reconstruction/nbv_rec_visualize/mis/output_test"
|
|
model_output_path = os.path.join(output_dir, f"{name}_model_mesh.obj")
|
|
rec_output_path = os.path.join(output_dir, f"{name}_rec_mesh.obj")
|
|
o3d.io.write_triangle_mesh(model_output_path, model_mesh)
|
|
o3d.io.write_triangle_mesh(rec_output_path, rec_mesh) |