Improve Rotation3d

pycolmap/geometry/bindings.h

@@ -5,6 +5,7 @@
 
 #include "pycolmap/geometry/homography_matrix.h"
 #include "pycolmap/helpers.h"
+#include "pycolmap/pybind11_extension.h"
 
 #include <sstream>
 
@@ -20,14 +21,20 @@ using namespace pybind11::literals;
 void BindGeometry(py::module& m) {
   BindHomographyGeometry(m);
 
-  py::class_<Eigen::Quaterniond> PyRotation3d(m, "Rotation3d");
+  py::class_ext_<Eigen::Quaterniond> PyRotation3d(m, "Rotation3d");
   PyRotation3d.def(py::init([]() { return Eigen::Quaterniond::Identity(); }))
       .def(py::init<const Eigen::Vector4d&>(),
            "xyzw"_a,
            "Quaternion in [x,y,z,w] format.")
       .def(py::init<const Eigen::Matrix3d&>(),
            "rotmat"_a,
            "3x3 rotation matrix.")
+      .def(py::init([](const Eigen::Vector3d& vec) {
+             return Eigen::Quaterniond(
+                 Eigen::AngleAxis<double>(vec.norm(), vec.normalized()));
+           }),
+           "axis_angle"_a,
+           "Axis-angle 3D vector.")
       .def_property(
           "quat",
           py::overload_cast<>(&Eigen::Quaterniond::coeffs),
@@ -46,6 +53,14 @@ void BindGeometry(py::module& m) {
       .def("normalize", &Eigen::Quaterniond::normalize)
       .def("matrix", &Eigen::Quaterniond::toRotationMatrix)
       .def("norm", &Eigen::Quaterniond::norm)
+      .def("angle",
+           [](const Eigen::Quaterniond& self) {
+             return Eigen::AngleAxis<double>(self).angle();
+           })
+      .def("angle_to",
+           [](const Eigen::Quaterniond& self, const Eigen::Quaterniond& other) {
+             return self.angularDistance(other);
+           })
       .def("inverse", &Eigen::Quaterniond::inverse)
       .def("__repr__", [](const Eigen::Quaterniond& self) {
         std::stringstream ss;
@@ -55,7 +70,7 @@ void BindGeometry(py::module& m) {
   py::implicitly_convertible<py::array, Eigen::Quaterniond>();
   MakeDataclass(PyRotation3d);
 
-  py::class_<Rigid3d> PyRigid3d(m, "Rigid3d");
+  py::class_ext_<Rigid3d> PyRigid3d(m, "Rigid3d");
   PyRigid3d.def(py::init<>())
       .def(py::init<const Eigen::Quaterniond&, const Eigen::Vector3d&>())
       .def(py::init([](const Eigen::Matrix3x4d& matrix) {
@@ -87,7 +102,7 @@ void BindGeometry(py::module& m) {
   py::implicitly_convertible<py::array, Rigid3d>();
   MakeDataclass(PyRigid3d);
 
-  py::class_<Sim3d> PySim3d(m, "Sim3d");
+  py::class_ext_<Sim3d> PySim3d(m, "Sim3d");
   PySim3d.def(py::init<>())
       .def(
           py::init<double, const Eigen::Quaterniond&, const Eigen::Vector3d&>())

pycolmap/pybind11_extension.h

@@ -96,6 +96,38 @@ struct type_caster<std::vector<Eigen::Matrix<Scalar, Size, 1>>> {
 
 }  // namespace detail
 
+template <typename type_, typename... options>
+class class_ext_ : public class_<type_, options...> {
+ public:
+  using Parent = class_<type_, options...>;
+  using Parent::class_;  // inherit constructors
+  using type = type_;
+
+  template <typename C, typename D, typename... Extra>
+  class_ext_& def_readwrite(const char* name, D C::*pm, const Extra&... extra) {
+    static_assert(
+        std::is_same<C, type>::value || std::is_base_of<C, type>::value,
+        "def_readwrite() requires a class member (or base class member)");
+    cpp_function fget([pm](type&c) -> D& { return c.*pm; }, is_method(*this)),
+        fset([pm](type&c, const D&value) { c.*pm = value; }, is_method(*this));
+    this->def_property(
+        name, fget, fset, return_value_policy::reference_internal, extra...);
+    return *this;
+  }
+
+  template <typename... Args>
+  class_ext_& def(Args&&... args) {
+    Parent::def(std::forward<Args>(args)...);
+    return *this;
+  }
+
+  template <typename... Args>
+  class_ext_& def_property(Args&&... args) {
+    Parent::def_property(std::forward<Args>(args)...);
+    return *this;
+  }
+};
+
 // Fix long-standing bug https://github.com/pybind/pybind11/issues/4529
 // TODO(sarlinpe): remove when https://github.com/pybind/pybind11/pull/4972
 // appears in the next release of pybind11.

pycolmap/scene/point2D.h

@@ -45,7 +45,7 @@ void BindPoint2D(py::module& m) {
         return repr;
       });
 
-  py::class_<Point2D, std::shared_ptr<Point2D>> PyPoint2D(m, "Point2D");
+  py::class_ext_<Point2D, std::shared_ptr<Point2D>> PyPoint2D(m, "Point2D");
   PyPoint2D.def(py::init<>())
       .def(py::init<const Eigen::Vector2d&, size_t>(),
            "xy"_a,

pycolmap/scene/point3D.h

@@ -26,7 +26,7 @@ void BindPoint3D(py::module& m) {
                std::to_string(self.size()) + ")";
       });
 
-  py::class_<Point3D, std::shared_ptr<Point3D>> PyPoint3D(m, "Point3D");
+  py::class_ext_<Point3D, std::shared_ptr<Point3D>> PyPoint3D(m, "Point3D");
   PyPoint3D.def(py::init<>())
       .def_readwrite("xyz", &Point3D::xyz)
       .def_readwrite("color", &Point3D::color)