diff --git a/picamera_coverage.zip b/picamera_coverage.zip index 0e85922e..bdf5503e 100644 Binary files a/picamera_coverage.zip and b/picamera_coverage.zip differ diff --git a/src/openflexure_microscope_server/things/autofocus.py b/src/openflexure_microscope_server/things/autofocus.py index 9001601b..0ee8d32b 100644 --- a/src/openflexure_microscope_server/things/autofocus.py +++ b/src/openflexure_microscope_server/things/autofocus.py @@ -7,6 +7,7 @@ of images (a 'z-stack'). See repository root for licensing information. """ +import enum import logging import os import time @@ -15,12 +16,12 @@ from types import TracebackType from typing import Literal, Mapping, Optional, Self, Sequence import numpy as np -from pydantic import BaseModel, computed_field, field_validator, model_validator +from pydantic import BaseModel, Field, computed_field, field_validator, model_validator import labthings_fastapi as lt from labthings_fastapi.types.numpy import NDArray -from .camera import BaseCamera +from .camera import BaseCamera, CaptureParams from .stage import BaseStage LOGGER = logging.getLogger(__name__) @@ -33,20 +34,37 @@ class NotStreamingError(RuntimeError): """No images captured from stream. The camera is almost certainly not streaming.""" -class SmartStackParams(BaseModel): - """A class for holding for smart stack parameters, and returning computed ones.""" +class SharpnessMethod(enum.Enum): + """The possible SharpnessMethods for autofocus.""" + + JPEG = enum.auto() + + +class AutofocusParams(BaseModel): + """A class for running autofocus routines.""" + + dz: int + sharpness_method: SharpnessMethod = SharpnessMethod.JPEG + + +class StackOrigin(enum.Enum): + """The position of the current location in the stack.""" + + START = enum.auto() + CENTER = enum.auto() + END = enum.auto() + + +class StackParams(BaseModel): + """A class for holding stack parameters, and returning computed ones.""" stack_dz: int - images_to_save: int - min_images_to_test: int - autofocus_dz: int - images_dir: str - save_resolution: tuple[int, int] + images_to_save: int = Field(gt=0) # Using docstrings under variables as this is how pdoc would expect # attributed to be documented - settling_time: float = 0.3 + settling_time: float = Field(default=0.3, ge=0) """Time (in seconds) between moving and capturing an image""" backlash_correction: int = 250 @@ -54,6 +72,25 @@ class SmartStackParams(BaseModel): Distance (in steps) to overshoot a move and then undo, to account for backlash """ + origin: StackOrigin = StackOrigin.START + """Where the stack is positioned relative to the current z position.""" + + +class SmartStackParams(StackParams): + """A class for holding smart stack parameters, and returning computed ones.""" + + min_images_to_test: int + + save_on_failure: bool = False + """Whether to save an image even if no focus was found.""" + + check_turning_points: bool = True + """ + 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) + """ + stack_height_limit: int = 15 """ How many images can be appended to the stack after the predicted peak to test @@ -457,8 +494,8 @@ class AutofocusThing(lt.Thing): def run_smart_stack( self, stack_parameters: SmartStackParams, - save_on_failure: bool = False, - check_turning_points: bool = True, + capture_parameters: CaptureParams, + autofocus_parameters: AutofocusParams, ) -> tuple[bool, int]: """Run a smart stack. @@ -470,10 +507,6 @@ class AutofocusThing(lt.Thing): :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 - the sharpnesses of the images in the stack is exactly 1. (May fail with - thick samples) :returns: A tuple containing: @@ -485,31 +518,29 @@ class AutofocusThing(lt.Thing): attempt += 1 success, captures, sharpest_id = self.smart_z_stack( stack_parameters=stack_parameters, - check_turning_points=check_turning_points, + check_turning_points=stack_parameters.check_turning_points, ) - if success: - break - - if attempt >= stack_parameters.max_attempts: + if success or attempt >= stack_parameters.max_attempts: break # The z position of the first images in the previous attempt. initial_z_pos = captures[0].position["z"] # If a stack is not successful, move to the start and autofocus try: - self.reset_stack(initial_z_pos, stack_parameters.autofocus_dz) + self.reset_stack(initial_z_pos, autofocus_parameters.dz) except NoFocusFoundError: break # Save stack_parameters.image_to_save images centred on the sharpest capture. # If the smart_stack failed the exact number of images saved may not be # stack_parameters.image_to_save - if success or save_on_failure: + if success or stack_parameters.save_on_failure: self.save_stack( sharpest_id=sharpest_id, captures=captures, stack_parameters=stack_parameters, + capture_parameters=capture_parameters, ) return success, _get_capture_by_id(captures, sharpest_id).position["z"] @@ -535,6 +566,7 @@ class AutofocusThing(lt.Thing): sharpest_id: int, captures: list[CaptureInfo], stack_parameters: SmartStackParams, + capture_parameters: CaptureParams, ) -> int: """Save the required captures to disk. @@ -552,8 +584,8 @@ class AutofocusThing(lt.Thing): # Loop through the range, saving each capture to disk for capture in captures[slice_to_save]: self._cam.save_from_memory( - jpeg_path=os.path.join(stack_parameters.images_dir, capture.filename), - save_resolution=stack_parameters.save_resolution, + jpeg_path=os.path.join(capture_parameters.images_dir, capture.filename), + save_resolution=capture_parameters.save_resolution, buffer_id=capture.buffer_id, ) self._cam.clear_buffers() @@ -712,6 +744,79 @@ class AutofocusThing(lt.Thing): return "success", capture_id + @lt.action + def run_basic_stack( + self, + stack_parameters: StackParams, + capture_parameters: CaptureParams, + ) -> tuple[int, list[int]]: + """Capture a simple z-stack with no focus testing or fitting. + + This performs a fixed stack of images spaced by `stack_dz`, + saving all images captured. No sharpness testing, restart + logic, or autofocus is performed. + + :param stack_parameters: StackParams defining stack spacing, + image count and backlash correction. + :param capture_parameters: CaptureParams defining save + resolution, images directory. + + :returns: + - Final z position + - List of z positions captured + """ + captures: list[CaptureInfo] = [] + z_positions: list[int] = [] + + total_range = stack_parameters.stack_dz * (stack_parameters.images_to_save - 1) + + # Determine starting offset based on stack origin + if stack_parameters.origin == StackOrigin.CENTER: + target_offset = -total_range // 2 + elif stack_parameters.origin == StackOrigin.END: + target_offset = -total_range + else: + target_offset = 0 + + # Apply backlash correction: overshoot and move back + overshoot = target_offset - stack_parameters.backlash_correction + if overshoot != 0: + self._stage.move_relative(z=overshoot) + + # Move back to starting point if needed + if stack_parameters.backlash_correction != 0: + self._stage.move_relative(z=stack_parameters.backlash_correction) + + # Capture images_to_save images + for move_count in range(stack_parameters.images_to_save): + time.sleep(stack_parameters.settling_time) + + capture = self.capture_stack_image( + buffer_max=stack_parameters.images_to_save + ) + captures.append(capture) + z_positions.append(capture.position["z"]) + + # Only move stage if not on the last image + if move_count < stack_parameters.images_to_save - 1: + self._stage.move_relative(z=stack_parameters.stack_dz) + + # Save all captures + for capture in captures: + self._cam.save_from_memory( + jpeg_path=os.path.join( + capture_parameters.images_dir, + capture.filename, + ), + save_resolution=capture_parameters.save_resolution, + buffer_id=capture.buffer_id, + ) + + self._cam.clear_buffers() + + final_z = self._stage.position["z"] + return final_z, z_positions + class NotAPeakError(lt.exceptions.InvocationError): """The data to fit isn't a peak.""" diff --git a/src/openflexure_microscope_server/things/camera/__init__.py b/src/openflexure_microscope_server/things/camera/__init__.py index cb599bdc..f656ed1b 100644 --- a/src/openflexure_microscope_server/things/camera/__init__.py +++ b/src/openflexure_microscope_server/things/camera/__init__.py @@ -15,11 +15,12 @@ import tempfile import time from datetime import datetime from types import TracebackType -from typing import Any, Literal, Mapping, Optional, Self, Tuple +from typing import Annotated, Any, Literal, Mapping, Optional, Self, Tuple import numpy as np import piexif from PIL import Image +from pydantic import BaseModel, Field import labthings_fastapi as lt from labthings_fastapi.types.numpy import NDArray @@ -47,6 +48,17 @@ class CaptureError(RuntimeError): """An error trying to capture from a CameraThing.""" +PositiveInt = Annotated[int, Field(ge=1)] +NonEmptyString = Annotated[str, Field(min_length=1)] + + +class CaptureParams(BaseModel): + """A class for capturing at least a single image.""" + + images_dir: NonEmptyString + save_resolution: tuple[PositiveInt, PositiveInt] + + class NoImageInMemoryError(RuntimeError): """An error called if no image is in memory when accessed.""" diff --git a/src/openflexure_microscope_server/things/scan_workflows.py b/src/openflexure_microscope_server/things/scan_workflows.py index 85198a9e..9bc61846 100644 --- a/src/openflexure_microscope_server/things/scan_workflows.py +++ b/src/openflexure_microscope_server/things/scan_workflows.py @@ -31,13 +31,14 @@ from openflexure_microscope_server.stitching import ( from openflexure_microscope_server.things.autofocus import ( MAX_TEST_IMAGE_COUNT, MIN_TEST_IMAGE_COUNT, + AutofocusParams, AutofocusThing, SmartStackParams, ) from openflexure_microscope_server.things.background_detect import ( ChannelDeviationLUV, ) -from openflexure_microscope_server.things.camera import BaseCamera +from openflexure_microscope_server.things.camera import BaseCamera, CaptureParams from openflexure_microscope_server.things.camera_stage_mapping import CameraStageMapper from openflexure_microscope_server.things.stage import BaseStage from openflexure_microscope_server.ui import PropertyControl, property_control_for @@ -101,7 +102,7 @@ class ScanWorkflow(Generic[SettingModelType], lt.Thing): is returned if it is not possible to stitch the scan. """ raise NotImplementedError( - "Each specific ScanWorkflow must implement a `all_settings`. method." + "Each specific ScanWorkflow must implement a `all_settings` method." ) def pre_scan_routine(self, settings: SettingModelType) -> None: @@ -167,7 +168,24 @@ class ScanWorkflow(Generic[SettingModelType], lt.Thing): ) -class RectGridWorkflow(ScanWorkflow[SettingModelType], Generic[SettingModelType]): +class RectGridSettingsModel(BaseModel): + """Base setting model for all RectGrid workflows.""" + + overlap: float + dx: int + dy: int + capture_params: CaptureParams + autofocus_params: AutofocusParams + + +RectGridSettingModelType = TypeVar( + "RectGridSettingModelType", bound=RectGridSettingsModel +) + + +class RectGridWorkflow( + ScanWorkflow[RectGridSettingModelType], Generic[RectGridSettingModelType] +): """A generic workflow for any scan that captures images on a rectilinear grid.""" # Redefine _csm Thing Slot, as CSM is required for any RectGridWorkflow @@ -239,22 +257,53 @@ class RectGridWorkflow(ScanWorkflow[SettingModelType], Generic[SettingModelType] correlation_resize=STITCHING_RESOLUTION[0] / self.save_resolution[0], ) + def _build_scan_settings(self, base_kwargs: dict) -> RectGridSettingModelType: + """Construct the _settings_model.""" + # Developer Note: This needs to be overridden if the settings model for this + # class contains extra keys. + return self._settings_model(**base_kwargs) + + def all_settings( + self, images_dir: str + ) -> tuple[RectGridSettingModelType, Optional[StitchingSettings]]: + """Return scan settings and the stitching settings. + + :param images_dir: The directory that images are to be written to. + :return: A tuple containing the settings model for this workflow and the + settings model for stitching. + """ + # Developer Note: When subclassing RectGridWorkflow rather than override + # this method first consider overriding _build_scan_settings + stitching_settings = self._get_stitching_settings_model() + dx, dy = self._calc_displacement_from_overlap(self.overlap) + + base_kwargs = { + "overlap": self.overlap, + "dx": dx, + "dy": dy, + "capture_params": CaptureParams( + images_dir=images_dir, save_resolution=self.save_resolution + ), + "autofocus_params": AutofocusParams(dz=self.autofocus_dz), + } + + scan_settings = self._build_scan_settings(base_kwargs) + + return scan_settings, stitching_settings + @lt.property def ready(self) -> bool: """Whether this scanworkflow is ready to start.""" return not self._csm.calibration_required -class HistoScanSettingsModel(BaseModel): +class HistoScanSettingsModel(RectGridSettingsModel): """The settings for a scan with the HistoScanWorkflow. This includes settings calculated when starting. This will be held by smart scan during a scan and serialised to disk. """ - overlap: float - dx: int - dy: int max_dist: int skip_background: bool smart_stack_params: SmartStackParams @@ -358,46 +407,19 @@ class HistoScanWorkflow(RectGridWorkflow[HistoScanSettingsModel]): return True return self._background_detector.ready - def all_settings( - self, images_dir: str - ) -> tuple[HistoScanSettingsModel, StitchingSettings]: - """Return the workflow and stitching settings. - - :param images_dir: The directory that images are to be written to. - :return: A tuple containing the settings model for this workflow and the - settings model for stitching. - """ - stitching_settings = self._get_stitching_settings_model() - dx, dy = self._calc_displacement_from_overlap(self.overlap) - - smart_stack_params = self.create_smart_stack_params( - images_dir=images_dir, - autofocus_dz=self.autofocus_dz, - save_resolution=self.save_resolution, - ) - - scan_settings = HistoScanSettingsModel( - overlap=self.overlap, + def _build_scan_settings(self, base_kwargs: dict) -> HistoScanSettingsModel: + """Construct the SettingModel for all_settings.""" + return HistoScanSettingsModel( + **base_kwargs, max_dist=self.max_range, - dx=dx, - dy=dy, skip_background=self.skip_background, - smart_stack_params=smart_stack_params, + smart_stack_params=self.create_smart_stack_params(), ) - return scan_settings, stitching_settings - def create_smart_stack_params( self, - images_dir: str, - autofocus_dz: int, - save_resolution: tuple[int, int], ) -> SmartStackParams: - """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 + """Set up the parameters used for all smart stacks in a scan. :returns: A StackSmartParams object with the required parameters. """ @@ -452,9 +474,7 @@ class HistoScanWorkflow(RectGridWorkflow[HistoScanSettingsModel]): 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, + save_on_failure=not self.skip_background, ) def pre_scan_routine(self, settings: HistoScanSettingsModel) -> None: @@ -463,7 +483,7 @@ class HistoScanWorkflow(RectGridWorkflow[HistoScanSettingsModel]): :param settings: The settings for this scan as a HistoScanSettingsModel """ self._autofocus.looping_autofocus( - dz=settings.smart_stack_params.autofocus_dz, start="centre" + dz=settings.autofocus_params.dz, start="centre" ) def new_scan_planner( @@ -513,15 +533,14 @@ class HistoScanWorkflow(RectGridWorkflow[HistoScanSettingsModel]): self.logger.info(msg) return False, None - save_on_failure = not settings.skip_background - focus_height: Optional[int] focused, focus_height = self._autofocus.run_smart_stack( stack_parameters=settings.smart_stack_params, - save_on_failure=save_on_failure, + capture_parameters=settings.capture_params, + autofocus_parameters=settings.autofocus_params, ) # An image was captured if we are focussed or we are not skipping background. - imaged = focused or save_on_failure + imaged = focused or settings.smart_stack_params.save_on_failure if not imaged: msg = f"Stack failed at {xyz_pos}. Treating as background." @@ -562,22 +581,16 @@ class HistoScanWorkflow(RectGridWorkflow[HistoScanSettingsModel]): ] -class RegularGridSettingsModel(BaseModel): +class RegularGridSettingsModel(RectGridSettingsModel): """The settings for a scan with a regular grid of dx and dy for x_count, y_count steps. This includes settings calculated when starting. This will be held by smart scan during a scan and serialised to disk. """ - overlap: float - dx: int - dy: int x_count: int y_count: int style: Literal["snake", "raster"] - images_dir: str - autofocus_dz: int - save_resolution: tuple[int, int] class RegularGridWorkflow(RectGridWorkflow[RegularGridSettingsModel]): @@ -592,32 +605,15 @@ class RegularGridWorkflow(RectGridWorkflow[RegularGridSettingsModel]): _planner_cls = RegularGridPlanner _grid_style: Literal["snake", "raster"] - def all_settings( - self, images_dir: str - ) -> tuple[RegularGridSettingsModel, Optional[StitchingSettings]]: - """Return the workflow and stitching settings. - - :param images_dir: The directory that images are to be written to. - :return: A tuple containing the settings model for this workflow and the - settings model for stitching. - """ - stitching_settings = self._get_stitching_settings_model() - dx, dy = self._calc_displacement_from_overlap(self.overlap) - - scan_settings = self._settings_model( - overlap=self.overlap, - dx=dx, - dy=dy, + def _build_scan_settings(self, base_kwargs: dict) -> RegularGridSettingsModel: + """Construct the SettingModel for all_settings.""" + return RegularGridSettingsModel( + **base_kwargs, x_count=self.x_count, y_count=self.y_count, style=self._grid_style, - images_dir=images_dir, - autofocus_dz=self.autofocus_dz, - save_resolution=self.save_resolution, ) - return scan_settings, stitching_settings - def pre_scan_routine(self, settings: RegularGridSettingsModel) -> None: """Perform these steps before starting the scan. @@ -625,7 +621,9 @@ class RegularGridWorkflow(RectGridWorkflow[RegularGridSettingsModel]): :param settings: The settings for this scan as as the relevant SettingsModel type. """ - self._autofocus.looping_autofocus(dz=settings.autofocus_dz, start="centre") + self._autofocus.looping_autofocus( + dz=settings.autofocus_params.dz, start="centre" + ) def new_scan_planner( self, settings: RegularGridSettingsModel, position: Mapping[str, int] @@ -659,9 +657,9 @@ class RegularGridWorkflow(RectGridWorkflow[RegularGridSettingsModel]): """ return self._autofocus_and_capture( xyz_pos=xyz_pos, - dz=settings.autofocus_dz, - images_dir=settings.images_dir, - save_resolution=settings.save_resolution, + dz=settings.autofocus_params.dz, + images_dir=settings.capture_params.images_dir, + save_resolution=settings.capture_params.save_resolution, ) @lt.property diff --git a/tests/unit_tests/test_scan_workflows.py b/tests/unit_tests/test_scan_workflows.py index 38ca6d61..50c27cdc 100644 --- a/tests/unit_tests/test_scan_workflows.py +++ b/tests/unit_tests/test_scan_workflows.py @@ -130,7 +130,7 @@ def test_histo_workflow_settings_generation(histo_workflow, mocker): assert workflow_settings.dx == 123 assert workflow_settings.dy == 456 # And that the input image dir is passed to stack the stack parameter for saving - assert workflow_settings.smart_stack_params.images_dir == "/this/img_dir" + assert workflow_settings.capture_params.images_dir == "/this/img_dir" # A CSM that is "normal" changing from camera maxtrix coordinates (y,x) to normal @@ -196,7 +196,7 @@ def test_histo_pre_scan_routine(histo_workflow, mocker): # Rather than create a whole Setting class, just create a mock with the value # we need set mock_settings = mocker.Mock() - mock_settings.smart_stack_params.autofocus_dz = 1234 + mock_settings.autofocus_params.dz = 1234 # Run the function histo_workflow.pre_scan_routine(mock_settings) # Check the autofocus was run using the mocked slot. @@ -258,6 +258,13 @@ def test_histo_acquisition_not_skipping_background(histo_workflow, mocker, caplo # Mocking for settings as above mock_settings = mocker.Mock() mock_settings.skip_background = False + mock_settings.smart_stack_params = SmartStackParams( + stack_dz=5, + images_to_save=3, + min_images_to_test=3, + save_on_failure=True, + check_turning_points=True, + ) # Set up background detector to say it is empty, this shouldn't stop stacking. histo_workflow._background_detector.image_is_sample.return_value = ( @@ -296,6 +303,13 @@ def test_histo_acquisition_on_sample(histo_workflow, mocker, caplog): # Mocking for settings as above mock_settings = mocker.Mock() mock_settings.skip_background = True + mock_settings.smart_stack_params = SmartStackParams( + stack_dz=5, + images_to_save=3, + min_images_to_test=3, + save_on_failure=False, + check_turning_points=True, + ) # Set up background detector to say there is sample. histo_workflow._background_detector.image_is_sample.return_value = (True, None) diff --git a/tests/unit_tests/test_stack.py b/tests/unit_tests/test_stack.py index 7cc4fe7e..900e58ab 100644 --- a/tests/unit_tests/test_stack.py +++ b/tests/unit_tests/test_stack.py @@ -20,11 +20,14 @@ from openflexure_microscope_server.things.autofocus import ( CaptureInfo, NotAPeakError, SmartStackParams, + StackOrigin, + StackParams, _count_turning_points, _get_capture_by_id, _get_capture_index_by_id, _get_peak_turning_point, ) +from openflexure_microscope_server.things.camera import CaptureParams from openflexure_microscope_server.things.scan_workflows import HistoScanWorkflow RANDOM_GENERATOR = np.random.default_rng() @@ -66,12 +69,7 @@ def test_stack_params_validation(save_ims, extra_ims): # to do automatically in hypothesis. This clamps the number between 3 and 9. min_images_to_test = max(min(save_ims + extra_ims, 9), 3) SmartStackParams( - stack_dz=50, - images_to_save=save_ims, - min_images_to_test=min_images_to_test, - autofocus_dz=2000, - images_dir="/this/is/fake", - save_resolution=(1640, 1232), + stack_dz=50, images_to_save=save_ims, min_images_to_test=min_images_to_test ) @@ -96,9 +94,6 @@ def test_stack_params_not_enough_test_images(save_ims, extra_ims): stack_dz=50, images_to_save=save_ims, min_images_to_test=save_ims + extra_ims, - autofocus_dz=2000, - images_dir="/this/is/fake", - save_resolution=(1640, 1232), ) @@ -114,16 +109,14 @@ def test_stack_params_negative_images_to_save(save_ims, extra_ims): # Depending on the values multiple messages are possible match = ( "(Can't test for focus with fewer than 3 images|" - "Images to save must be positive and odd)" + "Images to save must be positive and odd|" + "Input should be greater than 0)" ) with pytest.raises(ValueError, match=match): SmartStackParams( stack_dz=50, images_to_save=save_ims, min_images_to_test=save_ims + extra_ims, - autofocus_dz=2000, - images_dir="/this/is/fake", - save_resolution=(1640, 1232), ) @@ -147,9 +140,6 @@ def test_even_min_images_to_test(save_ims, extra_ims): stack_dz=50, images_to_save=save_ims, min_images_to_test=save_ims + extra_ims, - autofocus_dz=2000, - images_dir="/this/is/fake", - save_resolution=(1640, 1232), ) @@ -164,16 +154,14 @@ def test_even_images_to_save(save_ims, extra_ims): """ match = ( "(Can't test for focus with fewer than 3 images|" - "Images to save must be positive and odd)" + "Images to save must be positive and odd|" + "Input should be greater than 0)" ) with pytest.raises(ValueError, match=match): SmartStackParams( stack_dz=50, images_to_save=save_ims, min_images_to_test=save_ims + extra_ims, - autofocus_dz=2000, - images_dir="/this/is/fake", - save_resolution=(1640, 1232), ) @@ -183,12 +171,7 @@ def test_computed_stack_params(): Not using hypothesis or we will just copy in the same formulas. """ stack_parameters = SmartStackParams( - stack_dz=50, - images_to_save=5, - min_images_to_test=9, - autofocus_dz=2000, - images_dir="/this/is/fake", - save_resolution=(1640, 1232), + stack_dz=50, images_to_save=5, min_images_to_test=9 ) assert stack_parameters.stack_z_range == 8 * 50 @@ -271,7 +254,17 @@ def autofocus_thing(): @pytest.fixture def histo_scan_workflow(): """Return an autofocus thing connected to a server.""" - return create_thing_without_server(HistoScanWorkflow, mock_all_slots=True) + workflow = create_thing_without_server( + HistoScanWorkflow, + mock_all_slots=True, + ) + + # Minimal CSM setup so all_settings() works + workflow._csm.image_resolution = (1000, 1000) + workflow._csm.calibration_required = False + workflow._csm.convert_image_to_stage_coordinates = lambda x, y: {"x": x, "y": y} + + return workflow def test_create_stack(histo_scan_workflow, caplog): @@ -279,9 +272,7 @@ def test_create_stack(histo_scan_workflow, caplog): initial_min_images_to_test = histo_scan_workflow.stack_min_images_to_test initial_images_to_save = histo_scan_workflow.stack_images_to_save with caplog.at_level(logging.INFO): - stack_params = histo_scan_workflow.create_smart_stack_params( - autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232) - ) + stack_params = histo_scan_workflow.create_smart_stack_params() assert len(caplog.records) == 0 assert histo_scan_workflow.stack_min_images_to_test == initial_min_images_to_test @@ -308,9 +299,7 @@ def test_coercing_stack_test_ims( histo_scan_workflow.stack_min_images_to_test = initial_test_ims with caplog.at_level(logging.WARNING): - stack_params = histo_scan_workflow.create_smart_stack_params( - autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232) - ) + stack_params = histo_scan_workflow.create_smart_stack_params() assert len(caplog.records) == 1 assert str(caplog.records[0].msg).startswith(expected_log_start) @@ -338,9 +327,7 @@ def test_coercing_stack_save_ims( histo_scan_workflow.stack_images_to_save = initial_save_ims with caplog.at_level(logging.WARNING): - stack_params = histo_scan_workflow.create_smart_stack_params( - autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232) - ) + stack_params = histo_scan_workflow.create_smart_stack_params() assert len(caplog.records) == 1 assert str(caplog.records[0].msg).startswith(expected_log_start) @@ -353,10 +340,8 @@ def test_coercing_stack_save_ims( @pytest.mark.parametrize("pass_on", [1, 2, 3, 4]) def test_run_smart_stack(pass_on, histo_scan_workflow, autofocus_thing, mocker): """Test Running smart stack with the stack passing on different attempts.""" - stack_params = histo_scan_workflow.create_smart_stack_params( - autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232) - ) - assert stack_params.max_attempts == 3 + scan_settings, _ = histo_scan_workflow.all_settings(images_dir="dummy") + assert scan_settings.smart_stack_params.max_attempts == 3 # Set up returns from z-stack fake_captures = [ @@ -381,19 +366,19 @@ def test_run_smart_stack(pass_on, histo_scan_workflow, autofocus_thing, mocker): # Run it success, final_z = autofocus_thing.run_smart_stack( - stack_parameters=stack_params, - save_on_failure=False, - check_turning_points=True, + stack_parameters=scan_settings.smart_stack_params, + capture_parameters=scan_settings.capture_params, + autofocus_parameters=scan_settings.autofocus_params, ) # Only passes if the attempt it passes on is less than max attempts - assert success == (pass_on <= stack_params.max_attempts) + assert success == (pass_on <= scan_settings.smart_stack_params.max_attempts) # Final z is the one from the id returned by the stack "pick_me" assert final_z == 555 # smart_z_stack should run up until the time it passes. Running no more than # max_attempts - n_stacks = min(pass_on, stack_params.max_attempts) + n_stacks = min(pass_on, scan_settings.smart_stack_params.max_attempts) assert autofocus_thing.smart_z_stack.call_count == n_stacks # Move absolute should be 1 less time that the number of times z_stack_run assert autofocus_thing._stage.move_absolute.call_count == n_stacks - 1 @@ -417,9 +402,7 @@ def setup_and_run_smart_z_stack( is a list, it will be set as a side effect (and should be a list of tuples of results). If it a tuple (or anything else), it is set as a return value. """ - stack_params = histo_scan_workflow.create_smart_stack_params( - autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232) - ) + stack_params = histo_scan_workflow.create_smart_stack_params() stack_params.settling_time = 0 # Don't settle or tests take forever. autofocus_thing.capture_stack_image = mocker.Mock() @@ -680,3 +663,343 @@ def test_count_turning_points(): assert _count_turning_points(np.array([1, 2, 3, 4, 2, 4, 3, 2, 1])) == 3 # But only one if the dip isn't prominent assert _count_turning_points(np.array([1, 2, 3, 4, 3.8, 4, 3, 2, 1])) == 1 + + +@pytest.fixture +def fake_capture(autofocus_thing): + """Return a fake capture function using the current stage Z position.""" + + def _fake_capture(*_args, **_kwargs): + z = autofocus_thing._stage.position["z"] + return CaptureInfo( + buffer_id=f"id_{z}", + position={"x": 0, "y": 0, "z": z}, + sharpness=1, + ) + + return _fake_capture + + +@pytest.fixture +def fake_move_relative(autofocus_thing): + """Return a fake move_relative function updating the stage Z position.""" + + def _fake_move_relative(z, **_kwargs): + autofocus_thing._stage.position["z"] += z + + return _fake_move_relative + + +def test_run_basic_stack_simple( + autofocus_thing, mocker, fake_capture, fake_move_relative +): + """Basic stack captures the correct number of images at correct Z positions.""" + # Stack parameters: small 3-image stack, 10-step spacing + stack_params = StackParams( + stack_dz=10, + images_to_save=3, + settling_time=0, + backlash_correction=0, + origin=StackOrigin.START, + ) + + # Capture parameters for the test + capture_params = mocker.Mock() + capture_params.images_dir = "dummy" + capture_params.save_resolution = (100, 100) + + # Reset stage position + start_z = 0 + autofocus_thing._stage.position = {"x": 0, "y": 0, "z": start_z} + + # Patch capture and stage movement + autofocus_thing.capture_stack_image = mocker.Mock(side_effect=fake_capture) + autofocus_thing._stage.move_relative = mocker.Mock(side_effect=fake_move_relative) + autofocus_thing._cam.save_from_memory = mocker.Mock() + autofocus_thing._cam.clear_buffers = mocker.Mock() + + final_z, z_positions = autofocus_thing.run_basic_stack( + stack_parameters=stack_params, + capture_parameters=capture_params, + ) + + # Expected Z positions for the stack + expected_z_positions = [ + start_z + i * stack_params.stack_dz for i in range(stack_params.images_to_save) + ] + assert z_positions == expected_z_positions, ( + "Z positions captured do not match expected values" + ) + + # Final Z should be the last captured Z + expected_final_z = expected_z_positions[-1] + assert final_z == expected_final_z, "Final Z position is incorrect" + + # Check that capture_stack_image was called exactly images_to_save times + assert ( + autofocus_thing.capture_stack_image.call_count == stack_params.images_to_save + ), "Incorrect number of captures" + + # Check that stage moved correctly (should match relative increments) + moves = [ + call.kwargs["z"] for call in autofocus_thing._stage.move_relative.call_args_list + ] + expected_moves = [stack_params.stack_dz] * (stack_params.images_to_save - 1) + assert moves == expected_moves, ( + "Stage move_relative calls do not match expected increments" + ) + + +def test_run_basic_stack_center_origin( + autofocus_thing, mocker, fake_capture, fake_move_relative +): + """Stack should shift start position when origin is CENTER. + + CENTER should cause the stage to move down by half the z range before + the stack begins. + """ + stack_params = StackParams( + stack_dz=10, + images_to_save=5, + settling_time=0, + backlash_correction=0, + origin=StackOrigin.CENTER, + ) + + capture_params = mocker.Mock() + capture_params.images_dir = "dummy" + capture_params.save_resolution = (100, 100) + + start_z = 0 + autofocus_thing._stage.position = {"x": 0, "y": 0, "z": start_z} + + # Patch capture and stage movement + autofocus_thing.capture_stack_image = mocker.Mock(side_effect=fake_capture) + autofocus_thing._stage.move_relative = mocker.Mock(side_effect=fake_move_relative) + autofocus_thing._cam.save_from_memory = mocker.Mock() + autofocus_thing._cam.clear_buffers = mocker.Mock() + + # Run stack + final_z, z_positions = autofocus_thing.run_basic_stack( + stack_parameters=stack_params, + capture_parameters=capture_params, + ) + + # Calculate expected starting offset - half of stack range + total_range = stack_params.stack_dz * (stack_params.images_to_save - 1) + expected_offset = -total_range // 2 + + # First move should apply center offset + first_call = autofocus_thing._stage.move_relative.call_args_list[0] + assert first_call.kwargs["z"] == expected_offset, ( + "Center origin offset not applied correctly" + ) + + # Expected Z positions after CENTER offset + expected_z_positions = [ + expected_offset + i * stack_params.stack_dz + for i in range(stack_params.images_to_save) + ] + assert z_positions == expected_z_positions, ( + "Z positions captured do not match expected values" + ) + + # Final Z should be last captured Z + expected_final_z = expected_z_positions[-1] + assert final_z == expected_final_z, "Final Z position is incorrect" + + +def test_run_basic_stack_backlash_applied( + autofocus_thing, mocker, fake_capture, fake_move_relative +): + """Backlash correction should overshoot first, then move back before starting stack.""" + stack_params = StackParams( + stack_dz=10, + images_to_save=3, + settling_time=0, + backlash_correction=50, + origin=StackOrigin.START, + ) + + capture_params = mocker.Mock() + capture_params.images_dir = "dummy" + capture_params.save_resolution = (100, 100) + + start_z = 0 + autofocus_thing._stage.position = {"x": 0, "y": 0, "z": start_z} + + autofocus_thing.capture_stack_image = mocker.Mock(side_effect=fake_capture) + autofocus_thing._stage.move_relative = mocker.Mock(side_effect=fake_move_relative) + autofocus_thing._cam.save_from_memory = mocker.Mock() + autofocus_thing._cam.clear_buffers = mocker.Mock() + + autofocus_thing.run_basic_stack(stack_params, capture_params) + + calls = autofocus_thing._stage.move_relative.call_args_list + + # First move is overshoot for backlash + assert calls[0].kwargs["z"] == -stack_params.backlash_correction, ( + "Backlash overshoot not applied correctly" + ) + + # Second move corrects back to starting point + assert calls[1].kwargs["z"] == stack_params.backlash_correction, ( + "Backlash correction move not applied correctly" + ) + + +def test_run_basic_stack_backlash_and_offset( + autofocus_thing, mocker, fake_capture, fake_move_relative +): + """Backlash correction and stack offset both applied. + + Stack should begin by moving down by half the height of the stack, + plus backlash correction, then move up by backlash correction, then run the basic stack. + """ + stack_params = StackParams( + stack_dz=100, + images_to_save=3, + settling_time=0, + backlash_correction=50, + origin=StackOrigin.CENTER, + ) + + capture_params = mocker.Mock() + capture_params.images_dir = "dummy" + capture_params.save_resolution = (100, 100) + + start_z = 0 + autofocus_thing._stage.position = {"x": 0, "y": 0, "z": start_z} + + autofocus_thing.capture_stack_image = mocker.Mock(side_effect=fake_capture) + autofocus_thing._stage.move_relative = mocker.Mock(side_effect=fake_move_relative) + autofocus_thing._cam.save_from_memory = mocker.Mock() + autofocus_thing._cam.clear_buffers = mocker.Mock() + + autofocus_thing.run_basic_stack(stack_params, capture_params) + + calls = autofocus_thing._stage.move_relative.call_args_list + + # Combined initial offset + backlash overshoot + total_stack_range = stack_params.stack_dz * (stack_params.images_to_save - 1) + expected_first_move = -(total_stack_range // 2 + stack_params.backlash_correction) + assert calls[0].kwargs["z"] == expected_first_move, ( + "Combined overshoot not applied correctly" + ) + + # Backlash correction back to base of stack + assert calls[1].kwargs["z"] == stack_params.backlash_correction, ( + "Backlash correction move not applied correctly" + ) + + # Subsequent moves in stack + expected_stack_moves = [stack_params.stack_dz] * (stack_params.images_to_save - 1) + actual_stack_moves = [c.kwargs["z"] for c in calls[2:]] + assert actual_stack_moves == expected_stack_moves, ( + "Stack moves after backlash/offset not correct" + ) + + +def test_run_basic_stack_end_origin( + autofocus_thing, mocker, fake_capture, fake_move_relative +): + """END origin should shift stack down by full stack height before starting.""" + stack_params = StackParams( + stack_dz=10, + images_to_save=4, + settling_time=0, + backlash_correction=0, + origin=StackOrigin.END, + ) + + capture_params = mocker.Mock() + capture_params.images_dir = "dummy" + capture_params.save_resolution = (100, 100) + + start_z = 0 + autofocus_thing._stage.position = {"x": 0, "y": 0, "z": start_z} + + autofocus_thing.capture_stack_image = mocker.Mock(side_effect=fake_capture) + autofocus_thing._stage.move_relative = mocker.Mock(side_effect=fake_move_relative) + autofocus_thing._cam.save_from_memory = mocker.Mock() + autofocus_thing._cam.clear_buffers = mocker.Mock() + + autofocus_thing.run_basic_stack(stack_params, capture_params) + + # Calculate the stack offset based on StackOrigin.END + total_range = stack_params.stack_dz * (stack_params.images_to_save - 1) + expected_first_move = -total_range + first_call = autofocus_thing._stage.move_relative.call_args_list[0] + assert first_call.kwargs["z"] == expected_first_move, ( + "End origin offset not applied correctly" + ) + + # Check number of captures + assert ( + autofocus_thing.capture_stack_image.call_count == stack_params.images_to_save + ), "Incorrect number of captures for END origin" + + # Check final Z is equal to starting Z + final_z = autofocus_thing._stage.position["z"] + expected_final_z = start_z + assert final_z == expected_final_z, "Final Z position for END origin incorrect" + + +def test_invalid_stack_images_raises(): + """Test basic stack raises expected error for negative or zero image count.""" + for capture_count in [-3, 0]: + with pytest.raises(ValueError, match="Input should be greater than 0"): + StackParams( + stack_dz=10, + images_to_save=capture_count, + settling_time=0, + backlash_correction=0, + origin=StackOrigin.START, + ) + + +def test_invalid_stack_settling_raises(): + """Test basic stack raises expected error for negative settling time.""" + with pytest.raises(ValueError, match="Input should be greater than or equal to 0"): + StackParams( + stack_dz=10, + images_to_save=1, + settling_time=-1, + backlash_correction=0, + origin=StackOrigin.START, + ) + + +@pytest.mark.parametrize( + ("bad_path", "match_err"), + [ + ("", "String should have at least 1 character"), + (None, "Input should be a valid string"), + (67, "Input should be a valid string"), + ], +) +def test_invalid_capture_dir_raises(bad_path, match_err): + """Test basic stack raises expected error for bad image dir paths.""" + with pytest.raises(ValueError, match=match_err): + CaptureParams(images_dir=bad_path, save_resolution=(20, 20)) + + +@pytest.mark.parametrize( + ("bad_res", "match_err"), + [ + ((-100, 50), "Input should be greater than or equal to 1"), + ((20, 0), "Input should be greater than or equal to 1"), + ("", "Input should be a valid tuple"), + (None, "Input should be a valid tuple"), + (67, "Input should be a valid tuple"), + ( + ["path"], + "Input should be a valid integer, unable to parse string as an integer", + ), + ((20, 20, 20), "Tuple should have at most 2 items"), + ], +) +def test_invalid_capture_res_raises(bad_res, match_err): + """Test basic stack raises expected error for invalid save resolutions.""" + with pytest.raises(ValueError, match=match_err): + CaptureParams(images_dir="dummy", save_resolution=bad_res)