czbiohub-sf · jeskesen · Aug 9, 2024 · Aug 9, 2024 · Aug 13, 2024 · Aug 13, 2024
diff --git a/CITATION.cff b/CITATION.cff
@@ -25,7 +25,8 @@ authors:
     orcid: 'https://orcid.org/0000-0002-3033-674X'
   - given-names: 'Shalin B.'
     family-names: Mehta
-    affiliation: [email protected]
+    email: [email protected]
+    affiliation: Chan Zuckerberg Biohub San Francisco
     orcid: 'https://orcid.org/0000-0002-2542-3582'
 repository-code: 'https://github.com/czbiohub-sf/shrimPy'
 abstract: >-

diff --git a/examples/acquisition_settings/demo_acquisition_settings.yaml b/examples/acquisition_settings/demo_acquisition_settings.yaml
@@ -6,7 +6,7 @@ time_settings:
   time_interval_s: 5
 
 lf_channel_settings:
-  default_exposure_times_ms: [10, 10]
+  default_exposure_times_ms: [5, 5]
   channel_group: 'Channel-Multiband'
   channels: ['DAPI', 'FITC']
   use_sequencing: True
@@ -30,6 +30,8 @@ lf_microscope_settings:
   channel_sequencing_settings:
     - ['LED', 'Sequence', 'On']
   use_autofocus: False
+  autotracker_config: ['Channel-Multiband', 'DAPI']
+
 
 ls_channel_settings:
   default_exposure_times_ms: [20, 30]
@@ -56,3 +58,14 @@ ls_microscope_settings:
     - ['Camera', 'BitDepth', '11']
   z_sequencing_settings:
     - ['Z', 'UseSequences', 'Yes']
+
+ls_autoexposure_settings:
+  autoexposure_method: 'manual'
+  rerun_each_timepoint: True
+
+autotracker_settings:
+  tracking_method: 'phase_cross_correlation'
+  tracking_interval: 1
+  shift_limit: [30,200,200]
+  device: 'cpu'
+  scale_yx: 0.075
diff --git a/examples/acquisition_settings/example_acquisition_settings.yaml b/examples/acquisition_settings/example_acquisition_settings.yaml
@@ -120,6 +120,8 @@ lf_microscope_settings:
   use_autofocus: True
   autofocus_stage: 'ZDrive'
   autofocus_method: 'PFS'
+  # Autotracker parameters using the label-free channel
+  autotracker_config: ['Channel - LF', 'BF']
 
 # Microscope settings which will be applied when the light-sheet acquisition is
 # initialized
@@ -144,3 +146,48 @@ ls_microscope_settings:
 ls_autoexposure_settings:
   autoexposure_method: 'manual'
   rerun_each_timepoint: True
+
+autotracker_settings:
+  tracking_method: 'phase_cross_correlation'
+  reconstruction: ['phase', 'vs'] ##list of reconstruction methods to apply before tracking
+  tracking_interval: 1
+  shift_limit: [30,200,200]
+  device: 'cpu'
+  scale_yx: 0.075
+  phase_config: # apply phase reconstruction before tracking
+    transfer_function:
+      wavelength_illumination: 0.450
+      yx_pixel_size: 0.1494
+      z_pixel_size: 0.174
+      z_padding: 5
+      index_of_refraction_media: 1.4
+      numerical_aperture_detection: 1.35
+      numerical_aperture_illumination: 0.52
+      invert_phase_contrast: false
+    apply_inverse:
+      reconstruction_algorithm: Tikhonov
+      regularization_strength: 0.01
+      TV_rho_strength: 0.001
+      TV_iterations: 1
+  vs_config:
+    model:
+      class_path: viscy.translation.engine.VSUNet
+      init_args:
+        architecture: fcmae
+        model_config:
+          in_channels: 1
+          out_channels: 2
+          in_stack_depth: 21
+          encoder_blocks: [3, 3, 9, 3]
+          dims: [96, 192, 384, 768]
+          encoder_drop_path_rate: 0.0
+          stem_kernel_size: [7, 4, 4]
+          in_stack_depth: 21
+          decoder_conv_blocks: 2
+          pretraining: False
+          head_conv: True
+          head_conv_expansion_ratio: 4
+          head_conv_pool: False
+      test_time_augmentations: True
+      tta_type: median
+    ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/fcmae-3d/fit_v2/pretrain_end2end/lightning_logs/finetune_VS_end2end_v1_test6_prefetch2_nopersistwork_restart_2/checkpoints/epoch=64-step=24960.ckpt
diff --git a/mantis/acquisition/AcquisitionSettings.py b/mantis/acquisition/AcquisitionSettings.py
@@ -1,7 +1,8 @@
+import copy
 import warnings
 
 from dataclasses import field
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from typing import Dict, List, Literal, Optional, Tuple, Union, Any
 
 import numpy as np
 
@@ -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
@@ -73,6 +76,7 @@ class ChannelSettings:
     # dictionaries with following structure: {well_id: list_of_exposure_times}
     exposure_times_per_well: Dict = field(init=None, default_factory=dict)
     laser_powers_per_well: Dict = field(init=None, default_factory=dict)
+    min_exposure_time: NonNegativeFloat = 0  # in ms
 
     def __post_init__(self):
         self.num_channels = len(self.channels)
@@ -142,6 +146,7 @@ class MicroscopeSettings:
     o3_refocus_config: Optional[ConfigSettings] = None
     o3_refocus_interval_min: Optional[int] = None
     o3_refocus_skip_wells: List[str] = field(default_factory=list)
+    autotracker_config: Optional[ConfigSettings] = None
 
 
 @dataclass
@@ -194,3 +199,29 @@ 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  # TODO: add units
+    scale_yx: Optional[float] = 1.0  # in um per pixel
+    shift_limit: Optional[Union[Tuple[float, float, float], 'None']] = None # in um
+    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
+    reconstruction: Optional[List[str]] = field(default_factory=list)
+    phase_config: Optional[Dict[str, Any]] = field(default_factory=dict)
+    vs_config: Optional[Dict[str, Any]] = field(default_factory=dict)
+
+    @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