Start smart stack parameters in scan workflow

src/openflexure_microscope_server/scan_directories.py

@@ -8,11 +8,12 @@ import shutil
 import threading
 import zipfile
 from datetime import datetime, timedelta
-from typing import Any, Mapping, Optional, Self
+from typing import Annotated, Any, Mapping, Optional, Self
 
 from pydantic import (
     BaseModel,
     ConfigDict,
+    PlainSerializer,
     ValidationError,
     field_serializer,
     field_validator,
@@ -99,6 +100,14 @@ def _coerce_lecacy_scan_data(data: dict) -> dict:
     return data
 
 
+# This is a dictionary of a PyDantic model. It allows ScanData to hold arbitrary
+# workflow settings models during a scan.
+AnyModelOrDict = Annotated[
+    dict | BaseModel,
+    PlainSerializer(lambda v: v.model_dump(), return_type=dict),
+]
+
+
 class ScanData(BaseModel):
     """Data about a scan to be saved to a JSON file in the directory.
 
@@ -147,7 +156,7 @@ class ScanData(BaseModel):
     workflow: str
     """The class name of the workflow Thing."""
 
-    workflow_settings: dict[str, Any]
+    workflow_settings: AnyModelOrDict
     """The settings for this workflow."""
 
     stitching_settings: Optional[StitchingData]

src/openflexure_microscope_server/things/autofocus.py

@@ -33,8 +33,8 @@ class NotStreamingError(RuntimeError):
     """No images captured from stream. The camera is almost certainly not streaming."""
 
 
-class StackParams(BaseModel):
+class SmartStackParams(BaseModel):
-    """A class for holding for stack parameters, and returning computed ones."""
+    """A class for holding for smart stack parameters, and returning computed ones."""
 
     stack_dz: int
     images_to_save: int
@@ -80,7 +80,8 @@ class StackParams(BaseModel):
             )
         if min_images_to_test > MAX_TEST_IMAGE_COUNT:
             raise ValueError(
-                f"Testing with more than {MAX_TEST_IMAGE_COUNT} images is likely to focus on the cover slip, or strike the sample."
+                f"Testing with more than {MAX_TEST_IMAGE_COUNT} images is likely to "
+                "focus on the cover slip, or strike the sample."
             )
         if min_images_to_test % 2 == 0 or min_images_to_test <= 0:
             raise ValueError(
@@ -97,7 +98,7 @@ class StackParams(BaseModel):
         return images_to_save
 
     @model_validator(mode="after")
-    def check_image_limits(self) -> "StackParams":
+    def check_image_limits(self) -> "SmartStackParams":
         """Ensure the number of images to save isn't more than the minimum tested."""
         if self.images_to_save > self.min_images_to_test:
             raise ValueError("Can't save more images than the minimum number tested.")
@@ -452,107 +453,10 @@ class AutofocusThing(lt.Thing):
             "Looping autofocus couldn't converge on a focus location."
         )
 
-    stack_images_to_save: int = lt.setting(default=1)
-    """The number of images to save in a stack.
-
-    Defaults to 1 unless you need to see either side of focus
-    """
-
-    stack_min_images_to_test: int = lt.setting(default=9)
-    """The minimum number of images to capture in a stack.
-
-    This many images are captures and tested for focus, if the focus is not central
-    enough more images may be captured. After new images are captured the number sets
-    the number of images used for checking if focus is central.
-
-    Defaults to 9 which balances reliability and speed.
-    """
-
-    stack_dz: int = lt.setting(default=50)
-    """Distance in steps between images in a z-stack.
-
-    Suggested values:
-
-    * 50 for 60-100x
-    * 100 for 40x
-    * 200 for 20x
-    """
-
-    @lt.action
-    def create_stack_params(
-        self,
-        images_dir: str,
-        autofocus_dz: int,
-        save_resolution: tuple[int, int],
-    ) -> StackParams:
-        """Set up the parameters used for all stacks in a scan.
-
-        :param images_dir: the folder to save all images
-        :param autofocus_dz: the range to autofocus over if a stack fails
-        :param save_resolution: The resolution to save the captures to disk with
-
-        :returns: A StackParams object with the required parameters.
-        """
-        # Coerce min_images_to_test parameter
-        min_images_to_test = self.stack_min_images_to_test
-        if min_images_to_test < MIN_TEST_IMAGE_COUNT:
-            self.logger.warning(
-                f"Cannot test only {min_images_to_test} image(s) as this will fail. "
-                "Setting min images to test to lowest possible value of"
-                f"{MIN_TEST_IMAGE_COUNT}."
-            )
-            min_images_to_test = MIN_TEST_IMAGE_COUNT
-        elif min_images_to_test > MAX_TEST_IMAGE_COUNT:
-            self.logger.warning(
-                f"Testing {min_images_to_test} images will cause defocus. "
-                "Setting min images to test to highest possible value of "
-                f"{MAX_TEST_IMAGE_COUNT}."
-            )
-            min_images_to_test = MAX_TEST_IMAGE_COUNT
-        elif min_images_to_test % 2 == 0:
-            min_images_to_test += 1
-            self.logger.warning(
-                "Minimum number of images to test should be odd, setting to "
-                f"{min_images_to_test}."
-            )
-        # Set the Thing property to the coerced value
-        self.stack_min_images_to_test = min_images_to_test
-
-        # Coerce the images to save parameter to be positive, odd, and less than
-        # min_images_to_save
-        images_to_save = self.stack_images_to_save
-        if images_to_save <= 0:
-            self.logger.warning(
-                "At least 1 images must be saved. Setting images to save to 1."
-            )
-            images_to_save = 1
-        elif images_to_save > min_images_to_test:
-            self.logger.warning(
-                f"Cannot save {images_to_save} images as this above the minimum "
-                f"number to test. Setting images to save to {MAX_TEST_IMAGE_COUNT}."
-            )
-            images_to_save = min_images_to_test
-        elif images_to_save % 2 == 0:
-            images_to_save += 1
-            self.logger.warning(
-                f"Images to save should be odd, setting to {images_to_save}."
-            )
-        # Set the Thing property to the coerced value
-        self.stack_images_to_save = images_to_save
-
-        return StackParams(
-            stack_dz=self.stack_dz,
-            images_to_save=self.stack_images_to_save,
-            min_images_to_test=self.stack_min_images_to_test,
-            autofocus_dz=autofocus_dz,
-            images_dir=images_dir,
-            save_resolution=save_resolution,
-        )
-
     @lt.action
     def run_smart_stack(
         self,
-        stack_parameters: StackParams,
+        stack_parameters: SmartStackParams,
         save_on_failure: bool = False,
         check_turning_points: bool = True,
     ) -> tuple[bool, int]:
@@ -564,7 +468,7 @@ class AutofocusThing(lt.Thing):
         The sharpest image, and optionally images around the sharpest, will be saved
         to the images_dir with their coordinates in the filename.
 
-        :param stack_parameters: A StackParams object containing the required
+        :param stack_parameters: A SmartStackParams object containing the required
             parameters to run a stack.
         :param save_on_failure: Whether to save an image even if no focus was found.
         :param check_turning_points: Whether to check the number of turning points in
@@ -579,7 +483,7 @@ class AutofocusThing(lt.Thing):
         attempt = 0
         while True:
             attempt += 1
-            success, captures, sharpest_id = self.z_stack(
+            success, captures, sharpest_id = self.smart_z_stack(
                 stack_parameters=stack_parameters,
                 check_turning_points=check_turning_points,
             )
@@ -630,7 +534,7 @@ class AutofocusThing(lt.Thing):
         self,
         sharpest_id: int,
         captures: list[CaptureInfo],
-        stack_parameters: StackParams,
+        stack_parameters: SmartStackParams,
     ) -> int:
         """Save the required captures to disk.
 
@@ -640,7 +544,7 @@ class AutofocusThing(lt.Thing):
         :param sharpest_id: the buffer id index of the sharpest image
         :param captures: a list of captures, including file name, image data and
             metadata
-        :param stack_parameters: a StackParams object holding stack parameters
+        :param stack_parameters: a SmartStackParams object holding stack parameters
         """
         sharpest_index = _get_capture_index_by_id(captures, sharpest_id)
         slice_to_save = stack_parameters.slice_to_save(sharpest_index)
@@ -655,19 +559,23 @@ class AutofocusThing(lt.Thing):
         self._cam.clear_buffers()
         return sharpest_index
 
-    def z_stack(
+    def smart_z_stack(
         self,
-        stack_parameters: StackParams,
+        stack_parameters: SmartStackParams,
         check_turning_points: bool,
     ) -> tuple[bool, list[CaptureInfo], int]:
         """Capture a series of images checking that sharpest image central.
 
+        This is part of run_smart_stack. This is the actual z_stackng stacking method
+        called by the action run_smart_stack. The action also handles reseting,
+        autofocussing, and retrying.
+
         The images are separated in z offset by stack_parameters.stack_dz, as they
         are captured the last stack_parameters.min_images_to_test images are checked
         to see if the sharpest image is central enough in the stack. If it is the stack
         completes.
 
-        :param stack_parameters: a StackParams object holding stack parameters
+        :param stack_parameters: a SmartStackParams object holding stack parameters
         :param check_turning_points: Whether to check the number of turning points in
             the sharpnesses of the images in the stack is exactly 1. (May fail with
             thick samples)

src/openflexure_microscope_server/things/scan_workflows.py

@@ -1,11 +1,18 @@
 from typing import Mapping, Optional
 
+from pydantic import BaseModel
+
 import labthings_fastapi as lt
 
 from openflexure_microscope_server.scan_directories import StitchingData
 from openflexure_microscope_server.scan_planners import ScanPlanner, SmartSpiral
 from openflexure_microscope_server.stitching import STITCHING_RESOLUTION
-from openflexure_microscope_server.things.autofocus import AutofocusThing
+from openflexure_microscope_server.things.autofocus import (
+    MAX_TEST_IMAGE_COUNT,
+    MIN_TEST_IMAGE_COUNT,
+    AutofocusThing,
+    SmartStackParams,
+)
 from openflexure_microscope_server.things.background_detect import (
     ChannelDeviationLUV,
 )
@@ -43,7 +50,8 @@ class ScanWorkflow(lt.Thing):
             "Each specific ScanWorkflow must implement a ready property."
         )
 
-    def all_settings(self) -> tuple[dict, Optional[StitchingData]]:
+    # TODO should be a model
+    def all_settings(self, images_dir: str) -> tuple[dict, Optional[StitchingData]]:
         """Return the scan settings and the stitching settings.
 
         - The specific settings for this scan workflow are returned as a dict.
@@ -72,6 +80,21 @@ class ScanWorkflow(lt.Thing):
         )
 
 
+class HistoScanSettingsModel(BaseModel):
+    """The settings including needed for running a HistoScan.
+
+    This includes settings caluclated when starting. This will be serialised in
+    ScanData.
+    """
+
+    overlap: float
+    max_dist: int
+    dx: int
+    dy: int
+    skip_background: bool
+    smart_stack_prarams: SmartStackParams
+
+
 class HistoScanWorkflow(ScanWorkflow):
     # Thing Slots
     _background_detector: ChannelDeviationLUV = lt.thing_slot()
@@ -96,6 +119,34 @@ class HistoScanWorkflow(ScanWorkflow):
     overlap: float = lt.setting(default=0.45)
     """The fraction (0-1) that adjacent images should overlap in x or y."""
 
+    # Stacking settings
+
+    stack_images_to_save: int = lt.setting(default=1)
+    """The number of images to save in a stack.
+
+    Defaults to 1 unless you need to see either side of focus
+    """
+
+    stack_min_images_to_test: int = lt.setting(default=9)
+    """The minimum number of images to capture in a stack.
+
+    This many images are captures and tested for focus, if the focus is not central
+    enough more images may be captured. After new images are captured the number sets
+    the number of images used for checking if focus is central.
+
+    Defaults to 9 which balances reliability and speed.
+    """
+
+    stack_dz: int = lt.setting(default=50)
+    """Distance in steps between images in a z-stack.
+
+    Suggested values:
+
+    * 50 for 60-100x
+    * 100 for 40x
+    * 200 for 20x
+    """
+
     # noqa, scan_name is unused but is needed for equivalence with other workflows.
     def check_before_start(self, scan_name: str) -> None:  # noqa: ARG002
         """Before starting a scan, check that background and camera-stage-mapping are set.
@@ -110,7 +161,7 @@ class HistoScanWorkflow(ScanWorkflow):
             raise RuntimeError("Camera Stage Mapping is not calibrated.")
 
         if self.skip_background:
-            if not self.background_detector.ready:
+            if not self._background_detector.ready:
                 raise RuntimeError(
                     "Background is not set: you need to calibrate background detection."
                 )
@@ -132,7 +183,9 @@ class HistoScanWorkflow(ScanWorkflow):
             return True
         return self._background_detector.ready
 
-    def all_settings(self) -> tuple[dict, StitchingData]:
+    def all_settings(
+        self, images_dir: str
+    ) -> tuple[HistoScanSettingsModel, StitchingData]:
         stitching_settings = StitchingData(
             overlap=self.overlap,
             correlation_resize=STITCHING_RESOLUTION[0] / self.save_resolution[0],
@@ -154,15 +207,21 @@ class HistoScanWorkflow(ScanWorkflow):
             )
             autofocus_dz = 0
 
-        scan_settings = {
+        smart_stack_prarams = self.create_smart_stack_params(
-            "overlap": self.overlap,
+            images_dir=images_dir,
-            "max_dist": self.max_range,
+            autofocus_dz=autofocus_dz,
-            "dx": dx,
+            save_resolution=self.save_resolution,
-            "dy": dy,
+        )
-            "autofocus_dz": autofocus_dz,
+
-            "autofocus_on": bool(autofocus_dz),
+        scan_settings = HistoScanSettingsModel(
-            "skip_background": self.skip_background,
+            overlap=self.overlap,
-        }
+            max_dist=self.max_range,
+            dx=dx,
+            dy=dy,
+            skip_background=self.skip_background,
+            smart_stack_prarams=smart_stack_prarams,
+        )
+        autofocus_dz = (autofocus_dz,)
 
         return scan_settings, stitching_settings
 
@@ -190,18 +249,90 @@ class HistoScanWorkflow(ScanWorkflow):
         # If not use the other stage axes
         return x_move_stage["y"], y_move_stage["x"]
 
-    def pre_scan_routine(self, settings: dict) -> None:
+    def create_smart_stack_params(
-        self._autofocus.looping_autofocus(dz=settings["autofocus_dz"], start="centre")
+        self,
+        images_dir: str,
+        autofocus_dz: int,
+        save_resolution: tuple[int, int],
+    ) -> SmartStackParams:
+        """Set up the parameters used for all stacks in a scan.
 
-    def new_scan_planner(self, settings: dict, position: Mapping[str, int]) -> None:
+        :param images_dir: the folder to save all images
+        :param autofocus_dz: the range to autofocus over if a stack fails
+        :param save_resolution: The resolution to save the captures to disk with
+
+        :returns: A StackSmartParams object with the required parameters.
+        """
+        # Coerce min_images_to_test parameter
+        min_images_to_test = self.stack_min_images_to_test
+        if min_images_to_test < MIN_TEST_IMAGE_COUNT:
+            self.logger.warning(
+                f"Cannot test only {min_images_to_test} image(s) as this will fail. "
+                "Setting min images to test to lowest possible value of"
+                f"{MIN_TEST_IMAGE_COUNT}."
+            )
+            min_images_to_test = MIN_TEST_IMAGE_COUNT
+        elif min_images_to_test > MAX_TEST_IMAGE_COUNT:
+            self.logger.warning(
+                f"Testing {min_images_to_test} images will cause defocus. "
+                "Setting min images to test to highest possible value of "
+                f"{MAX_TEST_IMAGE_COUNT}."
+            )
+            min_images_to_test = MAX_TEST_IMAGE_COUNT
+        elif min_images_to_test % 2 == 0:
+            min_images_to_test += 1
+            self.logger.warning(
+                "Minimum number of images to test should be odd, setting to "
+                f"{min_images_to_test}."
+            )
+        # Set the Thing property to the coerced value
+        self.stack_min_images_to_test = min_images_to_test
+
+        # Coerce the images to save parameter to be positive, odd, and less than
+        # min_images_to_save
+        images_to_save = self.stack_images_to_save
+        if images_to_save <= 0:
+            self.logger.warning(
+                "At least 1 images must be saved. Setting images to save to 1."
+            )
+            images_to_save = 1
+        elif images_to_save > min_images_to_test:
+            self.logger.warning(
+                f"Cannot save {images_to_save} images as this above the minimum "
+                f"number to test. Setting images to save to {MAX_TEST_IMAGE_COUNT}."
+            )
+            images_to_save = min_images_to_test
+        elif images_to_save % 2 == 0:
+            images_to_save += 1
+            self.logger.warning(
+                f"Images to save should be odd, setting to {images_to_save}."
+            )
+        # Set the Thing property to the coerced value
+        self.stack_images_to_save = images_to_save
+
+        return SmartStackParams(
+            stack_dz=self.stack_dz,
+            images_to_save=self.stack_images_to_save,
+            min_images_to_test=self.stack_min_images_to_test,
+            autofocus_dz=autofocus_dz,
+            images_dir=images_dir,
+            save_resolution=save_resolution,
+        )
+
+    def pre_scan_routine(self, settings: HistoScanSettingsModel) -> None:
+        self._autofocus.looping_autofocus(dz=settings.autofocus_dz, start="centre")
+
+    def new_scan_planner(
+        self, settings: HistoScanSettingsModel, position: Mapping[str, int]
+    ) -> None:
         # The initial plan for the scan should be a single x,y position. All future
         # moves will be planned around this point. In future, route planner could
         # have multiple starting positions, each of which will be visited before the
         # scan can end.
         planner_settings = {
-            "dx": settings["dx"],
+            "dx": settings.dx,
-            "dy": settings["dy"],
+            "dy": settings.dy,
-            "max_dist": settings["max_dist"],
+            "max_dist": settings.max_dist,
         }
         return self._planner_cls(
             initial_position=(position["x"], position["y"]),
@@ -209,7 +340,7 @@ class HistoScanWorkflow(ScanWorkflow):
         )
 
     def aquisition_routine(
-        self, settings: dict, xyz_pos: tuple[int, int, int]
+        self, settings: HistoScanSettingsModel, xyz_pos: tuple[int, int, int]
     ) -> tuple[bool, Optional[int]]:
         """Perform aquisition routine. This is run at each scan location.
 
@@ -217,7 +348,7 @@ class HistoScanWorkflow(ScanWorkflow):
             If failed to find focus, returns for the focus z-position.
         """
         # If skipping background, take an image to check if current field of view is background
-        if settings["skip_background"]:
+        if settings.skip_background:
             capture_image, bg_message = self._cam.image_is_sample()
 
         if not capture_image:
@@ -225,10 +356,10 @@ class HistoScanWorkflow(ScanWorkflow):
             self.logger.info(msg)
             return False, None
 
-        save_on_failure = settings["skip_background"]
+        save_on_failure = settings.skip_background
 
         focused, focused_height = self._autofocus.run_smart_stack(
-            stack_parameters=self.stack_params,
+            stack_parameters=settings.smart_stack_params,
             save_on_failure=save_on_failure,
         )
         # An image was captured if we are focussed or we are not skipping background.

src/openflexure_microscope_server/things/smart_scan.py

@@ -253,7 +253,14 @@ class SmartScanThing(lt.Thing):
         # Record starting position so it can be returned to at end of scan.
         starting_position = self._stage.position
 
-        workflow_settings, stitching_settings = workflow.all_settings()
+        images_dir = self.ongoing_scan.images_dir
+        # Type narrowing
+        if images_dir is None:
+            raise RuntimeError("Couldn't run scan, images directory was not created.")
+
+        workflow_settings, stitching_settings = workflow.all_settings(
+            images_dir=images_dir
+        )
 
         # If stitching settings is None then this workflow doesn't support stitching.
         auto_stitch = self.stitch_automatically and stitching_settings is not None
@@ -305,15 +312,11 @@ class SmartScanThing(lt.Thing):
             workflow.pre_scan_routine(self._scan_data)
             self.ongoing_scan.save_scan_data(self._scan_data)
             images_dir = self.ongoing_scan.images_dir
+            # Type narrowing
             if images_dir is None:
                 raise RuntimeError(
                     "Couldn't run scan, images directory was not created."
                 )
-            self._stack_params = self._autofocus.create_stack_params(
-                images_dir=images_dir,
-                autofocus_dz=self.scan_data.autofocus_dz,
-                save_resolution=self.scan_data.save_resolution,
-            )
             self._preview_stitcher = stitching.PreviewStitcher(
                 images_dir,
                 overlap=self.scan_data.overlap,