From 6f8421d9ef959b9c8406b00ac0403d727dcfc8e1 Mon Sep 17 00:00:00 2001
From: SzabolcsGergely <szabi@luxonis.com>
Date: Fri, 13 Dec 2024 16:39:02 +0200
Subject: [PATCH] Add utility functions to set/get rotation matrix

---
 .../src/pipeline/datatype/ImgFrameBindings.cpp    |  4 +++-
 include/depthai/common/ImgTransformations.hpp     | 13 +++++++++++++
 src/pipeline/datatype/ImgTransformations.cpp      | 15 +++++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/bindings/python/src/pipeline/datatype/ImgFrameBindings.cpp b/bindings/python/src/pipeline/datatype/ImgFrameBindings.cpp
index 4c9451a7f..826ff2ed8 100644
--- a/bindings/python/src/pipeline/datatype/ImgFrameBindings.cpp
+++ b/bindings/python/src/pipeline/datatype/ImgFrameBindings.cpp
@@ -158,7 +158,9 @@ void bind_imgframe(pybind11::module& m, void* pCallstack) {
              DOC(dai, ImgTransformation, addPadding))
         .def("addFlipVertical", &ImgTransformation::addFlipVertical, DOC(dai, ImgTransformation, addFlipVertical))
         .def("addFlipHorizontal", &ImgTransformation::addFlipHorizontal, DOC(dai, ImgTransformation, addFlipHorizontal))
-        .def("addRotation", &ImgTransformation::addRotation, py::arg("angle"), py::arg("rotationPoint"), DOC(dai, ImgTransformation, addRotation))
+        .def("addRotation", py::overload_cast<float, dai::Point2f>(&ImgTransformation::addRotation), py::arg("angle"), py::arg("rotationPoint"), DOC(dai, ImgTransformation, addRotation))
+        .def("addRotation", py::overload_cast<std::array<std::array<float, 3>, 3> >(&ImgTransformation::addRotation), py::arg("rotationMatrix"), DOC(dai, ImgTransformation, addRotation, 2))
+        .def("getRotation", &ImgTransformation::getRotation, DOC(dai, ImgTransformation, getRotation))
         .def("addScale", &ImgTransformation::addScale, py::arg("scaleX"), py::arg("scaleY"), DOC(dai, ImgTransformation, addScale))
         .def("remapPointTo", &ImgTransformation::remapPointTo, py::arg("to"), py::arg("point"), DOC(dai, ImgTransformation, remapPointTo))
         .def("remapPointFrom", &ImgTransformation::remapPointFrom, py::arg("to"), py::arg("point"), DOC(dai, ImgTransformation, remapPointFrom))
diff --git a/include/depthai/common/ImgTransformations.hpp b/include/depthai/common/ImgTransformations.hpp
index 40ae2aa10..2616b35d8 100644
--- a/include/depthai/common/ImgTransformations.hpp
+++ b/include/depthai/common/ImgTransformations.hpp
@@ -197,6 +197,19 @@ struct ImgTransformation {
      * @param rotationPoint Point around which to rotate
      */
     ImgTransformation& addRotation(float angle, dai::Point2f rotationPoint);
+
+    /**
+     * Add a rotation transformation around optical center.
+     * @param rotationMatrix Rotation matrix
+     */
+    ImgTransformation& addRotation(std::array<std::array<float, 3>, 3> rotationMatrix);
+
+    /**
+     * Get the rotation matrix around optical center from the transformation matrix.
+     * @return Rotation matrix
+     */
+    std::array<std::array<float, 3>, 3> getRotation();
+
     /**
      * Add a scale transformation.
      * @param scaleX Scale factor in the horizontal direction
diff --git a/src/pipeline/datatype/ImgTransformations.cpp b/src/pipeline/datatype/ImgTransformations.cpp
index eef90507a..10730b4e7 100644
--- a/src/pipeline/datatype/ImgTransformations.cpp
+++ b/src/pipeline/datatype/ImgTransformations.cpp
@@ -341,6 +341,21 @@ ImgTransformation& ImgTransformation::addRotation(float angle, dai::Point2f rota
     addTransformation(translateMatrixInv);
     return *this;
 }
+
+ImgTransformation& ImgTransformation::addRotation(std::array<std::array<float, 3>, 3> rotationMatrix) {
+    auto intrinsicMatrix = getIntrinsicMatrix();
+    auto homography = matmul(intrinsicMatrix, matmul(rotationMatrix, getMatrixInverse(intrinsicMatrix)));
+    addTransformation(homography);
+    return *this;
+}
+
+std::array<std::array<float, 3>, 3> ImgTransformation::getRotation() {
+    auto intrinsicMatrix = getIntrinsicMatrix();
+    auto homography = getMatrix();
+    auto rotation = matmul(getMatrixInverse(intrinsicMatrix), matmul(homography, intrinsicMatrix));
+    return rotation;
+}
+
 ImgTransformation& ImgTransformation::addScale(float scaleX, float scaleY) {
     width *= scaleX;
     height *= scaleY;