Pybind11 wrapping of multibody/shapes (RobotLocomotion#7913)

bindings/pydrake/BUILD.bazel

@@ -87,6 +87,7 @@ drake_pybind_library(
         ":autodiffutils_py",
         ":common_py",
         ":parsers_py",
+        "//bindings/pydrake/multibody:shapes_py",
     ],
     py_srcs = ["rbtree.py"],
 )

bindings/pydrake/multibody/BUILD.bazel

@@ -44,20 +44,33 @@ drake_pybind_library(
     cc_srcs = ["rigid_body_plant_py.cc"],
     package_info = PACKAGE_INFO,
     py_deps = [
+        ":module_py",
         "//bindings/pydrake:rbtree_py",
         "//bindings/pydrake/systems:framework_py",
     ],
 )
 
+drake_pybind_library(
+    name = "shapes_py",
+    cc_so_name = "shapes",
+    cc_srcs = ["shapes_py.cc"],
+    package_info = PACKAGE_INFO,
+    py_deps = [
+        ":module_py",
+    ],
+)
+
 drake_py_library(
     name = "all_py",
     deps = [
         ":rigid_body_plant_py",
+        ":shapes_py",
     ],
 )
 
 PY_LIBRARIES_WITH_INSTALL = [
     ":rigid_body_plant_py",
+    ":shapes_py",
 ]
 
 PY_LIBRARIES = [
@@ -87,4 +100,13 @@ drake_py_test(
     ],
 )
 
+drake_py_test(
+    name = "shapes_test",
+    size = "small",
+    data = ["//examples/quadrotor:models"],
+    deps = [
+        ":shapes_py",
+    ],
+)
+
 add_lint_tests()

bindings/pydrake/multibody/all.py

@@ -1,4 +1,5 @@
 from .rigid_body_plant import *
+from .shapes import *
 
 try:
     from .drawing import *

bindings/pydrake/multibody/shapes_py.cc

@@ -0,0 +1,91 @@
+#include <pybind11/eigen.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "drake/bindings/pydrake/pydrake_pybind.h"
+#include "drake/multibody/shapes/drake_shapes.h"
+
+namespace drake {
+namespace pydrake {
+
+PYBIND11_MODULE(shapes, m) {
+  // NOLINTNEXTLINE(build/namespaces): Emulate placement in namespace.
+  using namespace DrakeShapes;
+
+  m.doc() = "Core geometry and shape types.";
+
+  py::enum_<Shape>(m, "Shape")
+    .value("UNKNOWN", Shape::UNKNOWN)
+    .value("BOX", Shape::BOX)
+    .value("SPHERE", Shape::SPHERE)
+    .value("CYLINDER", Shape::CYLINDER)
+    .value("MESH", Shape::MESH)
+    .value("MESH_POINTS", Shape::MESH_POINTS)
+    .value("CAPSULE", Shape::CAPSULE);
+
+  py::class_<Geometry>(m, "Geometry")
+    .def("getShape", &Geometry::getShape)
+    .def("hasFaces", &Geometry::hasFaces)
+    .def("getFaces", [](const Geometry* self) {
+          TrianglesVector triangles;
+          self->getFaces(&triangles);
+          return triangles;
+        })
+    .def("getPoints",
+        [](const Geometry* self) {
+          Eigen::Matrix3Xd pts(3, 0);
+          self->getPoints(pts);
+          return pts;
+        })
+    .def("getBoundingBoxPoints",
+        [](const Geometry* self) {
+          Eigen::Matrix3Xd pts(3, 0);
+          self->getBoundingBoxPoints(pts);
+          return pts;
+        });
+
+  py::class_<Box, Geometry>(m, "Box")
+    .def(py::init<const Eigen::Vector3d&>(), py::arg("size"))
+    .def_readonly("size", &Box::size);
+  py::class_<Sphere, Geometry>(m, "Sphere")
+    .def(py::init<double>(), py::arg("radius"))
+    .def_readonly("radius", &Sphere::radius);
+  py::class_<Cylinder, Geometry>(m, "Cylinder")
+    .def(py::init<double, double>(), py::arg("radius"), py::arg("length"))
+    .def_readonly("radius", &Cylinder::radius)
+    .def_readonly("length", &Cylinder::length);
+  py::class_<Mesh, Geometry>(m, "Mesh")
+    .def(py::init<const std::string&, const std::string&>(),
+         py::arg("uri"), py::arg("resolved_filename"))
+    .def_readonly("scale", &Mesh::scale_)
+    .def_readonly("uri", &Mesh::uri_)
+    .def_readonly("resolved_filename", &Mesh::resolved_filename_);
+  py::class_<MeshPoints, Geometry>(m, "MeshPoints")
+    .def(py::init<const Eigen::Matrix3Xd&>(),
+         py::arg("points"));
+  py::class_<Capsule, Geometry>(m, "Capsule")
+    .def(py::init<double, double>(), py::arg("radius"), py::arg("length"))
+    .def_readonly("radius", &Capsule::radius)
+    .def_readonly("length", &Capsule::length);
+
+  py::class_<Element>(m, "Element")
+    .def(py::init<const Geometry&, const Eigen::Isometry3d&>(),
+         py::arg("geometry_in"),
+         py::arg("T_element_to_local"))
+    .def("hasGeometry", &Element::hasGeometry)
+    .def("getGeometry", &Element::getGeometry, py_reference_internal)
+    .def("getLocalTransform", [](const Element* self) {
+      return self->getLocalTransform().matrix();
+    });
+  py::class_<VisualElement, Element>(m, "VisualElement")
+    .def(py::init<const Geometry&, const Eigen::Isometry3d&,
+                  const Eigen::Vector4d&>(),
+         py::arg("geometry_in"),
+         py::arg("T_element_to_local"),
+         py::arg("material_in"))
+    .def("setMaterial", &VisualElement::setMaterial, "Apply an RGBA material.")
+    .def("getMaterial", &VisualElement::getMaterial, "Get an RGBA material.");
+}
+
+}  // namespace pydrake
+}  // namespace drake

bindings/pydrake/multibody/test/shapes_test.py

@@ -0,0 +1,50 @@
+import os
+import unittest
+
+import numpy as np
+
+import pydrake
+from pydrake.multibody import shapes
+
+
+class TestShapes(unittest.TestCase):
+    def test_api(self):
+        box_size = [1., 2., 3.]
+        radius = 0.1
+        length = 0.2
+
+        box = shapes.Box(size=box_size)
+        self.assertTrue(np.allclose(box.size, box_size))
+        self.assertEqual(box.getPoints().shape, (3, 8))
+        self.assertEqual(len(box.getFaces()), 12)
+        self.assertEqual(len(box.getFaces()[0]), 3)
+
+        sphere = shapes.Sphere(radius=radius)
+        self.assertEqual(sphere.radius, radius)
+        self.assertEqual(sphere.getPoints().shape, (3, 1))
+        with self.assertRaises(RuntimeError):
+            sphere.getFaces()
+
+        cylinder = shapes.Cylinder(radius=radius, length=length)
+        self.assertEqual(cylinder.radius, radius)
+        self.assertEqual(cylinder.length, length)
+
+        capsule = shapes.Capsule(radius=radius, length=length)
+        self.assertEqual(capsule.radius, radius)
+        self.assertEqual(capsule.length, length)
+
+        pts = np.tile(box_size, (10, 1)).T
+        mesh_points = shapes.MeshPoints(pts)
+        self.assertEqual(mesh_points.getPoints().shape, (3, 10))
+
+        obj_mesh_path = os.path.join(
+            pydrake.getDrakePath(), "examples/quadrotor/quadrotor_base.obj")
+        obj_mesh_uri = "box_obj"
+        mesh = shapes.Mesh(uri=obj_mesh_uri, resolved_filename=obj_mesh_path)
+        self.assertTrue(np.allclose(mesh.scale, [1., 1., 1.]))
+        self.assertEqual(mesh.uri, obj_mesh_uri)
+        self.assertEqual(mesh.resolved_filename, obj_mesh_path)
+
+
+if __name__ == '__main__':
+    unittest.main()

bindings/pydrake/rbtree_py.cc

@@ -21,6 +21,7 @@ PYBIND11_MODULE(_rbtree_py, m) {
   using drake::parsers::PackageMap;
   namespace sdf = drake::parsers::sdf;
 
+  py::module::import("pydrake.multibody.shapes");
   py::module::import("pydrake.parsers");
 
   py::enum_<FloatingBaseType>(m, "FloatingBaseType")
@@ -123,6 +124,7 @@ PYBIND11_MODULE(_rbtree_py, m) {
          py::arg("in_terms_of_qdot") = false)
     .def("get_num_bodies", &RigidBodyTree<double>::get_num_bodies)
     .def("get_num_frames", &RigidBodyTree<double>::get_num_frames)
+    .def("get_num_actuators", &RigidBodyTree<double>::get_num_actuators)
     .def("getBodyOrFrameName",
          &RigidBodyTree<double>::getBodyOrFrameName,
          py::arg("body_or_frame_id"))
@@ -210,7 +212,8 @@ PYBIND11_MODULE(_rbtree_py, m) {
   py::class_<RigidBody<double> >(m, "RigidBody")
     .def("get_name", &RigidBody<double>::get_name)
     .def("get_body_index", &RigidBody<double>::get_body_index)
-    .def("get_center_of_mass", &RigidBody<double>::get_center_of_mass);
+    .def("get_center_of_mass", &RigidBody<double>::get_center_of_mass)
+    .def("get_visual_elements", &RigidBody<double>::get_visual_elements);
 
   py::class_<RigidBodyFrame<double>,
              std::shared_ptr<RigidBodyFrame<double> > >(m, "RigidBodyFrame")

bindings/pydrake/systems/framework_py.cc

@@ -219,7 +219,8 @@ PYBIND11_MODULE(framework, m) {
   // Value types.
   py::class_<VectorBase<T>>(m, "VectorBase")
     .def("CopyToVector", &VectorBase<T>::CopyToVector)
-    .def("SetFromVector", &VectorBase<T>::SetFromVector);
+    .def("SetFromVector", &VectorBase<T>::SetFromVector)
+    .def("size", &VectorBase<T>::size);
 
   // TODO(eric.cousineau): Make a helper function for the Eigen::Ref<> patterns.
   py::class_<BasicVector<T>, VectorBase<T>>(m, "BasicVector")

bindings/pydrake/systems/test/vector_test.py

@@ -28,11 +28,13 @@ def test_basic_vector_double(self):
         self.assertTrue(np.allclose(value, expected))
         self.assertTrue(np.allclose(value_data.get_value(), expected))
         self.assertTrue(np.allclose(value_data.get_mutable_value(), expected))
+        self.assertEqual(value_data.size(), 3)
         expected = [5., 6, 7]
         value_data.SetFromVector(expected)
         self.assertTrue(np.allclose(value, expected))
         self.assertTrue(np.allclose(value_data.get_value(), expected))
         self.assertTrue(np.allclose(value_data.get_mutable_value(), expected))
+        self.assertEqual(value_data.size(), 3)
 
 
 if __name__ == '__main__':

bindings/pydrake/test/atlas_constructor_test.py

@@ -16,6 +16,7 @@ def test_constructor(self):
         robot = rbtree.RigidBodyTree(
             model, package_map=pm,
             floating_base_type=rbtree.FloatingBaseType.kRollPitchYaw)
+        self.assertEqual(robot.get_num_actuators(), 30)
 
 
 if __name__ == '__main__':

bindings/pydrake/test/rbt_com_test.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pydrake
 from pydrake.rbtree import RigidBodyTree, FloatingBaseType
+import pydrake.multibody.shapes as shapes
 import os.path
 
 
@@ -59,14 +60,16 @@ def assert_sane(x, nonzero=True):
                 self.assertTrue(np.any(x != 0))
 
         num_q = num_v = 7
+        num_u = r.get_num_actuators()
+        self.assertEquals(num_u, 1)
         q = np.zeros(num_q)
         v = np.zeros(num_v)
         # Update kinematics.
         kinsol = r.doKinematics(q, v)
         # Sanity checks:
         # - Actuator map.
-        self.assertEquals(r.B.shape, (num_v, 1))
-        B_expected = np.zeros((num_v, 1))
+        self.assertEquals(r.B.shape, (num_v, num_u))
+        B_expected = np.zeros((num_v, num_u))
         B_expected[-1] = 1
         self.assertTrue(np.allclose(r.B, B_expected))
         # - Mass matrix.
@@ -87,6 +90,59 @@ def assert_sane(x, nonzero=True):
         self.assertTrue(friction_torques.shape, (num_v,))
         assert_sane(friction_torques, nonzero=False)
 
+    def testRigidBodyGeometry(self):
+        # TODO(gizatt) This test ought to have its reliance on
+        # the Pendulum model specifics stripped (so that e.g. material changes
+        # don't break the test), and split pure API testing of
+        # VisualElement and Geometry over to shapes_test while keeping
+        # RigidBodyTree and RigidBody visual element extraction here.
+        urdf_path = os.path.join(
+            pydrake.getDrakePath(), "examples/pendulum/Pendulum.urdf")
+
+        tree = RigidBodyTree(
+            urdf_path,
+            floating_base_type=FloatingBaseType.kFixed)
+
+        # "base_part2" should have a single visual element.
+        base_part2 = tree.FindBody("base_part2")
+        self.assertIsNotNone(base_part2)
+        visual_elements = base_part2.get_visual_elements()
+        self.assertEqual(len(visual_elements), 1)
+        self.assertIsInstance(visual_elements[0], shapes.VisualElement)
+
+        # It has green material by default
+        sphere_visual_element = visual_elements[0]
+        green_material = np.array([0.3, 0.6, 0.4, 1])
+        white_material = np.array([1., 1., 1., 1.])
+
+        self.assertTrue(np.allclose(sphere_visual_element.getMaterial(),
+                        green_material))
+        sphere_visual_element.setMaterial(white_material)
+        self.assertTrue(np.allclose(sphere_visual_element.getMaterial(),
+                        white_material))
+
+        # We expect this link TF to have positive z-component...
+        local_tf = sphere_visual_element.getLocalTransform()
+        self.assertAlmostEqual(local_tf[2, 3], 0.015)
+
+        # ... as well as sphere geometry.
+        self.assertTrue(sphere_visual_element.hasGeometry())
+        sphere_geometry = sphere_visual_element.getGeometry()
+        self.assertIsInstance(sphere_geometry, shapes.Sphere)
+        self.assertEqual(sphere_geometry.getShape(), shapes.Shape.SPHERE)
+        self.assertNotEqual(sphere_geometry.getShape(), shapes.Shape.BOX)
+
+        # For a sphere geometry, getPoints() should return
+        # one point at the center of the sphere.
+        sphere_geometry_pts = sphere_geometry.getPoints()
+        self.assertEqual(sphere_geometry_pts.shape, (3, 1))
+        sphere_geometry_bb = sphere_geometry.getBoundingBoxPoints()
+        self.assertEqual(sphere_geometry_bb.shape, (3, 8))
+        # Sphere's don't have faces supplied (yet?)
+        self.assertFalse(sphere_geometry.hasFaces())
+        with self.assertRaises(RuntimeError):
+            sphere_geometry.getFaces()
+
 
 if __name__ == '__main__':
     unittest.main()

multibody/shapes/visual_element.h

@@ -6,21 +6,33 @@
 #include "drake/multibody/shapes/element.h"
 
 namespace DrakeShapes {
+
 class VisualElement : public Element {
  public:
   explicit VisualElement(const Eigen::Isometry3d& T_element_to_local)
       : Element(T_element_to_local),
         material(Eigen::Vector4d(0.7, 0.7, 0.7, 1)) {}
 
+  /**
+   * Constructs a geometry at a specified transform with
+   * a given material (i.e. a color specified as a 4-vector
+   * of RGBA, each on [0, 1]).
+   */
   VisualElement(const Geometry& geometry,
                 const Eigen::Isometry3d& T_element_to_local,
                 const Eigen::Vector4d& material_in)
       : Element(geometry, T_element_to_local), material(material_in) {}
 
   virtual ~VisualElement() {}
 
+  /**
+  * Sets the element's material color, in RGBA format.
+  */
   void setMaterial(const Eigen::Vector4d& material);
 
+  /**
+  * Retrieves the element's material color, in RGBA format.
+  */
   const Eigen::Vector4d& getMaterial() const;
 
  protected: