tools/visualization.py

import numpy as np
import open3d as o3d
import plotly.graph_objects as go
import copy
import os
from pathlib import Path
from datetime import datetime

def get_color(deg):
    '''
    This function is used to determine the color of the arrow that shows angle error.
    It gets color based on the degree of rotation.
    '''
    deg = abs(deg)
    if deg < 5:
        return 'green'
    elif deg < 10:
        return 'orange'
    else:
        return 'red'
        
def plot_point_cloud(source, target, result = None, show_grid=False, x_diff=None, y_diff=None, z_diff=None):
    '''
    Visualizes two point clouds and optionally a result point cloud.

    args:
        source: o3d.geometry.PointCloud, the source point cloud to visualize 
        target: o3d.geometry.PointCloud, the target point cloud to visualize
        result: o3d.geometry.PointCloud, optional, the result point cloud to visualize
        show_grid: bool, optional, whether to show the grid in the 3D plot
        x_diff: float, optional, the difference in X angle to visualize
        y_diff: float, optional, the difference in Y angle to visualize
        z_diff: float, optional, the difference in Z angle to visualize

    result:
        A 3D plotly figure visualizing the source and target point clouds, and optionally
        the result point cloud, with arrows indicating angle differences.
    '''
    # Point cloud data
    source_points = np.asarray(source.points)
    target_points = np.asarray(target.points)

    if result is not None:
        result_points = np.asarray(result.points)

    source_scatter = go.Scatter3d(
        x=source_points[:, 0], y=source_points[:, 1], z=source_points[:, 2],
        mode='markers', marker=dict(size=2, color='rgb(255, 180, 0)'), name='Source (Yellow)')

    target_scatter = go.Scatter3d(
        x=target_points[:, 0], y=target_points[:, 1], z=target_points[:, 2],
        mode='markers', marker=dict(size=2, color='rgb(0, 166, 237)'), name='Target (Blue)')

    data = [source_scatter, target_scatter]

    if result is not None:
        result_scatter = go.Scatter3d(
            x=result_points[:, 0], y=result_points[:, 1], z=result_points[:, 2],
        # mode='markers', marker=dict(size=2, color='rgb(153, 0, 76)'), name='Result (Purple)')
        mode='markers', marker=dict(size=2, color='rgb(255, 180, 0)'), name='Result (Yellow)')
        data.append(result_scatter)

    # Create the figure
    fig = go.Figure(data=data)

    # Plots arrows to show angle difference 
    if x_diff and y_diff and z_diff:
        all_points = [source_points, target_points]
        if result is not None:
            all_points.append(result_points)

        center = np.mean(np.vstack(all_points), axis=0)
        scale_factor = 0.01  # Adjust this for visual scaling

        angles = [x_diff, y_diff, z_diff]
        axes = [
            {"axis": "X", "vec": np.array([1, 0, 0]), "angle": angles[0], "label": f'ΔX (Roll): {angles[0]:.2f}°', "color": get_color(angles[0])},
            {"axis": "Y", "vec": np.array([0, 1, 0]), "angle": angles[1], "label": f'ΔY (Pitch): {angles[1]:.2f}°', "color": get_color(angles[1])},
            {"axis": "Z", "vec": np.array([0, 0, 1]), "angle": angles[2], "label": f'ΔZ (Yaw): {angles[2]:.2f}°', "color": get_color(angles[2])},
        ]

        for axis in axes:
            magnitude = abs(axis["angle"]) * scale_factor
            vec = axis["vec"] * axis["angle"] * scale_factor
            tip = center + vec
            color = axis["color"]

            # Line (arrow)
            fig.add_trace(go.Scatter3d(
                x=[center[0], tip[0]], y=[center[1], tip[1]], z=[center[2], tip[2]],
                mode='lines',
                line=dict(color=color, width=5),
                showlegend=False
            ))

            # Arrow Head (cone)
            fig.add_trace(go.Cone(
                x=[tip[0]], y=[tip[1]], z=[tip[2]],
                u=[vec[0]], v=[vec[1]], w=[vec[2]],
                colorscale=[[0, color], [1, color]],
                showscale=False,
                sizemode="absolute",
                sizeref=0.04 * magnitude, 
                anchor="tip"
            ))

            # Annotation
            text_pos = tip + 0.02 * np.sign(vec)
            fig.add_trace(go.Scatter3d(
                x=[text_pos[0]], y=[text_pos[1]], z=[text_pos[2]],
                mode='text',
                text=[axis["label"]],
                textposition="top center",
                showlegend=False,
                textfont=dict(size=14, color=color)
            ))

    fig.update_layout(
        scene=dict(
            xaxis=dict(visible=show_grid),
            yaxis=dict(visible=show_grid),
            zaxis=dict(visible=show_grid),
            aspectmode='data'
        ),
        width=900,  
        height=700  
    )

    fig.show()

# Open3d helper function to draw 2 point clouds
def draw_point_clouds(pcd1, pcd2):
    '''
    args:
        pcd1: o3d.geometry.PointCloud
        pcd2: o3d.geometry.PointCloud
    
    result: 
        Visualizes pcd2 with yellow and pcd1 with cyan ransformed with an alignment transformation. 
    '''
    pcd1_temp = copy.deepcopy(pcd1)
    pcd2_temp = copy.deepcopy(pcd2)
    pcd1_temp.paint_uniform_color([1, 0.706, 0])
    pcd2_temp.paint_uniform_color([0, 0.651, 0.929])
    o3d.visualization.draw_geometries([pcd1_temp, pcd2_temp],
                                      zoom=0.4459,
                                      front=[0.9288, -0.2951, -0.2242],
                                      lookat=[1.6784, 2.0612, 1.4451],
                                      up=[-0.3402, -0.9189, -0.1996])
    
# open3d helper function to draw registration result
def draw_registration_result(source, target, transformation, method_name = None):
    '''
    args:
        source: o3d.geometry.PointCloud
        target: o3d.geometry.PointCloud
        transformation: np.array, 4x4 matrix (an intial guess of the transformation to roughly align PCs on top of each other) --> Global registration
    
    result: 
        Saves source, target, and transformed source point clouds as .ply files (for later viewing).
        If a display is available, also visualizes the registration result.
    '''
    # Make copies so we never mutate the inputs
    source_orig = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp = copy.deepcopy(source)

    # Color for easier later inspection
    # source_orig.paint_uniform_color([1, 0.706, 0])
    # source_temp.paint_uniform_color([1, 0.706, 0])
    # target_temp.paint_uniform_color([0, 0.651, 0.929])

    # Apply transformation to the "result" copy
    source_temp.transform(transformation)

    # Save point clouds for offline visualization
    out_dir = Path(__file__).resolve().parents[1] / "results" / "registration_plys"
    out_dir.mkdir(parents=True, exist_ok=True)

    stamp = datetime.now().strftime("%d_%H%M_%f")
    source_path = out_dir / f"{stamp}_source.ply"
    target_path = out_dir / f"{stamp}_target.ply"
    transformed_path = out_dir / f"{stamp}_source_transformed.ply"

    try:
        if method_name is not None:
            source_path = out_dir / f"{stamp}_{method_name}_source.ply"
            target_path = out_dir / f"{stamp}_{method_name}_target.ply"
            transformed_path = out_dir / f"{stamp}_{method_name}_source_transformed.ply"

        o3d.io.write_point_cloud(str(source_path), source_orig)
        o3d.io.write_point_cloud(str(target_path), target_temp)
        o3d.io.write_point_cloud(str(transformed_path), source_temp)
        print(f"[visualization] Saved .ply files to: {out_dir}")
    except Exception as e:
        print(f"[visualization] WARNING: Failed to write .ply files to {out_dir}: {e}")

    # Only try to open a window if a display exists (avoid headless GLFW warnings)
    if os.environ.get("DISPLAY"):
        o3d.visualization.draw_geometries([source_temp, target_temp],
                                          zoom=0.4459,
                                          front=[0.9288, -0.2951, -0.2242],
                                          lookat=[1.6784, 2.0612, 1.4451],
                                          up=[-0.3402, -0.9189, -0.1996])