This shows examples of:
MultibodyPlant
and SceneGraph
to a diagramMultibodyPlant
Meshcat
visualizationSceneGraphInspector
to query SceneGraph
geometriesSceneGraph
geometries with MultibodyPlant
bodiesRenderLabel
s from given geometriesModelInstanceIndex
.import os
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from pydrake.common import FindResourceOrThrow
from pydrake.geometry import MeshcatVisualizer, StartMeshcat
from pydrake.geometry.render import (
ClippingRange,
ColorRenderCamera,
DepthRange,
DepthRenderCamera,
RenderCameraCore,
RenderLabel,
MakeRenderEngineVtk,
RenderEngineVtkParams,
)
from pydrake.math import RigidTransform, RollPitchYaw
from pydrake.multibody.parsing import Parser
from pydrake.multibody.plant import AddMultibodyPlantSceneGraph
from pydrake.multibody.tree import BodyIndex
from pydrake.systems.analysis import Simulator
from pydrake.systems.framework import DiagramBuilder
from pydrake.systems.sensors import (
CameraInfo,
RgbdSensor,
)
We'll also start a Meshcat
instance here. It is typical to do this once at the beginning of the notebook, and use the same instance throughout. You should open a second browser window to the URL that is displayed in the cell output.
meshcat = StartMeshcat()
def xyz_rpy_deg(xyz, rpy_deg):
"""Shorthand for defining a pose."""
rpy_deg = np.asarray(rpy_deg)
return RigidTransform(RollPitchYaw(rpy_deg * np.pi / 180), xyz)
reserved_labels = [
RenderLabel.kDoNotRender,
RenderLabel.kDontCare,
RenderLabel.kEmpty,
RenderLabel.kUnspecified,
]
def colorize_labels(image):
"""Colorizes labels."""
# TODO(eric.cousineau): Revive and use Kuni's palette.
cc = mpl.colors.ColorConverter()
color_cycle = plt.rcParams["axes.prop_cycle"]
colors = np.array([cc.to_rgb(c["color"]) for c in color_cycle])
bg_color = [0, 0, 0]
image = np.squeeze(image)
background = np.zeros(image.shape[:2], dtype=bool)
for label in reserved_labels:
background |= image == int(label)
color_image = colors[image % len(colors)]
color_image[background] = bg_color
return color_image
def colorize_depth(depth, invalid_color=(100, 0, 0)):
"""
Colorizes depth (in meters) as an RGB image.
In general, it'd be better to use a matplotlib colormap, but
this simplifies handling for invalid depth values.
"""
assert depth.dtype == np.float32
invalid = (depth <= 0) | ~np.isfinite(depth)
scale_min = np.min(depth[~invalid])
scale_max = np.max(depth[~invalid])
# Normalize.
depth_scaled = (depth - scale_min) / (scale_max - scale_min)
h, w = depth_scaled.shape
# Convert to rgb.
depth_rgb = np.tile(depth_scaled.reshape((h, w, 1)), (1, 1, 3))
# Invert coloring (white is up front).
depth_rgb = 1 - depth_rgb
# Recolor.
depth_rgb = (depth_rgb * 255).astype(np.uint8)
depth_rgb[invalid] = np.array(invalid_color)
return depth_rgb
builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder, 0.0)
parser = Parser(plant)
iiwa_file = FindResourceOrThrow(
"drake/manipulation/models/iiwa_description/sdf/"
"iiwa14_no_collision.sdf"
)
iiwa_1 = parser.AddModelFromFile(iiwa_file, model_name="iiwa_1")
plant.WeldFrames(
frame_on_parent_F=plant.world_frame(),
frame_on_child_M=plant.GetFrameByName("iiwa_link_0", iiwa_1),
X_FM=xyz_rpy_deg([0, -0.5, 0], [0, 0, 0]),
)
iiwa_2 = parser.AddModelFromFile(iiwa_file, model_name="iiwa_2")
plant.WeldFrames(
frame_on_parent_F=plant.world_frame(),
frame_on_child_M=plant.GetFrameByName("iiwa_link_0", iiwa_2),
X_FM=xyz_rpy_deg([0, 0.5, 0], [0, 0, 0]),
)
renderer_name = "renderer"
scene_graph.AddRenderer(
renderer_name, MakeRenderEngineVtk(RenderEngineVtkParams()))
# N.B. These properties are chosen arbitrarily.
intrinsics = CameraInfo(
width=640,
height=480,
fov_y=np.pi/4,
)
core = RenderCameraCore(
renderer_name,
intrinsics,
ClippingRange(0.01, 10.0),
RigidTransform(),
)
color_camera = ColorRenderCamera(core, show_window=False)
depth_camera = DepthRenderCamera(core, DepthRange(0.01, 10.0))
Because we add this to plant.world_body()
, this is a scene-fixed camera.
To make it a moving camera, simply add it to a body that is not anchored against the world.
world_id = plant.GetBodyFrameIdOrThrow(plant.world_body().index())
X_WB = xyz_rpy_deg([2, 0, 0.75], [-90, 0, 90])
sensor = RgbdSensor(
world_id,
X_PB=X_WB,
color_camera=color_camera,
depth_camera=depth_camera,
)
builder.AddSystem(sensor)
builder.Connect(
scene_graph.get_query_output_port(),
sensor.query_object_input_port(),
)
Meshcat
Visualizer.¶meshcat_vis = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)
plant.Finalize()
diagram = builder.Build()
# TODO(eric.cousineau): Replace this with `diagram.Publish(diagram_context)`
# once all visualizers no longer use initialization events.
Simulator(diagram).Initialize()
diagram_context = diagram.CreateDefaultContext()
sensor_context = sensor.GetMyMutableContextFromRoot(diagram_context)
sg_context = scene_graph.GetMyMutableContextFromRoot(diagram_context)
Note that this uses the default labeling scheme, using body.index()
.
color = sensor.color_image_output_port().Eval(sensor_context).data
depth = sensor.depth_image_32F_output_port().Eval(sensor_context).data.squeeze(2)
label = sensor.label_image_output_port().Eval(sensor_context).data
fig, ax = plt.subplots(1, 3, figsize=(15, 10))
ax[0].imshow(color)
ax[1].imshow(colorize_depth(depth))
ax[2].imshow(colorize_labels(label))
We will loop through each geometry item, get it's corresponding body, and then remap the original label to a label that is only distinct for model instances.
NOTE: If the labels in the given plant merge the two model instances together, then this will not fix that. The correct behavior would be to update the RenderLabel values in the geometry instances themselves.
TODO(eric.cousineau): Add an example of updating only the RenderLabel.
query_object = scene_graph.get_query_output_port().Eval(sg_context)
inspector = query_object.inspector()
label_by_model = label.copy()
for geometry_id in inspector.GetAllGeometryIds():
body = plant.GetBodyFromFrameId(inspector.GetFrameId(geometry_id))
geometry_label = inspector.GetPerceptionProperties(
geometry_id).GetProperty("label", "id")
# WARNING: If you do not cast the `geometry_label` to `int`, this
# comparison will take a long time since NumPy will do
# element-by-element comparison using `RenderLabel.__eq__`.
mask = (label == int(geometry_label))
label_by_model[mask] = int(body.model_instance())
plt.figure(figsize=(5, 5))
plt.imshow(colorize_labels(label_by_model))