Coverage for src/pyroboplan/visualization/meshcat_utils.py: 27%

1"""Utilities for visualization using MeshCat.""" 

2 

3import meshcat.geometry as mg 

4import numpy as np 

5 

6 

7FRAME_AXIS_POSITIONS = ( 

8 np.array([[0, 0, 0], [1, 0, 0], [0, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]]) 

9 .astype(np.float32) 

10 .T 

11) 

12FRAME_AXIS_COLORS = ( 

13 np.array([[1, 0, 0], [1, 0.6, 0], [0, 1, 0], [0.6, 1, 0], [0, 0, 1], [0, 0.6, 1]]) 

14 .astype(np.float32) 

15 .T 

16) 

17 

18 

19def visualize_frame( 

20 visualizer, name, tform, line_length=0.2, line_width=3, line_color=None 

21): 

22 """ 

23 Visualizes a coordinate frame as an axis triad at a specified pose. 

24 

25 Parameters 

26 ---------- 

27 visualizer : `pinocchio.visualize.meshcat_visualizer.MeshcatVisualizer` 

28 The visualizer instance. 

29 name : str 

30 The name of the MeshCat component to add. 

31 tform : `pinocchio.SE3` 

32 The transform at which to display the frame. 

33 line_length : float, optional 

34 The length of the axes in the triad. 

35 line_width : float, optional 

36 The width of the axes in the triad. 

37 line_color : array-like, optional 

38 The line colors to use. If None, chooses default axes colors. 

39 """ 

40 if line_color: 

41 color = np.array([line_color] * 6).T 

42 else: 

43 color = FRAME_AXIS_COLORS 

44 

45 visualizer.viewer[name].set_object( 

46 mg.LineSegments( 

47 mg.PointsGeometry( 

48 position=line_length * FRAME_AXIS_POSITIONS, 

49 color=color, 

50 ), 

51 mg.LineBasicMaterial( 

52 linewidth=line_width, 

53 vertexColors=True, 

54 ), 

55 ) 

56 ) 

57 visualizer.viewer[name].set_transform(tform.homogeneous) 

58 

59 

60def visualize_frames( 

61 visualizer, prefix_name, tforms, line_length=0.2, line_width=3, line_color=None 

62): 

63 """ 

64 Visualizes a set of coordinate frames as axis triads at specified poses. 

65 

66 Parameters 

67 ---------- 

68 visualizer : `pinocchio.visualize.meshcat_visualizer.MeshcatVisualizer` 

69 The visualizer instance. 

70 prefix_name : str 

71 The name of the MeshCat component to add. 

72 tforms : list[`pinocchio.SE3`] 

73 A list of transforms at which to display frames. 

74 line_length : float, optional 

75 The length of the axes in the triad. 

76 line_width : float, optional 

77 The width of the axes in the triad. 

78 line_color : array-like, optional 

79 The line colors to use. If None, chooses default axes colors. 

80 """ 

81 visualizer.viewer[prefix_name].delete() 

82 for idx, tform in enumerate(tforms): 

83 visualize_frame( 

84 visualizer, 

85 f"{prefix_name}/frame{idx}", 

86 tform, 

87 line_length=line_length, 

88 line_width=line_width, 

89 line_color=line_color, 

90 ) 

91 

92 

93def visualize_path(visualizer, name, tforms, line_width=3, line_color=[0.0, 0.0, 0.0]): 

94 """ 

95 Visualizes a path of poses as lines containing only the translation component. 

96 

97 Parameters 

98 ---------- 

99 visualizer : `pinocchio.visualize.meshcat_visualizer.MeshcatVisualizer` 

100 The visualizer instance. 

101 name : str 

102 The name of the MeshCat component to add. 

103 tforms : list[`pinocchio.SE3`] 

104 A list of transforms representing the vertices of the path. 

105 line_width : float, optional 

106 The width of the axes in the triad. 

107 line_color : array-like, optional 

108 The line color to use. 

109 """ 

110 visualize_paths( 

111 visualizer, name, [tforms], line_width=line_width, line_color=line_color 

112 ) 

113 

114 

115def visualize_paths(visualizer, name, paths, line_width=3, line_color=[0.0, 0.0, 0.0]): 

116 """ 

117 Visualizes a list of paths of poses as lines containing only the translation component. 

118 

119 Parameters 

120 ---------- 

121 visualizer : `pinocchio.visualize.meshcat_visualizer.MeshcatVisualizer` 

122 The visualizer instance. 

123 name : str 

124 The name of the MeshCat component to add. 

125 paths : list[list[`pinocchio.SE3`]] 

126 A list of paths, each of which is a list of transforms representing the vertices of the path. 

127 line_width : float, optional 

128 The width of the axes in the triad. 

129 line_color : array-like, optional 

130 The line color to use. 

131 """ 

132 positions = [] 

133 for tforms in paths: 

134 for idx in range(1, len(tforms)): 

135 positions.extend([tforms[idx - 1].translation, tforms[idx].translation]) 

136 

137 visualizer.viewer[name].set_object( 

138 mg.LineSegments( 

139 mg.PointsGeometry( 

140 position=np.array(positions).T, 

141 color=np.array([line_color for _ in positions]).T, 

142 ), 

143 mg.LineBasicMaterial( 

144 linewidth=line_width, 

145 vertexColors=True, 

146 ), 

147 ) 

148 )