SimplerEnv/visualizer.py at fsd · hilookas/SimplerEnv · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import open3d as o3d
import numpy as np
from pytransform3d import transformations as pt
from pytransform3d import rotations as pr
import pickle as pkl

def create_coordinate_frame(translation, rotation, scale=0.1):
    # Create coordinate frame mesh
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=scale)

    # Transform frame to desired pose
    frame.translate(translation)
    frame.rotate(rotation, center=translation)

    return frame

# Example usage
if __name__ == "__main__":
    # Replace with your point cloud file path
    pointcloud_file = "pcd.ply"  # or .xyz, .pcd, etc.

    # Example EEF poses (translation + euler angles)
    eef_poses_euler = [
        np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),  # x,y,z, rx,ry,rz
        # np.array([0.5, 0.0, 0.0, 0.0, 0.0, 0.0]),  # x,y,z, rx,ry,rz
        # np.array([0.0, 0.5, 0.0, np.pi/2, 0.0, 0.0])
    ]

    eef_poses = []
    for pose in eef_poses_euler:
        translation = pose[:3]
        rotation = pose[3:]  # Assuming last 3 values are euler angles
        rotation = pr.matrix_from_euler(rotation, 0, 1, 2, True)
        eef_poses.append((translation, rotation))

    pick_translation, pick_rotation = pkl.load(open("pick.pkl", "rb"))
    eef_poses.append((pick_translation, pick_rotation))

    place_translation, place_rotation = pkl.load(open("place.pkl", "rb"))
    eef_poses.append((place_translation, place_rotation))

    # Load point cloud from file
    pcd = o3d.io.read_point_cloud(pointcloud_file)

    # Create visualization window
    vis = o3d.visualization.Visualizer()
    vis.create_window()

    # Add point cloud to visualizer
    vis.add_geometry(pcd)

    # Add coordinate frames for each EEF pose if provided
    for translation, rotation in eef_poses:
        frame = create_coordinate_frame(translation, rotation)
        vis.add_geometry(frame)

    # Set default view
    vis.get_render_option().point_size = 1
    vis.get_render_option().background_color = np.asarray([0, 0, 0])


    # Run visualizer
    vis.run()
    vis.destroy_window()

    # To save a point cloud as a PLY file:
    # 1. Create an Open3D point cloud object
    # 2. Use o3d.io.write_point_cloud("pcd.ply", pcd)
    # Example:
    # pcd = o3d.geometry.PointCloud()
    # pcd.points = o3d.utility.Vector3dVector(points)
    # o3d.io.write_point_cloud("pcd.ply", pcd)