98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
|
|
import numpy as np
|
|
import open3d as o3d
|
|
|
|
from preprocessing.plane import PlaneResult
|
|
|
|
|
|
class Visualizer:
|
|
"""Open3D visualization utilities."""
|
|
|
|
def show_registration(
|
|
self,
|
|
source: o3d.geometry.PointCloud,
|
|
target: o3d.geometry.PointCloud,
|
|
recovered: o3d.geometry.PointCloud,
|
|
) -> None:
|
|
source_colored = copy.deepcopy(source)
|
|
target_colored = copy.deepcopy(target)
|
|
recovered_colored = copy.deepcopy(recovered)
|
|
source_colored.paint_uniform_color([0, 0, 1])
|
|
target_colored.paint_uniform_color([0, 1, 0])
|
|
recovered_colored.paint_uniform_color([1, 0, 0])
|
|
|
|
o3d.visualization.draw_geometries(
|
|
[source_colored, target_colored],
|
|
window_name="Before Registration (Blue: Source, Green: Target)",
|
|
width=1024,
|
|
height=768,
|
|
)
|
|
o3d.visualization.draw_geometries(
|
|
[recovered_colored, target_colored],
|
|
window_name="After Registration (Red: Recovered, Green: Target)",
|
|
width=1024,
|
|
height=768,
|
|
)
|
|
|
|
def show_denoise(
|
|
self,
|
|
clean: o3d.geometry.PointCloud,
|
|
noisy: o3d.geometry.PointCloud,
|
|
filtered: o3d.geometry.PointCloud | None = None,
|
|
filter_name: str = "Filtered",
|
|
) -> None:
|
|
clean_colored = copy.deepcopy(clean)
|
|
noisy_colored = copy.deepcopy(noisy)
|
|
clean_colored.paint_uniform_color([0, 0, 1])
|
|
noisy_colored.paint_uniform_color([1.0, 0.7, 0.5])
|
|
o3d.visualization.draw_geometries(
|
|
[clean_colored, noisy_colored],
|
|
window_name=f"Clean (Blue, {len(clean.points)}) vs Noisy (Orange, {len(noisy.points)})",
|
|
)
|
|
|
|
if filtered is not None:
|
|
filtered_colored = copy.deepcopy(filtered)
|
|
filtered_colored.paint_uniform_color([0, 1, 0])
|
|
o3d.visualization.draw_geometries(
|
|
[clean_colored, filtered_colored],
|
|
window_name=f"{filter_name}: Clean (Blue) vs Filtered (Green)",
|
|
)
|
|
|
|
def show_curvature(self, pointcloud: o3d.geometry.PointCloud, window_name: str = "Curvature Visualization") -> None:
|
|
o3d.visualization.draw_geometries([pointcloud], window_name=window_name)
|
|
|
|
def show_plane(self, plane_result: PlaneResult) -> None:
|
|
pointcloud = copy.deepcopy(plane_result.pointcloud)
|
|
pointcloud.paint_uniform_color([0.2, 0.4, 1.0])
|
|
inlier_cloud = pointcloud.select_by_index(plane_result.inliers)
|
|
inlier_cloud.paint_uniform_color([1.0, 0.2, 0.2])
|
|
|
|
plane_mesh = self._create_plane_mesh(plane_result.model)
|
|
o3d.visualization.draw_geometries(
|
|
[pointcloud, inlier_cloud, plane_mesh],
|
|
window_name="Plane Fitting",
|
|
width=1024,
|
|
height=768,
|
|
)
|
|
|
|
def _create_plane_mesh(self, plane_model: tuple[float, float, float, float]) -> o3d.geometry.TriangleMesh:
|
|
a, b, c, d = plane_model
|
|
size = 1.0
|
|
xx, yy = np.meshgrid(np.linspace(-size, size, 2), np.linspace(-size, size, 2))
|
|
if abs(c) < 1e-12:
|
|
zz = np.zeros_like(xx)
|
|
else:
|
|
zz = -(a * xx + b * yy + d) / c
|
|
|
|
vertices = np.column_stack([xx.ravel(), yy.ravel(), zz.ravel()])
|
|
triangles = np.array([[0, 1, 2], [1, 3, 2]])
|
|
mesh = o3d.geometry.TriangleMesh()
|
|
mesh.vertices = o3d.utility.Vector3dVector(vertices)
|
|
mesh.triangles = o3d.utility.Vector3iVector(triangles)
|
|
mesh.paint_uniform_color([0.1, 0.9, 0.1])
|
|
mesh.compute_vertex_normals()
|
|
return mesh
|