From 7c16823f03013e82d9c3b22876570d7d5254f7ab Mon Sep 17 00:00:00 2001
From: niksirbi <niko.sirbiladze@gmail.com>
Date: Thu, 23 May 2024 18:28:11 +0100
Subject: [PATCH] add ability to save midplane transform

---
 .../napari/align_widget.py                    | 27 +++++++++++++++++++
 .../preproc/alignment.py                      | 10 ++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/brainglobe_template_builder/napari/align_widget.py b/brainglobe_template_builder/napari/align_widget.py
index 0c8c96e..e1f3bb9 100644
--- a/brainglobe_template_builder/napari/align_widget.py
+++ b/brainglobe_template_builder/napari/align_widget.py
@@ -1,8 +1,11 @@
+from pathlib import Path
+
 from napari.layers import Image, Labels, Layer, Points
 from napari.utils.notifications import show_info
 from napari.viewer import Viewer
 from qtpy.QtWidgets import (
     QComboBox,
+    QFileDialog,
     QFormLayout,
     QGroupBox,
     QPushButton,
@@ -95,6 +98,16 @@ def _create_align_group(self):
         self.align_image_button.clicked.connect(self._on_align_button_click)
         self.align_groupbox.layout().addRow(self.align_image_button)
 
+        # Add button to save transformed image
+        self.save_transform_button = QPushButton(
+            "Save transform", parent=self.align_groupbox
+        )
+        self.save_transform_button.setEnabled(False)
+        self.save_transform_button.clicked.connect(
+            self._on_save_transform_click
+        )
+        self.align_groupbox.layout().addRow(self.save_transform_button)
+
     def _get_layers_by_type(self, layer_type: Layer) -> list:
         """Return a list of napari layers of a given type."""
         return [
@@ -183,6 +196,8 @@ def _on_align_button_click(self):
         self.viewer.layers[aligned_mask_name].visible = False
         # Make aligner object accessible to other methods
         self.aligner = aligner
+        # Enable save transform button
+        self.save_transform_button.setEnabled(True)
 
     def _on_dropdown_selection_change(self):
         # Enable estimate button if mask dropdown has a value
@@ -199,3 +214,15 @@ def _on_dropdown_selection_change(self):
             self.align_image_button.setEnabled(False)
         else:
             self.align_image_button.setEnabled(True)
+
+    def _on_save_transform_click(self):
+        """Save the midplane alignment tranform to a text file."""
+        if not hasattr(self, "aligner"):
+            show_info("Please align the image to the midplane first")
+            return
+        dlg = QFileDialog()
+        dlg.setFileMode(QFileDialog.AnyFile)
+        dlg.setAcceptMode(QFileDialog.AcceptSave)
+        if dlg.exec_():
+            path = dlg.selectedFiles()[0]
+            self.aligner.save_transform(Path(path))
diff --git a/brainglobe_template_builder/preproc/alignment.py b/brainglobe_template_builder/preproc/alignment.py
index 27c1368..1007fc7 100644
--- a/brainglobe_template_builder/preproc/alignment.py
+++ b/brainglobe_template_builder/preproc/alignment.py
@@ -1,4 +1,5 @@
 from itertools import product
+from pathlib import Path
 from typing import Literal
 
 import numpy as np
@@ -193,9 +194,16 @@ def transform_image(self, image: np.ndarray = None):
             self._compute_transform()
         if image is None:
             image = self.image
-        self.transformed_image = apply_transform(image, self.transform)
+        transformed_image = apply_transform(image, self.transform)
+        self.transformed_image = transformed_image.astype(image.dtype)
         return self.transformed_image
 
+    def save_transform(self, dest_path: Path):
+        """Save the midplane alignment transform to a text file."""
+        if not hasattr(self, "transform"):
+            raise ValueError("Please align the image to the midplane first")
+        np.savetxt(dest_path, self.transform)
+
     def label_halves(self, image: np.ndarray) -> np.ndarray:
         """Label each half of the image along the symmetry axis with different
         integer values, to help diagnose issues with the splitting process.