File size: 1,782 Bytes
68720c1 |
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 |
import numpy as np
from matplotlib import colormaps
import rerun as rr
from rerun.components import Material
from scipy.spatial import transform
def color_fn(x, cmap="tab10"):
return colormaps[cmap](x % colormaps[cmap].N)
def log_sample(
root_name: str,
traj: np.ndarray,
char_traj: np.ndarray,
K: np.ndarray,
vertices: np.ndarray,
faces: np.ndarray,
normals: np.ndarray,
caption: str,
mesh_masks: np.ndarray,
):
num_cameras = traj.shape[0]
rr.log(root_name, rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, timeless=True)
rr.log(
f"{root_name}/trajectory/points",
rr.Points3D(traj[:, :3, 3]),
timeless=True,
)
rr.log(
f"{root_name}/trajectory/line",
rr.LineStrips3D(
np.stack((traj[:, :3, 3][:-1], traj[:, :3, 3][1:]), axis=1),
colors=[(1.0, 0.0, 1.0, 1.0)],
),
timeless=True,
)
for k in range(num_cameras):
rr.set_time_sequence("frame_idx", k)
translation = traj[k][:3, 3]
rotation_q = transform.Rotation.from_matrix(traj[k][:3, :3]).as_quat()
rr.log(
f"{root_name}/camera/image",
rr.Pinhole(
image_from_camera=K,
width=K[0, -1] * 2,
height=K[1, -1] * 2,
),
)
rr.log(
f"{root_name}/camera",
rr.Transform3D(
translation=translation,
rotation=rr.Quaternion(xyzw=rotation_q),
),
)
rr.set_time_sequence("image", k)
# Log character trajectory points
rr.log(
f"{root_name}/char_traj/points",
rr.Points3D(char_traj.T, colors=[(1.0, 0.0, 0.0, 1.0)]),
timeless=True,
)
|