From 3c6cd2ba9a6bd950e6f1bbae9c1e9f5363bb3b34 Mon Sep 17 00:00:00 2001
From: Patrick Cleeve <cleeve@delmic.com>
Date: Tue, 11 Feb 2025 10:03:42 +1100
Subject: [PATCH] fix review

---
 src/odemis/acq/drift/__init__.py          |  5 +++-
 src/odemis/acq/milling/millmng.py         | 32 ++++++++---------------
 src/odemis/acq/milling/test/tasks_test.py |  3 ---
 3 files changed, 15 insertions(+), 25 deletions(-)

diff --git a/src/odemis/acq/drift/__init__.py b/src/odemis/acq/drift/__init__.py
index f6db843e3a..52a5236026 100644
--- a/src/odemis/acq/drift/__init__.py
+++ b/src/odemis/acq/drift/__init__.py
@@ -364,7 +364,10 @@ def align_reference_image(
     if (ref_image.ndim != 2 or new_image.ndim != 2 or ref_image.shape != new_image.shape):
         raise ValueError(f"Only equally sized 2D images are supported for alignment. {ref_image.shape}, {new_image.shape}")
 
-    shift_px = MeasureShift(ref_image, new_image, 10)
+    if ref_image.metadata[model.MD_PIXEL_SIZE] != new_image.metadata[model.MD_PIXEL_SIZE]:
+        raise ValueError("The images must have the same pixel size.")
+
+    shift_px = MeasureShift(ref_image, new_image, 2)
 
     pixelsize = ref_image.metadata[model.MD_PIXEL_SIZE]
     shift_m = (shift_px[0] * pixelsize[0], shift_px[1] * pixelsize[1])
diff --git a/src/odemis/acq/milling/millmng.py b/src/odemis/acq/milling/millmng.py
index 014bb49c7a..e0a3f8895f 100644
--- a/src/odemis/acq/milling/millmng.py
+++ b/src/odemis/acq/milling/millmng.py
@@ -49,7 +49,7 @@
 class MillingTaskManager:
     """This class manages running milling tasks."""
 
-    def __init__(self, future: Future, tasks: List[MillingTaskSettings]):
+    def __init__(self, future: Future, tasks: List[MillingTaskSettings], fib_stream: FIBStream):
         """
         :param future: the future that will be executing the task
         :param tasks: The milling tasks to run (in order)
@@ -58,18 +58,8 @@ def __init__(self, future: Future, tasks: List[MillingTaskSettings]):
         self.microscope = model.getComponent(role="fibsem")
         self.tasks = tasks
 
-        # for reference image alignment,
-        self.ion_beam = model.getComponent(role="ion-beam")
-        self.ion_det = model.getComponent(role="se-detector-ion")
-        self.ion_focus = model.getComponent(role="ion-focus")
-
-        self.fib_stream = FIBStream(
-            name="FIB",
-            detector=self.ion_det,
-            dataflow=self.ion_det.data,
-            emitter=self.ion_beam,
-            focuser=self.ion_focus,
-        )
+        # for reference image alignment
+        self.fib_stream = fib_stream
 
         self._future = future
         if future is not None:
@@ -118,7 +108,7 @@ def run_milling(self, settings: MillingTaskSettings):
         imaging_fov = microscope.get_field_of_view(milling_channel)
 
         # error management
-        e, ce = False, False
+        exception = False
         try:
 
             # acquire a reference image at the imaging settings
@@ -171,8 +161,10 @@ def run_milling(self, settings: MillingTaskSettings):
 
         except CancelledError as ce:
             logging.debug(f"Cancelled milling: {ce}")
+            exception = ce
         except Exception as e:
             logging.error(f"Error while milling: {e}")
+            exception = e
         finally:
             # restore imaging state
             microscope.set_beam_current(imaging_current, milling_channel)
@@ -182,10 +174,8 @@ def run_milling(self, settings: MillingTaskSettings):
             microscope.set_active_view(2)
             microscope.clear_patterns()
 
-            if e:
-                raise e
-            if ce:
-                raise ce # future excepts error to be raised
+            if exception:
+                raise exception
         return
 
     def run(self):
@@ -210,14 +200,14 @@ def run(self):
             logging.debug("Stopping because milling was cancelled")
             raise
         except Exception:
-            logging.exception("The milling failed")
+            logging.warning("The milling failed")
             raise
         finally:
             self._future._task_state = FINISHED
 
 
 # TODO: replace with run_milling_tasks_openfibsem
-def run_milling_tasks(tasks: List[MillingTaskSettings]) -> Future:
+def run_milling_tasks(tasks: List[MillingTaskSettings], fib_stream: FIBStream) -> Future:
     """
     Run multiple milling tasks in order.
     :param tasks: List of milling tasks to be executed in order.
@@ -226,7 +216,7 @@ def run_milling_tasks(tasks: List[MillingTaskSettings]) -> Future:
     # Create a progressive future with running sub future
     future = model.ProgressiveFuture()
     # create acquisition task
-    milling_task_manager = MillingTaskManager(future, tasks)
+    milling_task_manager = MillingTaskManager(future, tasks, fib_stream)
     # add the ability of cancelling the future during execution
     future.task_canceller = milling_task_manager.cancel
 
diff --git a/src/odemis/acq/milling/test/tasks_test.py b/src/odemis/acq/milling/test/tasks_test.py
index 0b0fc8e030..1669bf6262 100644
--- a/src/odemis/acq/milling/test/tasks_test.py
+++ b/src/odemis/acq/milling/test/tasks_test.py
@@ -101,6 +101,3 @@ def test_milling_task_settings(self):
         self.assertEqual(milling_task_settings_from_json.patterns[0].depth.value, trench_pattern.depth.value)
         self.assertEqual(milling_task_settings_from_json.patterns[0].spacing.value, trench_pattern.spacing.value)
         self.assertEqual(milling_task_settings_from_json.patterns[0].center.value, trench_pattern.center.value)
-
-    def test_save_load_milling_tasks(self):
-        pass