From 189e82e8dcbccdddc35f7f91937571519f581669 Mon Sep 17 00:00:00 2001 From: Julian Stirling Date: Tue, 16 Dec 2025 16:14:16 +0000 Subject: [PATCH] Remove deps from autofocus --- .../things/autofocus.py | 133 ++++++------------ tests/test_autofocus.py | 37 +++-- tests/test_stack.py | 36 ++--- 3 files changed, 75 insertions(+), 131 deletions(-) diff --git a/src/openflexure_microscope_server/things/autofocus.py b/src/openflexure_microscope_server/things/autofocus.py index ee023ddb..1d310db0 100644 --- a/src/openflexure_microscope_server/things/autofocus.py +++ b/src/openflexure_microscope_server/things/autofocus.py @@ -10,20 +10,18 @@ See repository root for licensing information. import logging import os import time -from contextlib import contextmanager from dataclasses import dataclass -from typing import Annotated, Literal, Mapping, Optional, Sequence +from types import TracebackType +from typing import Literal, Mapping, Optional, Self, Sequence import numpy as np -from fastapi import Depends from pydantic import BaseModel, computed_field, field_validator, model_validator import labthings_fastapi as lt from labthings_fastapi.types.numpy import NDArray -from .camera import CameraDependency as CameraClient -from .camera import RawCameraDependency as RawCamera -from .stage import StageDependency as Stage +from .camera import BaseCamera +from .stage import BaseStage LOGGER = logging.getLogger(__name__) MIN_TEST_IMAGE_COUNT = 3 @@ -224,20 +222,16 @@ class JPEGSharpnessMonitor: SharpnessMonitorDep as an argument is called. """ - def __init__( - self, stage: Stage, camera: RawCamera, portal: lt.deps.BlockingPortal - ) -> None: + def __init__(self, stage: BaseStage, camera: BaseCamera) -> None: """Initialise a new JPEGSharpnessMonitor. The args are injected automatically. :param stage: A direct_thing_client dependency for the the microscope stage. :param camera: A raw_thing_client depeendency for the camera. This is a raw dependency as the underlying class needs to be - :param portal: The asyncio blocking portal for asynchronous task scheduling. """ self.camera = camera self.stage = stage - self.portal = portal - LOGGER.debug(f"Created sharpness monitor with {stage}, {camera}, {portal}") + LOGGER.debug(f"Created sharpness monitor with {stage}, {camera}") self.stage_positions: list[Mapping[str, int]] = [] self.stage_times: list[float] = [] self.jpeg_times: list[float] = [] @@ -254,14 +248,23 @@ class JPEGSharpnessMonitor: if not self.running: break - @contextmanager - def run(self) -> None: - """Context manager, during which we will monitor sharpness from the camera.""" - self.portal.start_task_soon(self.monitor_sharpness) - try: - yield - finally: - self.running = False + def __enter__(self) -> Self: + """Start context manager, during which sharpness from the camera is monitored.""" + # Use the cameras _thing_server_interface to get access to the server event loop + # and start a task. + self.camera._thing_server_interface.start_async_task_soon( + self.monitor_sharpness + ) + return self + + def __exit__( + self, + _exc_type: type[BaseException], + _exc_value: Optional[BaseException], + _traceback: Optional[TracebackType], + ) -> None: + """Clean up after context manager is closed.""" + self.running = False def focus_rel(self, dz: int, block_cancellation: int = False) -> tuple[int, int]: """Move the stage by dz, monitoring the position over time. @@ -339,9 +342,6 @@ class JPEGSharpnessMonitor: return SharpnessDataArrays(**data) -SharpnessMonitorDep = Annotated[JPEGSharpnessMonitor, Depends()] - - class AutofocusThing(lt.Thing): """The Thing concerned with combinations of z axis movements and the camera. @@ -350,10 +350,12 @@ class AutofocusThing(lt.Thing): field of view to assess focus (autofocus and testing the success of a z-stack) """ + _cam: BaseCamera = lt.thing_slot() + _stage: BaseStage = lt.thing_slot() + @lt.action def fast_autofocus( self, - sharpness_monitor: SharpnessMonitorDep, dz: int = 2000, start: Literal["centre", "base"] = "centre", ) -> SharpnessDataArrays: @@ -363,7 +365,7 @@ class AutofocusThing(lt.Thing): the position where the image was sharpest. We'll then move back down, and finally up to the sharpest point. """ - with sharpness_monitor.run(): + with JPEGSharpnessMonitor(self._stage, self._cam) as sharpness_monitor: # Move to (-dz / 2) if start == "centre": sharpness_monitor.focus_rel(-dz // 2) @@ -384,7 +386,6 @@ class AutofocusThing(lt.Thing): @lt.action def z_move_and_measure_sharpness( self, - sharpness_monitor: SharpnessMonitorDep, dz: Sequence[int], wait: float = 0, ) -> SharpnessDataArrays: @@ -400,7 +401,7 @@ class AutofocusThing(lt.Thing): If ``wait`` is specified, we will wait for that many seconds between moves. """ - with sharpness_monitor.run(): + with JPEGSharpnessMonitor(self._stage, self._cam) as sharpness_monitor: for move_index, current_dz in enumerate(dz): if move_index > 0 and wait > 0: time.sleep(wait) @@ -410,8 +411,6 @@ class AutofocusThing(lt.Thing): @lt.action def looping_autofocus( self, - stage: Stage, - sharpness_monitor: SharpnessMonitorDep, dz: int = 2000, start: Literal["centre", "base"] = "centre", ) -> tuple[list[float], list[float]]: @@ -425,12 +424,12 @@ class AutofocusThing(lt.Thing): attempt = 0 backlash = 200 - with sharpness_monitor.run(): + with JPEGSharpnessMonitor(self._stage, self._cam) as sharpness_monitor: while attempt < 10: attempt += 1 if start == "centre": - stage.move_relative(x=0, y=0, z=-(backlash + dz / 2)) - stage.move_relative(x=0, y=0, z=backlash) + self._stage.move_relative(x=0, y=0, z=-(backlash + dz / 2)) + self._stage.move_relative(x=0, y=0, z=backlash) # Always start centrally for future runs start = "centre" @@ -445,8 +444,8 @@ class AutofocusThing(lt.Thing): target_max = np.max(heights) - dz / 5 # move to the peak - stage.move_absolute(z=peak_height - backlash) - stage.move_absolute(z=peak_height) + self._stage.move_absolute(z=peak_height - backlash) + self._stage.move_absolute(z=peak_height) if target_min < peak_height < target_max: # If it is within the target range then return @@ -557,9 +556,6 @@ class AutofocusThing(lt.Thing): @lt.action def run_smart_stack( self, - cam: CameraClient, - stage: Stage, - sharpness_monitor: SharpnessMonitorDep, stack_parameters: StackParams, save_on_failure: bool = False, check_turning_points: bool = True, @@ -572,11 +568,6 @@ 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 cam: Camera Dependency supplied by LabThings dependency injection - :param stage: Stage Dependency supplied by LabThings dependency injection - :param sharpness_monitor: Sharpness Monitor Dependency (for focus detection) - supplied by LabThings dependency injection :param stack_parameters: A StackParams object containing the required parameters to run a stack. :param save_on_failure: Whether to save an image even if no focus was found. @@ -595,8 +586,6 @@ class AutofocusThing(lt.Thing): success, captures, sharpest_id = self.z_stack( stack_parameters=stack_parameters, check_turning_points=check_turning_points, - cam=cam, - stage=stage, ) if success: @@ -609,12 +598,7 @@ class AutofocusThing(lt.Thing): 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, - stage, - sharpness_monitor, - ) + self.reset_stack(initial_z_pos, stack_parameters.autofocus_dz) except NoFocusFoundError: break @@ -626,7 +610,6 @@ class AutofocusThing(lt.Thing): sharpest_id=sharpest_id, captures=captures, stack_parameters=stack_parameters, - cam=cam, ) return success, _get_capture_by_id(captures, sharpest_id).position["z"] @@ -635,8 +618,6 @@ class AutofocusThing(lt.Thing): self, initial_z_pos: list[int], autofocus_dz: int, - stage: Stage, - sharpness_monitor: SharpnessMonitorDep, ) -> None: """Return to the initial z position and run a looping autofocus. @@ -646,19 +627,14 @@ class AutofocusThing(lt.Thing): ``stage`` and ``sharpness_monitor`` are Thing dependencies passed through from the calling action. """ - stage.move_absolute(z=initial_z_pos) - self.looping_autofocus( - stage=stage, - sharpness_monitor=sharpness_monitor, - dz=autofocus_dz, - ) + self._stage.move_absolute(z=initial_z_pos) + self.looping_autofocus(dz=autofocus_dz) def save_stack( self, sharpest_id: int, captures: list[list], stack_parameters: StackParams, - cam: CameraClient, ) -> int: """Save the required captures to disk. @@ -669,28 +645,24 @@ class AutofocusThing(lt.Thing): :param captures: a list of captures, including file name, image data and metadata :param stack_parameters: a StackParams object holding stack parameters - :param cam: is a Thing dependency passed through from the calling action - """ sharpest_index = _get_capture_index_by_id(captures, sharpest_id) slice_to_save = stack_parameters.slice_to_save(sharpest_index) # Loop through the range, saving each capture to disk for capture in captures[slice_to_save]: - cam.save_from_memory( + self._cam.save_from_memory( jpeg_path=os.path.join(stack_parameters.images_dir, capture.filename), save_resolution=stack_parameters.save_resolution, buffer_id=capture.buffer_id, ) - cam.clear_buffers() + self._cam.clear_buffers() return sharpest_index def z_stack( self, stack_parameters: StackParams, check_turning_points: bool, - cam: CameraClient, - stage: Stage, ) -> tuple[bool, list[CaptureInfo], int]: """Capture a series of images checking that sharpest image central. @@ -703,8 +675,6 @@ class AutofocusThing(lt.Thing): :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) - :param cam: Camera Dependency to be passed through from the calling action - :param stage: Stage Dependency to be passed through from the calling action :returns: A tuple of @@ -714,14 +684,14 @@ class AutofocusThing(lt.Thing): """ # Move down by the height of the z stack, plus an overshoot # Better to start too low and take too many images than too high and need to refocus - stage.move_relative( + self._stage.move_relative( z=-( stack_parameters.steps_undershoot + stack_parameters.backlash_correction + stack_parameters.stack_z_range / 2 ) ) - stage.move_relative(z=stack_parameters.backlash_correction) + self._stage.move_relative(z=stack_parameters.backlash_correction) captures = [] # Always check for focus using the the last `min_images_to_test` in the @@ -735,11 +705,7 @@ class AutofocusThing(lt.Thing): # Append a new image to the stack captures.append( - self.capture_stack_image( - cam, - stage, - buffer_max=stack_parameters.min_images_to_test, - ) + self.capture_stack_image(buffer_max=stack_parameters.min_images_to_test) ) # If the number of images is enough to test, test them @@ -754,33 +720,26 @@ class AutofocusThing(lt.Thing): if result == "restart": return False, captures, capture_id # If reached here the result was "continue" - stage.move_relative(z=stack_parameters.stack_dz) + self._stage.move_relative(z=stack_parameters.stack_dz) return False, captures, capture_id - def capture_stack_image( - self, - cam: CameraClient, - stage: Stage, - buffer_max: int, - ) -> CaptureInfo: + def capture_stack_image(self, buffer_max: int) -> CaptureInfo: """Capture another image and return the capture information. The capture is stored by the camera Thing, and can be saved by ID. - :param cam: Camera Dependency to be passed through from the calling action - :param stage: Stage Dependency to be passed through from the calling action :param buffer_max: The maximum number of images to tell the camera to keep in memory for saving once the stack is complete :returns: A CaptureInfo object containing the capture information including its camera buffer_id needed for saving. """ - stage_location = stage.position - buffer_id = cam.capture_to_memory(buffer_max=buffer_max) + stage_location = self._stage.position + buffer_id = self._cam.capture_to_memory(buffer_max=buffer_max) return CaptureInfo( buffer_id=buffer_id, position=stage_location, - sharpness=cam.grab_jpeg_size(stream_name="lores"), + sharpness=self._cam.grab_jpeg_size(stream_name="lores"), ) # Silence too many returns in this situation as refactoring to reduce returns is diff --git a/tests/test_autofocus.py b/tests/test_autofocus.py index 58a5f577..1c3f6e08 100644 --- a/tests/test_autofocus.py +++ b/tests/test_autofocus.py @@ -48,22 +48,23 @@ def fake_sharpness_data( def test_looping_autofocus(start_z, max_loc, centre, attempts_expected, passes, mocker): """Test the high level looping autofocus algorithm.""" dz = 2000 + autofocus_thing = create_thing_without_server(AutofocusThing, mock_all_slots=True) + # Make a mock stage where move_absolute abs and relative updates the position counter. - stage = mocker.Mock() - stage.position = {"x": 0, "y": 0, "z": start_z} + autofocus_thing._stage.position = {"x": 0, "y": 0, "z": start_z} def set_pos(**kwargs: int) -> None: """Move absolute should update position. So make a side effect for the mock.""" for axis, value in kwargs.items(): - stage.position[axis] = value + autofocus_thing._stage.position[axis] = value def adjust_pos(**kwargs: int) -> None: """Move relative should update position. So make a side effect for the mock.""" for axis, value in kwargs.items(): - stage.position[axis] += value + autofocus_thing._stage.position[axis] += value - stage.move_absolute.side_effect = set_pos - stage.move_relative.side_effect = adjust_pos + autofocus_thing._stage.move_absolute.side_effect = set_pos + autofocus_thing._stage.move_relative.side_effect = adjust_pos # Make a mock sharpness monitor that can generate sharpness data. sharpness_monitor = mocker.MagicMock() @@ -73,27 +74,25 @@ def test_looping_autofocus(start_z, max_loc, centre, attempts_expected, passes, """Generate sharpnesses based on parameterised input, and mock stage position.""" return fake_sharpness_data( dz=dz, - start_z=stage.position["z"], + start_z=autofocus_thing._stage.position["z"], max_loc=max_loc, ) sharpness_monitor.move_data.side_effect = return_sharpness - autofocus_thing = create_thing_without_server(AutofocusThing) + # Mock the context manager call so out MagicMock sharpness monitor is returned. + mock_context = mocker.patch( + "openflexure_microscope_server.things.autofocus.JPEGSharpnessMonitor" + ) + mock_context.return_value.__enter__.return_value = sharpness_monitor + mock_context.return_value.__exit__.return_value = None + if passes: - autofocus_thing.looping_autofocus( - stage=stage, - sharpness_monitor=sharpness_monitor, - dz=dz, - start="centre" if centre else "base", - ) + autofocus_thing.looping_autofocus(dz=dz, start="centre" if centre else "base") else: with pytest.raises(NoFocusFoundError): autofocus_thing.looping_autofocus( - stage=stage, - sharpness_monitor=sharpness_monitor, - dz=dz, - start="centre" if centre else "base", + dz=dz, start="centre" if centre else "base" ) assert sharpness_monitor.focus_rel.call_count == attempts_expected - assert abs(max_loc - stage.position["z"]) < dz / 40 + assert abs(max_loc - autofocus_thing._stage.position["z"]) < dz / 40 diff --git a/tests/test_stack.py b/tests/test_stack.py index e7f1df27..17ccec94 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -265,7 +265,7 @@ def test_retrieval_of_captures(start): @pytest.fixture def autofocus_thing(): """Return an autofocus thing connected to a server.""" - return create_thing_without_server(AutofocusThing) + return create_thing_without_server(AutofocusThing, mock_all_slots=True) def test_create_stack(autofocus_thing, caplog): @@ -352,9 +352,6 @@ def test_coercing_stack_save_ims( @pytest.mark.parametrize("pass_on", [1, 2, 3, 4]) def test_run_smart_stack(pass_on, autofocus_thing, mocker): """Test Running smart stack with the stack passing on different attempts.""" - cam = mocker.Mock() - stage = mocker.Mock() - sharpness_monitor = mocker.MagicMock() stack_params = autofocus_thing.create_stack_params( autofocus_dz=2000, images_dir="/this/is/fake", @@ -386,9 +383,6 @@ def test_run_smart_stack(pass_on, autofocus_thing, mocker): # Run it success, final_z = autofocus_thing.run_smart_stack( - cam=cam, - stage=stage, - sharpness_monitor=sharpness_monitor, stack_parameters=stack_params, save_on_failure=False, check_turning_points=True, @@ -403,16 +397,16 @@ def test_run_smart_stack(pass_on, autofocus_thing, mocker): n_stacks = min(pass_on, stack_params.max_attempts) assert autofocus_thing.z_stack.call_count == n_stacks # Move absolute should be 1 less time that the number of times z_stack_run - assert stage.move_absolute.call_count == n_stacks - 1 + assert autofocus_thing._stage.move_absolute.call_count == n_stacks - 1 # As should looping autofocus assert autofocus_thing.looping_autofocus.call_count == n_stacks - 1 # Check rest stack is moving to the first image in the stack. if n_stacks > 1: - assert stage.move_absolute.call_args.kwargs["z"] == -99 + assert autofocus_thing._stage.move_absolute.call_args.kwargs["z"] == -99 # Mock called to save image - assert cam.save_from_memory.call_count == (1 if success else 0) + assert autofocus_thing._cam.save_from_memory.call_count == (1 if success else 0) def setup_and_run_z_stack(check_returns, check_turning_points, autofocus_thing, mocker): @@ -430,8 +424,6 @@ def setup_and_run_z_stack(check_returns, check_turning_points, autofocus_thing, ) stack_params.settling_time = 0 # Don't settle or tests take forever. - stage = mocker.Mock() - cam = mocker.Mock() autofocus_thing.capture_stack_image = mocker.Mock() if isinstance(check_returns, list): autofocus_thing.check_stack_result = mocker.Mock(side_effect=check_returns) @@ -440,8 +432,6 @@ def setup_and_run_z_stack(check_returns, check_turning_points, autofocus_thing, return autofocus_thing.z_stack( stack_parameters=stack_params, check_turning_points=check_turning_points, - cam=cam, - stage=stage, ) @@ -506,20 +496,16 @@ def test_z_stack_return(autofocus_thing, mocker): assert ret[0] -def test_capture_stack_image(autofocus_thing, mocker): +def test_capture_stack_image(autofocus_thing): """Check that capture stack image calls the expected functions and returns the expected data.""" - stage = mocker.Mock() - stage.position = {"x": 123, "y": 456, "z": 789} - cam = mocker.Mock() - cam.capture_to_memory.return_value = "fake_buffer_id" - cam.grab_jpeg_size.return_value = 54321 + autofocus_thing._stage.position = {"x": 123, "y": 456, "z": 789} + autofocus_thing._cam.capture_to_memory.return_value = "fake_buffer_id" + autofocus_thing._cam.grab_jpeg_size.return_value = 54321 buffer_max = 11 - info = autofocus_thing.capture_stack_image( - cam=cam, stage=stage, buffer_max=buffer_max - ) - assert cam.capture_to_memory.call_count == 1 - assert cam.grab_jpeg_size.call_count == 1 + info = autofocus_thing.capture_stack_image(buffer_max=buffer_max) + assert autofocus_thing._cam.capture_to_memory.call_count == 1 + assert autofocus_thing._cam.grab_jpeg_size.call_count == 1 assert info.buffer_id == "fake_buffer_id" assert info.position == {"x": 123, "y": 456, "z": 789} assert info.sharpness == 54321