check_quaternion_from_vec2vec.py

import contrib

import numpy as np
import pickpp
import trimesh


env = contrib.RlightPickAndPlaceEnv(version="v3", num_instances=1).env
obs = env.reset()

mask = obs["mask"] == obs["objects"][0, 0]

pcd = pickpp.geometry.pointcloud_from_depth(obs["depth"], *obs["intrinsic"])
pcd_masked = pcd[mask]

pos, qua = obs["extrinsic"][:3], obs["extrinsic"][3:]
T_cam2world = pickpp.geometry.transformation_matrix(qua, pos)
pcd_masked = pickpp.geometry.transform_points(pcd_masked, T_cam2world)

pos, qua = obs["robot"][:3], obs["robot"][3:]
T_ee2world = pickpp.geometry.transformation_matrix(qua, pos)
pcd_masked = pickpp.geometry.transform_points(
    pcd_masked, np.linalg.inv(T_ee2world)
)

normals = pickpp.geometry.normals_from_pointcloud(pcd_masked)

geom = trimesh.PointCloud(vertices=pcd_masked)
scene = trimesh.Scene()
scene.add_geometry(geom)

for point, normal in zip(pcd_masked, normals):
    path = trimesh.load_path([point, point + normal * 0.02])
    scene.add_geometry(path)
del path, point, normal

for index in range(len(pcd_masked)):
    if index % 10 != 0:
        continue
    position = pcd_masked[index]
    normal = normals[index]

    # ee coord
    quaternion = pickpp.geometry.quaternion_from_vec2vec([0, 0, 1], normal)
    coord = pickpp.geometry.Coordinate(
        position=position, quaternion=quaternion
    )

    axis = trimesh.creation.axis(
        origin_size=0, axis_length=0.01, axis_radius=0.001
    )
    axis.apply_transform(coord.matrix)
    scene.add_geometry(axis)

scene.show()