Merging autotracker

mantis/acquisition/AcquisitionSettings.py

@@ -1,3 +1,4 @@
+import copy
 import warnings
 
 from dataclasses import field
@@ -36,10 +37,12 @@ class PositionSettings:
     position_labels: List[str] = field(default_factory=list)
     num_positions: int = field(init=False, default=0)
     well_ids: List[str] = field(init=False, default_factory=list)
+    xyz_positions_shift: list = field(init=False, default_factory=list)
 
     def __post_init__(self):
         assert len(self.xyz_positions) == len(self.position_labels)
         self.num_positions = len(self.xyz_positions)
+        self.xyz_positions_shift = copy.deepcopy(self.xyz_positions)
 
         try:
             # Look for "'A1-Site_0', 'H12-Site_1', ... " format
@@ -141,6 +144,7 @@ class MicroscopeSettings:
     use_o3_refocus: bool = False
     o3_refocus_config: Optional[ConfigSettings] = None
     o3_refocus_interval_min: Optional[int] = None
+    autotracker_config: Optional[ConfigSettings] = None
 
 
 @dataclass
@@ -193,3 +197,26 @@ def __post_init__(self):
             attr_val = getattr(self, attr)
             if attr_val is not None:
                 setattr(self, attr, round(attr_val, 1))
+
+
+@dataclass
+class AutotrackerSettings:
+    tracking_method: Literal['phase_cross_correlation', 'template_matching', 'multi_otsu']
+    tracking_interval: Optional[int] = 1
+    scale_yx: Optional[float] = 1.0
+    shift_limit: Optional[Union[Tuple[float, float, float], 'None']] = None
+    device: Optional[str] = 'cpu'
+    zyx_dampening_factor: Optional[Union[Tuple[float, float, float], None]] = None
+    # TODO: maybe do the ROI like in the ls_microscope_settings
+    template_roi_zyx: Optional[Tuple[int, int, int]] = None
+    template_channel: Optional[str] = None
+
+    @validator("tracking_method")
+    def check_tracking_method_options(cls, v):
+        # Check if template matching options are provided and are not None
+        if v == 'template_matching':
+            if not all([cls.template_roi_zyx, cls.template_channel]):
+                raise ValueError(
+                    'template_roi_zyx and template_channel must be provided for template matching'
+                )
+            return v

mantis/acquisition/acq_engine.py

@@ -31,6 +31,7 @@
     SliceSettings,
     MicroscopeSettings,
     AutoexposureSettings,
+    AutotrackerSettings,
 )
 from mantis.acquisition.hook_functions.pre_hardware_hook_functions import (
     log_preparing_acquisition,
@@ -49,7 +50,7 @@
     check_ls_acq_finished,
 )
 
-# isort: on
+from mantis.acquisition.autotracker import autotracker_hook_fn
 
 
 # Define constants
@@ -102,6 +103,7 @@ def __init__(
         self._slice_settings = SliceSettings()
         self._microscope_settings = MicroscopeSettings()
         self._autoexposure_settings = None
+        self._autotracker_settings = None
         self._z0 = None
         self.headless = False if mm_app_path is None else True
         self.type = 'light-sheet' if self.headless else 'label-free'
@@ -173,6 +175,10 @@ def microscope_settings(self):
     def autoexposure_settings(self):
         return self._autoexposure_settings
 
+    @property
+    def autotracker_settings(self):
+        return self._autotracker_settings
+
     @channel_settings.setter
     def channel_settings(self, settings: ChannelSettings):
         logger.debug(
@@ -204,6 +210,16 @@ def autoexposure_settings(self, settings: AutoexposureSettings):
         )
         self._autoexposure_settings = settings
 
+    @autotracker_settings.setter
+    def autotracker_settings(self, settings: AutotrackerSettings):
+        if settings is None:
+            logger.debug('Autotracker settings are not provided')
+        else:
+            logger.debug(
+                f"{self.type.capitalize()} acquisition will have the following settings:{asdict(settings)}"
+            )
+        self._autotracker_settings = settings
+
     def setup(self):
         """
         Apply acquisition settings as specified by the class properties
@@ -276,6 +292,7 @@ def reset(self):
                 )
 
 
+# TODO: check the enable_ls_acq and enable_lf_acq work independently
 class MantisAcquisition(object):
     """
     Acquisition class for simultaneous label-free and light-sheet acquisition on
@@ -336,6 +353,8 @@ def __init__(
         self._lf_acq_obj = None
         self._ls_acq_obj = None
 
+        globals.demo_run = demo_run
+
         if not enable_lf_acq or not enable_ls_acq:
             raise Exception('Disabling LF or LS acquisition is not currently supported')
 
@@ -453,6 +472,8 @@ def update_position_settings(self):
                 xyz_positions=xyz_positions,
                 position_labels=position_labels,
             )
+            self.position_settings.xyz_positions_shift = deepcopy(xyz_positions)
+
         else:
             logger.debug('Position list is already populated and will not be updated')
 
@@ -671,6 +692,12 @@ def setup_autoexposure(self):
                         )
                     )
 
+    def update_position_autotracker(self):
+        # Update the position list from the backup
+        self.position_settings.xyz_positions = deepcopy(
+            self.position_settings.xyz_positions_shift
+        )
+
     def go_to_position(self, position_index: int):
         # Move slowly for short distances such that autofocus can stay engaged.
         # Autofocus typically fails when moving long distances, so we can move
@@ -938,6 +965,18 @@ def run_autoexposure(
                         f'Autoexposure method {method} is not yet implemented.'
                     )
 
+    def run_autotracker(
+        self,
+        acq: BaseChannelSliceAcquisition,
+        well_id: str,
+        method: str = 'manual',
+    ):
+        logging.debug('running autotracker')
+        if not any(acq.channel_settings.use_autoexposure):
+            return
+        # TODO: implement autotracker
+        microscope_operations.autotracker(acq.mmc, acq.autoexposure_settings)
+
     def setup(self):
         """
         Setup the mantis acquisition. This method sets up the label-free
@@ -984,7 +1023,21 @@ def acquire(self):
                 start_daq_counters, [self._lf_z_ctr_task, self._lf_channel_ctr_task]
             )
         lf_post_hardware_hook_fn = log_acquisition_start
-        lf_image_saved_fn = check_lf_acq_finished
+
+        # TODO: implement logic for the autotracker_img_saved_hook_fn
+        if self.lf_acq.microscope_settings.autotracker_config is not None:
+            lf_image_saved_fn = partial(
+                autotracker_hook_fn,
+                'lf',
+                self.lf_acq.autotracker_settings,
+                self._position_settings,
+                self.lf_acq.microscope_settings.autotracker_config,
+                self.lf_acq.slice_settings,
+                self._acq_dir,
+            )
+        else:
+            logger.info('No autotracker config found. Using default image saved hook')
+            lf_image_saved_fn = check_lf_acq_finished
 
         # define LF acquisition
         self._lf_acq_obj = Acquisition(
@@ -1014,7 +1067,19 @@ def acquire(self):
                 self.ls_acq.channel_settings.channels,
             )
             ls_post_camera_hook_fn = partial(start_daq_counters, [self._ls_z_ctr_task])
-        ls_image_saved_fn = check_ls_acq_finished
+
+        # TODO: implement logic for the autotracker_img_saved_hook_fn
+        if self.ls_acq.microscope_settings.autotracker_config is not None:
+            ls_image_saved_fn = partial(
+                autotracker_hook_fn,
+                'ls',
+                self.ls_acq.autotracker_settings,
+                self.ls_acq.slice_settings,
+                self._acq_dir,
+            )
+        else:
+            logger.info('No autotracker config found. Using default image saved hook')
+            ls_image_saved_fn = check_ls_acq_finished
 
         # define LS acquisition
         self._ls_acq_obj = Acquisition(
@@ -1050,8 +1115,23 @@ def acquire(self):
 
                 # move to the given position
                 if p_label != previous_position_label:
+                    # Check if autotracker is on either arm
+                    if (
+                        self.ls_acq.microscope_settings.autotracker_config is not None
+                        or self.lf_acq.microscope_settings.autotracker_config is not None
+                    ):
+                        # TODO: Should we get the corods from the csv file or the modified xyz_positions_shifts
+                        logger.debug('Updating the positions for autotracker')
+                        logger.debug(
+                            'Previous position: %f,%f ',
+                            *self.position_settings.xyz_positions[p_idx][0:2],
+                        )
+                        self.update_position_autotracker()
                     self.go_to_position(p_idx)
 
+                    # TODO get the delta shifts
+                    # read the files here and move separately
+
                 # autofocus
                 if self.lf_acq.microscope_settings.use_autofocus:
                     autofocus_success = microscope_operations.autofocus(