From e04fb15f143bf768b6f675f5da5598db8dfa9a04 Mon Sep 17 00:00:00 2001
From: Igor Tatarnikov <61896994+IgorTatarnikov@users.noreply.github.com>
Date: Mon, 9 Dec 2024 17:43:02 +0000
Subject: [PATCH] Fix bug in rotation code (#64)

* Fixed rotation bug

* Reset the grid to display annotations correctly

* Add test for rotation
---
 .../registration_widget.py                    | 24 +++++++++++--------
 brainglobe_registration/utils/utils.py        |  2 +-
 tests/test_registration_widget.py             |  1 +
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/brainglobe_registration/registration_widget.py b/brainglobe_registration/registration_widget.py
index 4d0f367..131e71d 100644
--- a/brainglobe_registration/registration_widget.py
+++ b/brainglobe_registration/registration_widget.py
@@ -508,37 +508,38 @@ def _on_adjust_atlas_rotation(self, pitch: float, yaw: float, roll: float):
 
         # Create the rotation matrix
         roll_matrix = active_matrix_from_angle(0, np.deg2rad(roll))
-        pitch_matrix = active_matrix_from_angle(2, np.deg2rad(pitch))
         yaw_matrix = active_matrix_from_angle(1, np.deg2rad(yaw))
+        pitch_matrix = active_matrix_from_angle(2, np.deg2rad(pitch))
 
-        rot_matrix = roll_matrix @ pitch_matrix @ yaw_matrix
+        # Combine rotation matrices
+        rotation_matrix = yaw_matrix @ pitch_matrix @ roll_matrix
 
         full_matrix = np.eye(4)
-        full_matrix[:-1, :-1] = rot_matrix
+        full_matrix[:3, :3] = rotation_matrix
 
         # Translate the origin to the center of the image
-        origin = np.asarray(self._atlas.reference.shape) // 2
+        origin = np.asarray(self._atlas.reference.shape) / 2
         translate_matrix = np.eye(4)
-        translate_matrix[:-1, -1] = origin
+        translate_matrix[:-1, -1] = -origin
 
         bounding_box = calculate_rotated_bounding_box(
             self._atlas.reference.shape, full_matrix
         )
-        new_translation = np.asarray(bounding_box) // 2
+        new_translation = np.asarray(bounding_box) / 2
         post_rotate_translation = np.eye(4)
-        post_rotate_translation[:3, -1] = -new_translation
+        post_rotate_translation[:3, -1] = new_translation
 
         # Combine the matrices. The order of operations is:
         # 1. Translate the origin to the center of the image
         # 2. Rotate the image
         # 3. Translate the origin back to the top left corner
         transform_matrix = (
-            translate_matrix @ full_matrix @ post_rotate_translation
+            post_rotate_translation @ full_matrix @ translate_matrix
         )
 
         self._atlas_data_layer.data = dask_affine_transform(
             self._atlas.reference,
-            transform_matrix,
+            np.linalg.inv(transform_matrix),
             order=2,
             output_shape=bounding_box,
             output_chunks=(2, bounding_box[1], bounding_box[2]),
@@ -546,14 +547,17 @@ def _on_adjust_atlas_rotation(self, pitch: float, yaw: float, roll: float):
 
         self._atlas_annotations_layer.data = dask_affine_transform(
             self._atlas.annotation,
-            transform_matrix,
+            np.linalg.inv(transform_matrix),
             order=0,
             output_shape=bounding_box,
             output_chunks=(2, bounding_box[1], bounding_box[2]),
         )
 
         # Resets the viewer grid to update the grid to the new atlas
+        # The grid is disabled and re-enabled to force the grid to update
         self._viewer.reset_view()
+        self._viewer.grid.enabled = False
+        self._viewer.grid.enabled = True
 
         worker = self.compute_atlas_rotation(self._atlas_data_layer.data)
         worker.returned.connect(self.set_atlas_layer_data)
diff --git a/brainglobe_registration/utils/utils.py b/brainglobe_registration/utils/utils.py
index 858ebab..1f17cfd 100644
--- a/brainglobe_registration/utils/utils.py
+++ b/brainglobe_registration/utils/utils.py
@@ -150,7 +150,7 @@ def calculate_rotated_bounding_box(
         ]
     )
 
-    transformed_corners = rotation_matrix @ corners.T
+    transformed_corners = np.dot(rotation_matrix, corners.T)
     min_corner = np.min(transformed_corners, axis=1)
     max_corner = np.max(transformed_corners, axis=1)
 
diff --git a/tests/test_registration_widget.py b/tests/test_registration_widget.py
index 828c4aa..e7309d9 100644
--- a/tests/test_registration_widget.py
+++ b/tests/test_registration_widget.py
@@ -164,6 +164,7 @@ def test_scale_moving_image(
         (45, 0, 0, (150, 150, 114)),
         (0, 45, 0, (174, 80, 174)),
         (0, 0, 45, (132, 137, 137)),
+        (0, 90, 90, (80, 114, 132)),
     ],
 )
 def test_on_adjust_atlas_rotation(