-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[visualization] Add frame (aka triad) visualization function
- Loading branch information
1 parent
3f9800b
commit e112f3f
Showing
5 changed files
with
196 additions
and
68 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,79 @@ | ||
| import numpy as np | ||
|
|
||
| from pydrake.common.eigen_geometry import AngleAxis | ||
| from pydrake.geometry import ( | ||
| Cylinder, | ||
| FrameId, | ||
| GeometryInstance, | ||
| MakePhongIllustrationProperties, | ||
| SceneGraph, | ||
| ) | ||
| from pydrake.math import RigidTransform | ||
| from pydrake.multibody.plant import MultibodyPlant | ||
| from pydrake.multibody.tree import Body | ||
|
|
||
|
|
||
| def AddFrameTriadIllustration( | ||
| *, | ||
| scene_graph: SceneGraph, | ||
| body: Body = None, | ||
| frame_id: FrameId = None, | ||
| plant: MultibodyPlant = None, | ||
| name: str = "frame", | ||
| length: float = 0.3, | ||
| radius: float = 0.005, | ||
| opacity: float = 0.9, | ||
| X_FT: RigidTransform = RigidTransform(), | ||
| ): | ||
| """ | ||
| Adds illustration geometry representing the given frame using an RGB triad, | ||
| with the x-axis drawn in red, the y-axis in green and the z-axis in blue. | ||
| The given frame can be either a geometry FrameId or a multibody Body. | ||
| Args: | ||
| scene_graph: the SceneGraph where geometry will be added. | ||
| body: when provided, illustrates the frame of the given body; | ||
| either body= or frame_id= must be provided, but not both. | ||
| frame_id: when provided, illustrates the given geometry.FrameId | ||
| registered with the given plant and scene_graph; | ||
| either body= or frame_id= must be provided, but not both. | ||
| plant: MultibodyPlant associated with the given frame_id; | ||
| required if and only if a frame_id= is being used (not body=). | ||
| name: the added geometries will have names "{name} x-axis", etc. | ||
| length: the length of each axis in meters. | ||
| radius: the radius of each axis in meters. | ||
| opacity: the opacity each axis, between 0.0 and 1.0. | ||
| X_FT: rigid transform relating frame T (the triad geometry) and | ||
| frame F (the given frame_id= or body= frame). | ||
| Returns: | ||
| The newly-added geometry ids for (x, y, z) respectively. | ||
| """ | ||
| if sum([body is not None, frame_id is not None]) != 1: | ||
| raise ValueError("Must provide either body= xor frame_id=") | ||
| if body is not None: | ||
| if plant is not None: | ||
| if plant is not body.GetParentPlant(): | ||
| raise ValueError( | ||
| "Mismatched body= and plant=; remove the plant= arg") | ||
| else: | ||
| plant = body.GetParentPlant() | ||
| frame_id = plant.GetBodyFrameIdOrThrow(body.index()) | ||
| source_id = plant.get_source_id() | ||
| eye = np.eye(3) | ||
| result = [] | ||
| for i, char in enumerate(("x", "y", "z")): | ||
| geom_name = f"{name} {char}-axis" | ||
| # p_TG centers the cylinder halfway along the i'th axis. | ||
| p_TG = 0.5 * length * eye[i] | ||
| # R_TG rotates the canonical cylinder (aligned with +z) to align with | ||
| # the i'th axis instead. When i == 2, it spins the cylinder around the | ||
| # z axis, but this is effectively a no-op. | ||
| R_TG = AngleAxis(angle=np.pi/2, axis=eye[1-i]) | ||
| X_FG = X_FT @ RigidTransform(R_TG, p_TG) | ||
| geom = GeometryInstance(X_FG, Cylinder(radius, length), geom_name) | ||
| phong = MakePhongIllustrationProperties(np.append(eye[i], [opacity])) | ||
| geom.set_illustration_properties(phong) | ||
| geometry_id = scene_graph.RegisterGeometry(source_id, frame_id, geom) | ||
| result.append(geometry_id) | ||
| return tuple(result) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,94 @@ | ||
| import unittest | ||
|
|
||
| import numpy as np | ||
|
|
||
| from pydrake.geometry import Meshcat, Rgba | ||
| from pydrake.math import RigidTransform, RotationMatrix | ||
| from pydrake.multibody.plant import MultibodyPlant | ||
| from pydrake.planning import RobotDiagramBuilder | ||
| import pydrake.visualization as mut | ||
|
|
||
|
|
||
| class TestTriad(unittest.TestCase): | ||
|
|
||
| def setUp(self): | ||
| # Create a simple robot. | ||
| builder = RobotDiagramBuilder() | ||
| builder.parser().AddModels( | ||
| url="package://drake/multibody/benchmarks/acrobot/acrobot.sdf") | ||
| robot = builder.Build() | ||
| self._plant = robot.plant() | ||
| self._scene_graph = robot.scene_graph() | ||
|
|
||
| def test_via_body(self): | ||
| # Illustrate the Link2 body frame. | ||
| x, y, z = mut.AddFrameTriadIllustration( | ||
| scene_graph=self._scene_graph, | ||
| body=self._plant.GetBodyByName("Link2"), | ||
| name="foo", | ||
| length=0.2, | ||
| radius=0.001, | ||
| opacity=0.5, | ||
| X_FT=RigidTransform()) | ||
|
|
||
| # Check the geometry pose of each axis. | ||
| inspect = self._scene_graph.model_inspector() | ||
| np.testing.assert_equal( | ||
| inspect.GetPoseInFrame(x).GetAsMatrix34(), | ||
| RigidTransform( | ||
| RotationMatrix.MakeYRotation(np.pi/2), | ||
| [0.1, 0, 0]).GetAsMatrix34()) | ||
| np.testing.assert_equal( | ||
| inspect.GetPoseInFrame(y).GetAsMatrix34(), | ||
| RigidTransform( | ||
| RotationMatrix.MakeXRotation(np.pi/2), | ||
| [0, 0.1, 0]).GetAsMatrix34()) | ||
| np.testing.assert_equal( | ||
| inspect.GetPoseInFrame(z).GetAsMatrix34(), | ||
| RigidTransform( | ||
| RotationMatrix.MakeZRotation(np.pi/2), | ||
| [0, 0, 0.1]).GetAsMatrix34()) | ||
|
|
||
| # Check the geometry details of each axis. | ||
| for i, char in enumerate(("x", "y", "z")): | ||
| geom_id = [x, y, z][i] | ||
| frame_name = inspect.GetName(inspect.GetFrameId(geom_id)) | ||
| self.assertEqual(frame_name, "acrobot::Link2") | ||
| self.assertEqual(inspect.GetName(geom_id), f"foo {char}-axis") | ||
| self.assertEqual(inspect.GetShape(geom_id).length(), 0.2) | ||
| self.assertEqual(inspect.GetShape(geom_id).radius(), 0.001) | ||
| props = inspect.GetIllustrationProperties(geom_id) | ||
| self.assertEqual(props.GetProperty("phong", "diffuse"), | ||
| Rgba(r=(i == 0), g=(i == 1), b=(i == 2), a=0.5)) | ||
|
|
||
| def test_via_frame_id(self): | ||
| # Illustrate the Link2 geometry frame. | ||
| body = self._plant.GetBodyByName("Link2") | ||
| expected_frame = self._plant.GetBodyFrameIdIfExists(body.index()) | ||
| geom_ids = mut.AddFrameTriadIllustration( | ||
| plant=self._plant, | ||
| scene_graph=self._scene_graph, | ||
| frame_id=expected_frame) | ||
|
|
||
| # Check that all three geoms were added. | ||
| self.assertEqual(len(geom_ids), 3) | ||
| inspect = self._scene_graph.model_inspector() | ||
| for geom_id in geom_ids: | ||
| actual_frame = inspect.GetFrameId(geom_id) | ||
| self.assertEqual(actual_frame, expected_frame) | ||
|
|
||
| def test_redundant_plant_arg(self): | ||
| # Pass a plant= even though it's not needed. | ||
| # It's sliently cross-checked and accepted. | ||
| mut.AddFrameTriadIllustration( | ||
| plant=self._plant, | ||
| scene_graph=self._scene_graph, | ||
| body=self._plant.GetBodyByName("Link2")) | ||
|
|
||
| def test_wrong_plant_arg(self): | ||
| # Pass the wrong plant= value and get rejected. | ||
| with self.assertRaisesRegex(ValueError, "Mismatched"): | ||
| mut.AddFrameTriadIllustration( | ||
| plant=MultibodyPlant(0.0), | ||
| scene_graph=self._scene_graph, | ||
| body=self._plant.GetBodyByName("Link2")) |