diff --git a/drake_ros/core/rgbd_system.cc b/drake_ros/core/rgbd_system.cc index d40505ce..712b5de8 100644 --- a/drake_ros/core/rgbd_system.cc +++ b/drake_ros/core/rgbd_system.cc @@ -70,6 +70,43 @@ void RGBDSystem::CalcDepthImage(const drake::systems::Context& context, } } +const InputPort& RGBDSystem::DeclareDepthInputPort( + PixelType pixel_type, std::string port_name, + double publish_period, double start_time) { + switch (pixel_type) { + case PixelType::kRgb8U: + break; + case PixelType::kBgr8U: + break; + case PixelType::kRgba8U: { + break; + } + case PixelType::kBgra8U: + break; + case PixelType::kDepth16U: { + return this->template DeclareDepthInputPort( + std::move(port_name), publish_period, + start_time); + } + case PixelType::kDepth32F: { + return this->template DeclareDepthInputPort( + std::move(port_name), publish_period, + start_time); + } + case PixelType::kLabel16I: { + break; + } + case PixelType::kGrey8U: { + break; + } + default: + break; + } + throw std::logic_error(fmt::format( + "RGBDSystem::DeclareDepthInputPort does not support pixel_type={}", + static_cast(pixel_type))); +} + template const InputPort& RGBDSystem::DeclareDepthInputPort( std::string port_name, double publish_period, double start_time) { @@ -102,6 +139,43 @@ const InputPort& RGBDSystem::DeclareDepthInputPort( return port; } +const InputPort& RGBDSystem::DeclareImageInputPort( + PixelType pixel_type, std::string port_name, + double publish_period, double start_time) { + switch (pixel_type) { + case PixelType::kRgb8U: + break; + case PixelType::kBgr8U: + break; + case PixelType::kRgba8U: { + return this->template DeclareImageInputPort( + std::move(port_name), publish_period, + start_time); + } + case PixelType::kBgra8U: + break; + case PixelType::kDepth16U: { + break; + } + case PixelType::kDepth32F: { + break; + } + case PixelType::kLabel16I: { + break; + } + case PixelType::kGrey8U: { + return this->template DeclareImageInputPort( + std::move(port_name), publish_period, + start_time); + } + default: + break; + } + throw std::logic_error(fmt::format( + "RGBDSystem::DeclareImageInputPort does not support pixel_type={}", + static_cast(pixel_type))); +} + template const InputPort& RGBDSystem::DeclareImageInputPort( std::string port_name, double publish_period, double start_time) { diff --git a/drake_ros/core/rgbd_system.h b/drake_ros/core/rgbd_system.h index 828a8ec7..d62e5491 100644 --- a/drake_ros/core/rgbd_system.h +++ b/drake_ros/core/rgbd_system.h @@ -32,11 +32,20 @@ class RGBDSystem : public drake::systems::LeafSystem { ~RGBDSystem() override; + const InputPort& DeclareImageInputPort(PixelType pixel_type, + std::string port_name, + double publish_period, + double start_time); + template const InputPort& DeclareImageInputPort(std::string port_name, double publish_period, double start_time); + const InputPort& DeclareDepthInputPort(PixelType pixel_type, + std::string port_name, + double publish_period, + double start_time); template const InputPort& DeclareDepthInputPort(std::string port_name, double publish_period, diff --git a/drake_ros/drake_ros/core/__init__.py b/drake_ros/drake_ros/core/__init__.py index b81bf9d7..a50174ba 100644 --- a/drake_ros/drake_ros/core/__init__.py +++ b/drake_ros/drake_ros/core/__init__.py @@ -6,6 +6,7 @@ from drake_ros._cc.core import Isometry3ToRosPose from drake_ros._cc.core import Isometry3ToRosTransform from drake_ros._cc.core import QuaternionToRosQuaternion +from drake_ros._cc.core import RGBDSystem from drake_ros._cc.core import RigidTransformToRosPose from drake_ros._cc.core import RigidTransformToRosTransform from drake_ros._cc.core import RosAccelToSpatialAcceleration @@ -106,6 +107,7 @@ def _make_ros_subscriber_system( 'Isometry3ToRosTransform', 'PySerializer', 'QuaternionToRosQuaternion', + 'RGBDSystem', 'RigidTransformToRosPose', 'RigidTransformToRosTransform', 'RosAccelToSpatialAcceleration', diff --git a/drake_ros/drake_ros/core/cc_pybind.cc b/drake_ros/drake_ros/core/cc_pybind.cc index a2b51ce3..bfdafc6d 100644 --- a/drake_ros/drake_ros/core/cc_pybind.cc +++ b/drake_ros/drake_ros/core/cc_pybind.cc @@ -11,6 +11,7 @@ #include "drake_ros/core/drake_ros.h" #include "drake_ros/core/geometry_conversions.h" #include "drake_ros/core/geometry_conversions_pybind.h" +#include "drake_ros/core/rgbd_system.h" #include "drake_ros/core/qos_pybind.h" #include "drake_ros/core/ros_interface_system.h" #include "drake_ros/core/ros_publisher_system.h" @@ -25,6 +26,7 @@ using drake_ros::core::CameraInfoSystem; using drake_ros::core::ClockSystem; using drake_ros::core::DrakeRos; using drake_ros::core::init; +using drake_ros::core::RGBDSystem; using drake_ros::core::RosInterfaceSystem; using drake_ros::core::RosPublisherSystem; using drake_ros::core::RosSubscriberSystem; @@ -36,6 +38,29 @@ using drake::systems::TriggerType; using py_rvp = pybind11::return_value_policy; +// Implementation for `overload_cast_explicit`. We must use this structure so +// that we can constrain what is inferred. Otherwise, the ambiguity confuses +// the compiler. +template +struct overload_cast_impl { + auto operator()(Return (*func)(Args...)) const { return func; } + + template + auto operator()(Return (Class::*method)(Args...)) const { + return method; + } + + template + auto operator()(Return (Class::*method)(Args...) const) const { + return method; + } +}; + +/// Provides option to provide explicit signature when +/// `py::overload_cast` fails to infer the Return argument. +template +constexpr auto overload_cast_explicit = overload_cast_impl{}; + // A (de)serialization interface implementation for Python ROS messages // that can be overriden from Python itself. class PySerializerInterface : public py::wrapper { @@ -154,6 +179,64 @@ void DefCore(py::module m) { RosPublisherSystem::kDefaultTriggerTypes}, py::arg("publish_period") = 0.0); + py::class_>(m, "RGBDSystem") + .def(py::init<>()) + .def( + "DeclareImageInputPort", + [](RGBDSystem& self, + drake::systems::sensors::PixelType pixel_type, + std::string port_name, + double publish_period, + double start_time) -> const drake::systems::InputPort& { + return self.DeclareImageInputPort(pixel_type, + std::move(port_name), + publish_period, start_time); + }, + py::arg("pixel_type"), py::arg("port_name"), + py::arg("publish_period"), + py::arg("start_time"), py_rvp::reference_internal) + .def( + "DeclareDepthInputPort", + [](RGBDSystem& self, + drake::systems::sensors::PixelType pixel_type, + std::string port_name, + double publish_period, + double start_time) -> const drake::systems::InputPort& { + return self.DeclareDepthInputPort(pixel_type, + std::move(port_name), + publish_period, start_time); + }, + py::arg("pixel_type"), py::arg("port_name"), + py::arg("publish_period"), + py::arg("start_time"), py_rvp::reference_internal) + .def_static( + "AddToBuilder", + [](drake::systems::DiagramBuilder* builder, DrakeRos* ros, + const std::string& topic_name, const std::string& depth_topic_name, + const QoS& qos, + const std::unordered_set& + pub_triggers, + double publish_period) { + auto [rgba_system, pub_system] = RGBDSystem::AddToBuilder( + builder, ros, topic_name, depth_topic_name, qos, pub_triggers, publish_period); + + py::object py_builder = py::cast(builder, py_rvp::reference); + py::list result; + result.append( + py::cast(rgba_system, py_rvp::reference_internal, py_builder)); + result.append( + py::cast(pub_system, py_rvp::reference_internal, py_builder)); + return result; + }, + py::arg("builder"), py::arg("ros"), py::kw_only(), + py::arg("topic_name") = std::string{"/image"}, + py::arg("depth_topic_name") = std::string{"/depth"}, + py::arg("qos") = drake_ros::QoS(rclcpp::SystemDefaultsQoS()), + py::arg("publish_triggers") = + std::unordered_set{ + RosPublisherSystem::kDefaultTriggerTypes}, + py::arg("publish_period") = 0.0); + m.def( "init", [](std::vector args) { diff --git a/drake_ros_examples/examples/rgbd/rgbd.py b/drake_ros_examples/examples/rgbd/rgbd.py new file mode 100644 index 00000000..fe4608c5 --- /dev/null +++ b/drake_ros_examples/examples/rgbd/rgbd.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +import numpy as np +import os + +from ament_index_python.packages import get_package_share_directory + +import drake_ros.core +from drake_ros.core import ClockSystem +from drake_ros.core import CameraInfoSystem +from drake_ros.core import RGBDSystem +from drake_ros.core import RosInterfaceSystem + +from pydrake.geometry import (ClippingRange, ColorRenderCamera, DepthRange, DepthRenderCamera, + MakeRenderEngineVtk, RenderCameraCore, RenderEngineVtkParams) +from pydrake.math import RigidTransform, RollPitchYaw +from pydrake.multibody.parsing import Parser +from pydrake.multibody.plant import AddMultibodyPlant, MultibodyPlantConfig +from pydrake.systems.analysis import Simulator +from pydrake.systems.framework import DiagramBuilder +from pydrake.systems.sensors import CameraInfo, PixelType, RgbdSensor +from pydrake.visualization import ( + ColorizeDepthImage, + ColorizeLabelImage, +) + +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) + + +def main(): + # Create a Drake diagram + builder = DiagramBuilder() + + # Add a multibody plant and a scene graph to hold the robots + plant, scene_graph = AddMultibodyPlant( + MultibodyPlantConfig(time_step=0.001), + builder, + ) + + renderer_name = "renderer" + scene_graph.AddRenderer( + renderer_name, MakeRenderEngineVtk(RenderEngineVtkParams())) + parser = Parser(plant) + + package_path = get_package_share_directory("drake_ros_examples") + parser.package_map().Add("drake_ros_examples", package_path) + fs_path = parser.package_map().GetPath("drake_ros_examples") + sdf_url = os.path.join(fs_path, 'rgbd', 'rgbd.sdf') + + Parser(plant, scene_graph).AddModels(sdf_url) + + # Initialise the ROS infrastructure + drake_ros.core.init() + # Create a Drake system to interface with ROS + sys_ros_interface = builder.AddSystem(RosInterfaceSystem('cart_pole')) + ClockSystem.AddToBuilder(builder, sys_ros_interface.get_ros_interface()) + + camera_info_system = CameraInfoSystem.AddToBuilder( + builder, + sys_ros_interface.get_ros_interface()) + + depth_camera_info_system = CameraInfoSystem.AddToBuilder( + builder, + sys_ros_interface.get_ros_interface(), + topic_name='/depth/camera_info') + + 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)) + + camera_info_system[0].set_camera_info(intrinsics) + depth_camera_info_system[0].set_camera_info(intrinsics) + + world_id = plant.GetBodyFrameIdOrThrow(plant.world_body().index()) + X_WB = xyz_rpy_deg([0, 1, -0.12], [90, 180, 0]) + 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(), + ) + + colorize_depth = builder.AddSystem(ColorizeDepthImage()) + colorize_label = builder.AddSystem(ColorizeLabelImage()) + colorize_label.background_color.set([0,0,0]) + builder.Connect(sensor.GetOutputPort("depth_image_32f"), + colorize_depth.GetInputPort("depth_image_32f")) + builder.Connect(sensor.GetOutputPort("label_image"), + colorize_label.GetInputPort("label_image")) + + rgbd_publisher = builder.AddSystem(RGBDSystem()) + + image_publish_period = 1. / 30. + rgbd_port = rgbd_publisher.DeclareImageInputPort( + PixelType.kRgba8U, "color", image_publish_period, 0.) + builder.Connect(sensor.color_image_output_port(), rgbd_port) + + depth_port = rgbd_publisher.DeclareDepthInputPort( + PixelType.kDepth32F, "depth", image_publish_period, 0.) + builder.Connect(sensor.depth_image_32F_output_port(), depth_port) + + [pub_color_system, pub_depth_system] = RGBDSystem.AddToBuilder( + builder, sys_ros_interface.get_ros_interface()) + + builder.Connect(rgbd_publisher.GetOutputPort("rgbd_color"), + pub_color_system.get_input_port()) + + builder.Connect(rgbd_publisher.GetOutputPort("rgbd_depth"), + pub_depth_system.get_input_port()) + + plant.Finalize() + + # Build the complete system from the diagram + diagram = builder.Build() + + diagram_context = diagram.CreateDefaultContext() + cart_pole_context = plant.GetMyMutableContextFromRoot(diagram_context) + + plant.get_actuation_input_port().FixValue(cart_pole_context, 0) + + cart_slider = plant.GetJointByName("CartSlider") + pole_pin = plant.GetJointByName("PolePin") + + cart_slider.set_translation(context=cart_pole_context, translation=0.0) + pole_pin.set_angle(context=cart_pole_context, angle=2.0) + + # Create a simulator for the system + simulator = Simulator(diagram, diagram_context) + simulator.Initialize() + simulator_context = simulator.get_mutable_context() + simulator.set_publish_every_time_step(False) + simulator.set_target_realtime_rate(1.0) + + # Step the simulator in 0.1s intervals + step = 0.1 + while simulator_context.get_time() < float('inf'): + next_time = min( + simulator_context.get_time() + step, float('inf'), + ) + simulator.AdvanceTo(next_time) + + +if __name__ == '__main__': + main()