From 133b53011666737c39c7a9d413f62ee710892af4 Mon Sep 17 00:00:00 2001
From: Patrick Cleeve <pcleeve@gmail.com>
Date: Fri, 6 Sep 2024 12:26:03 +1000
Subject: [PATCH] [feat] add adaptive alignment

---
 autolamella/workflows/core.py | 17 ++++++++++++++---
 autolamella/workflows/ui.py   |  4 +++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/autolamella/workflows/core.py b/autolamella/workflows/core.py
index 9b1e368..105a758 100644
--- a/autolamella/workflows/core.py
+++ b/autolamella/workflows/core.py
@@ -51,7 +51,8 @@
     LamellaBottomEdge,
     detect_features,
     DetectedFeatures,
-    VolumeBlockCentre
+    VolumeBlockCentre,
+    AdaptiveLamellaCentre,
 )
 from autolamella.ui.AutoLamellaUI import AutoLamellaUI
 from fibsem import config as fcfg
@@ -372,6 +373,11 @@ def mill_lamella(
 
     # align sem using ml and beam shift
     if settings.protocol["options"].get("use_sem_beam_shift_alignment", False):
+
+        # get adaptive model # TODO: generalise this better once models are finalised
+        checkpoint = settings.protocol.get("adaptive_polish", {}).get("model_path", None)
+
+        # align sem using ml and beam shift
         lamella = align_feature_beam_shift(
             microscope,
             settings,
@@ -380,6 +386,8 @@ def mill_lamella(
             validate=validate,
             beam_type=BeamType.ELECTRON,
             hfw=settings.image.hfw,
+            feature=AdaptiveLamellaCentre(),
+            checkpoint=checkpoint,
         )
 
 
@@ -801,7 +809,8 @@ def align_feature_beam_shift(microscope: FibsemMicroscope,
                             validate: bool, 
                             beam_type: BeamType = BeamType.ELECTRON,
                             hfw: float = fcfg.REFERENCE_HFW_MEDIUM,
-                            feature: Feature = LamellaCentre()) -> Lamella:
+                            feature: Feature = LamellaCentre(), 
+                            checkpoint: str = None) -> Lamella:
     """Align the feature to the centre of the image using the beamshift."""
 
     # bookkeeping
@@ -818,7 +827,9 @@ def align_feature_beam_shift(microscope: FibsemMicroscope,
     set_images_ui(parent_ui, eb_image, ib_image)
 
     # detect
-    det = update_detection_ui(microscope, settings, features, parent_ui, validate, msg=lamella.info, position=lamella.state.microscope_state.stage_position)
+    det = update_detection_ui(microscope, settings, features, parent_ui, validate, 
+                              msg=lamella.info, position=None, 
+                              checkpoint=checkpoint)
 
     # TODO: add movement modes; stable move, vertical move, beam shift
 
diff --git a/autolamella/workflows/ui.py b/autolamella/workflows/ui.py
index 0b73dda..3eebd67 100644
--- a/autolamella/workflows/ui.py
+++ b/autolamella/workflows/ui.py
@@ -102,7 +102,8 @@ def _update_mill_stages_ui(
         time.sleep(0.5)
 
 def update_detection_ui(
-    microscope, settings, features, parent_ui, validate: bool, msg: str = "Lamella", position: FibsemStagePosition = None,
+    microscope, settings, features, parent_ui, validate: bool, msg: str = "Lamella", 
+    position: FibsemStagePosition = None, checkpoint: str = None
 ) -> DetectedFeatures:
     feat_str = ", ".join([f.name for f in features])
     update_status_ui(parent_ui, f"{msg}: Detecting Features ({feat_str})...")
@@ -112,6 +113,7 @@ def update_detection_ui(
         settings=settings,
         features=features,
         point=position,
+        checkpoint=checkpoint,
     )
 
     if validate: