Skip to content

Commit

Permalink
Added RGBDSystem Python wrapper
Browse files Browse the repository at this point in the history 
Signed-off-by: Alejandro Hernández Cordero <[email protected]>
  • Loading branch information
@ahcorde
ahcorde committed Dec 27, 2023
1 parent 21c8ace commit 8fab471
Show file tree
Hide file tree
Showing 5 changed files with 332 additions and 0 deletions.
74 changes: 74 additions & 0 deletions drake_ros/core/rgbd_system.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,43 @@ void RGBDSystem::CalcDepthImage(const drake::systems::Context<double>& context,
}
}

const InputPort<double>& 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<PixelType::kDepth32F>(
std::move(port_name), publish_period,
start_time);
}
case PixelType::kDepth32F: {
return this->template DeclareDepthInputPort<PixelType::kDepth32F>(
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<int>(pixel_type)));
}

template <PixelType kPixelType>
const InputPort<double>& RGBDSystem::DeclareDepthInputPort(
std::string port_name, double publish_period, double start_time) {
Expand Down Expand Up @@ -102,6 +139,43 @@ const InputPort<double>& RGBDSystem::DeclareDepthInputPort(
return port;
}

const InputPort<double>& 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<PixelType::kRgba8U>(
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<PixelType::kGrey8U>(
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<int>(pixel_type)));
}

template <PixelType kPixelType>
const InputPort<double>& RGBDSystem::DeclareImageInputPort(
std::string port_name, double publish_period, double start_time) {
Expand Down
9 changes: 9 additions & 0 deletions drake_ros/core/rgbd_system.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,20 @@ class RGBDSystem : public drake::systems::LeafSystem<double> {

~RGBDSystem() override;

const InputPort<double>& DeclareImageInputPort(PixelType pixel_type,
std::string port_name,
double publish_period,
double start_time);

template <PixelType kPixelType>
const InputPort<double>& DeclareImageInputPort(std::string port_name,
double publish_period,
double start_time);

const InputPort<double>& DeclareDepthInputPort(PixelType pixel_type,
std::string port_name,
double publish_period,
double start_time);
template <PixelType kPixelType>
const InputPort<double>& DeclareDepthInputPort(std::string port_name,
double publish_period,
Expand Down
2 changes: 2 additions & 0 deletions drake_ros/drake_ros/core/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -106,6 +107,7 @@ def _make_ros_subscriber_system(
'Isometry3ToRosTransform',
'PySerializer',
'QuaternionToRosQuaternion',
'RGBDSystem',
'RigidTransformToRosPose',
'RigidTransformToRosTransform',
'RosAccelToSpatialAcceleration',
Expand Down
83 changes: 83 additions & 0 deletions drake_ros/drake_ros/core/cc_pybind.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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;
Expand All @@ -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 <typename Return, typename... Args>
struct overload_cast_impl {
auto operator()(Return (*func)(Args...)) const { return func; }

template <typename Class>
auto operator()(Return (Class::*method)(Args...)) const {
return method;
}

template <typename Class>
auto operator()(Return (Class::*method)(Args...) const) const {
return method;
}
};

/// Provides option to provide explicit signature when
/// `py::overload_cast<Args...>` fails to infer the Return argument.
template <typename Return, typename... Args>
constexpr auto overload_cast_explicit = overload_cast_impl<Return, Args...>{};

// A (de)serialization interface implementation for Python ROS messages
// that can be overriden from Python itself.
class PySerializerInterface : public py::wrapper<SerializerInterface> {
Expand Down Expand Up @@ -154,6 +179,64 @@ void DefCore(py::module m) {
RosPublisherSystem::kDefaultTriggerTypes},
py::arg("publish_period") = 0.0);

py::class_<RGBDSystem, LeafSystem<double>>(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<double>& {
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<double>& {
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<double>* builder, DrakeRos* ros,
const std::string& topic_name, const std::string& depth_topic_name,
const QoS& qos,
const std::unordered_set<drake::systems::TriggerType>&
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<drake::systems::TriggerType>{
RosPublisherSystem::kDefaultTriggerTypes},
py::arg("publish_period") = 0.0);

m.def(
"init",
[](std::vector<std::string> args) {
Expand Down
Loading

0 comments on commit 8fab471

Please sign in to comment.