"""OpenFlexure Microscope autofocus module This module defines a Thing that is responsible for using the stage and camera together to perform an autofocus routine. See repository root for licensing information. """ from __future__ import annotations from contextlib import contextmanager import logging import time from typing import Annotated, Mapping, Optional, Sequence import os from fastapi import Depends import numpy as np from pydantic import BaseModel from labthings_fastapi.thing import Thing from labthings_fastapi.dependencies.blocking_portal import BlockingPortal from labthings_fastapi.decorators import thing_action, thing_property from labthings_fastapi.dependencies.metadata import GetThingStates from labthings_fastapi.types.numpy import NDArray from labthings_fastapi.dependencies.invocation import InvocationLogger from .camera import RawCameraDependency as Camera from .camera import CameraDependency as WrappedCamera from .stage import StageDependency as Stage from .capture import RawCaptureDependency as CaptureDep SETTLING_TIME = 0.3 BACKLASH_CORRECTION = 250 class JPEGSharpnessMonitor: __globals__ = globals() # Required for FastAPI dependency def __init__(self, stage: Stage, camera: Camera, portal: BlockingPortal): self.camera = camera self.stage = stage self.portal = portal print(f"Created sharpness monitor with {stage}, {camera}, {portal}") self.stage_positions: list[Mapping[str, int]] = [] self.stage_times: list[float] = [] self.jpeg_times: list[float] = [] self.jpeg_sizes: list[int] = [] running = False async def monitor_sharpness(self): """Start monitoring the frame sizes""" self.running = True async for frame in self.camera.lores_mjpeg_stream.frame_async_generator(): self.jpeg_times.append(time.time()) self.jpeg_sizes.append(len(frame)) if not self.running: break @contextmanager def run(self): """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 focus_rel(self, dz: int, **kwargs) -> tuple[int, int]: # Store the start time and position self.stage_times.append(time.time()) self.stage_positions.append(self.stage.position) # Main move self.stage.move_relative(z=dz, **kwargs) # Store the end time and position self.stage_times.append(time.time()) self.stage_positions.append(self.stage.position) # Index of the data for this movement data_index: int = len(self.stage_positions) - 2 # Final z position after move final_z_position: int = self.stage_positions[-1]["z"] return data_index, final_z_position def move_data( self, istart: int, istop: Optional[int] = None ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Extract sharpness as a function of (interpolated) z""" if istop is None: istop = istart + 2 jpeg_times: np.ndarray = np.array(self.jpeg_times) jpeg_sizes: np.ndarray = np.array(self.jpeg_sizes) stage_times: np.ndarray = np.array(self.stage_times)[istart:istop] stage_zs: np.ndarray = np.array( [p["z"] for p in self.stage_positions[istart:istop]] ) try: start: int = int(np.argmax(jpeg_times > stage_times[0])) stop: int = int(np.argmax(jpeg_times > stage_times[1])) except ValueError as e: if np.sum(jpeg_times > stage_times[0]) == 0: errmsg = ( "No images were captured during the move of the stage. " "Perhaps the camera is not streaming images?" ) raise ValueError(errmsg) from e raise e if stop < 1: stop = len(jpeg_times) logging.debug("changing stop to %s", (stop)) jpeg_times = jpeg_times[start:stop] jpeg_zs: np.ndarray = np.interp( jpeg_times, stage_times, stage_zs ) # np.ndarray[float] return jpeg_times, jpeg_zs, jpeg_sizes[start:stop] def sharpest_z_on_move(self, index: int) -> int: """Return the z position of the sharpest image on a given move""" _, jz, js = self.move_data(index) if len(js) == 0: raise ValueError( "No images were captured during the move of the stage. Perhaps the camera is not streaming images?" ) return jz[np.argmax(js)] def data_dict(self) -> SharpnessDataArrays: """Return the gathered data as a single convenient dictionary""" data = {} for k in ["jpeg_times", "jpeg_sizes", "stage_times", "stage_positions"]: data[k] = getattr(self, k) return SharpnessDataArrays(**data) SharpnessMonitorDep = Annotated[JPEGSharpnessMonitor, Depends()] class SharpnessDataArrays(BaseModel): jpeg_times: NDArray jpeg_sizes: NDArray stage_times: NDArray stage_positions: list[dict[str, int]] class AutofocusThing(Thing): """The Thing concerned with combinations of z axis movements and the camera. Actions here involve moving a stage in z, and using the camera to either capture images (generally, z-stacking) and measuring the sharpness of the field of view to assess focus (autofocus and testing the success of a z-stack)""" @thing_action def fast_autofocus( self, sharpness_monitor: SharpnessMonitorDep, dz: int = 2000, start: str = "centre", ) -> SharpnessDataArrays: """Sweep the stage up and down, then move to the sharpest point This method will will move down by dz/2, sweep up by dz, and then evaluate the position where the image was sharpest. We'll then move back down, and finally up to the sharpest point. """ with sharpness_monitor.run(): # Move to (-dz / 2) if start == "centre": sharpness_monitor.focus_rel(-dz / 2) # Move to dz while monitoring sharpness # i: Sharpness monitor index for this move # z: Final z position after move i, z = sharpness_monitor.focus_rel(dz, block_cancellation=True) # Get the z position with highest sharpness from the previous move (index i) fz: int = sharpness_monitor.sharpest_z_on_move(i) # Move all the way to the start so it's consistent i, z = sharpness_monitor.focus_rel(-dz) # Move to the target position fz (relative move of (fz - z)) sharpness_monitor.focus_rel(fz - z) # Return all focus data return sharpness_monitor.data_dict() @thing_action def z_move_and_measure_sharpness( self, sharpness_monitor: SharpnessMonitorDep, dz: Sequence[int], wait: float = 0, ) -> SharpnessDataArrays: """Make a move (or a series of moves) and monitor sharpness This method will will make a series of relative moves in z, and return the sharpness (JPEG size) vs time, along with timestamps for the moves. This can be used to calibrate autofocus. Each move is relative to the last one, i.e. we will finish at `sum(dz)` relative to the starting position. If `wait` is specified, we will wait for that many seconds between moves. """ with sharpness_monitor.run(): for i, current_dz in enumerate(dz): if i > 0 and wait > 0: time.sleep(wait) sharpness_monitor.focus_rel(current_dz) return sharpness_monitor.data_dict() @thing_action def looping_autofocus( self, stage: Stage, sharpness_monitor: SharpnessMonitorDep, dz=2000, start="centre", ): """Repeatedly autofocus the stage until it looks focused. This action will run the `fast_autofocus` action until it settles on a point in the middle 3/5 of its range. Such logic can be helpful if the microscope is close to focus, but not quite within `dz/2`. It will attempt to autofocus up to 10 times. """ repeat = True attempts = 0 backlash = 200 with sharpness_monitor.run(): while repeat and attempts < 10: if start == "centre": stage.move_relative(x=0, y=0, z=-(backlash + dz / 2)) stage.move_relative(x=0, y=0, z=backlash) i, z = sharpness_monitor.focus_rel(dz, block_cancellation=True) _, heights, sizes = sharpness_monitor.move_data(i) peak_height = heights[np.argmax(sizes)] height_min = np.min(heights) height_max = np.max(heights) if ( peak_height - height_min < dz / 5 or height_max - peak_height < dz / 5 ): attempts += 1 start = "centre" stage.move_absolute(z=peak_height - backlash) stage.move_absolute(z=peak_height) else: repeat = False stage.move_relative(x=0, y=0, z=-(dz + backlash)) stage.move_absolute(z=peak_height) return heights.tolist(), sizes.tolist() @thing_property def stack_images_to_capture(self) -> int: """The number of images to capture and save in a stack Defaults to 1 unless you need to see either side of focus""" return self.thing_settings.get("stack_images_to_capture", 1) @stack_images_to_capture.setter def stack_images_to_capture(self, value: int) -> None: self.thing_settings["stack_images_to_capture"] = value @thing_property def stack_images_to_test(self) -> int: """The number of images to test for successful focusing in a stack Defaults to 9, which balances reliability and speed""" return self.thing_settings.get("stack_images_to_test", 9) @stack_images_to_test.setter def stack_images_to_test(self, value: int) -> None: self.thing_settings["stack_images_to_test"] = value @thing_property def stack_dz(self) -> int: """Space in steps between images in a z-stack Suggested is 50 for 60-100x 100 for 40x 200 for 20x""" return self.thing_settings.get("stack_dz", 50) @stack_dz.setter def stack_dz(self, value: int) -> None: self.thing_settings["stack_dz"] = value @thing_action def run_smart_stack( self, cam: WrappedCamera, stage: Stage, logger: InvocationLogger, metadata_getter: GetThingStates, capture: CaptureDep, sharpness_monitor: SharpnessMonitorDep, images_dir: str, autofocus_dz: int, ) -> None: """Run a smart stack, which captures images offset in z, testing whether the sharpest image is towards the centre of the stack. The sharpest image, and optionally images around the sharpest, will be saved using their coordinates to images_dir Arguments: images_dir: the folder to save all images autofocus_dz: should the stack fail, the range to refocus over before retrying variables cam to sharpness_monitor are Thing dependencies injected automatically by LabThings FastAPI """ # Set the variables to prevent changes from the GUI or other windows stack_dz = self.stack_dz images_to_capture = self.stack_images_to_capture images_to_test = self.stack_images_to_test stack_z_range = stack_dz * (images_to_test - 1) # Starting too low by "overshoot" makes smart stacking faster. # Starting a stack too high requires it to move to the start, # autofocus and then re-stack. Starting slightly too low only # requires extra +z movements and captures. overshoot = stack_dz * 5 # Ensure the stack settings are appropriate self.validate_stack_inputs(images_to_test, images_to_capture) success = False # Loop until a stack is successful while not success: result, heights, captures, sharpest_index = self.z_stack( stack_dz, images_to_test, images_dir, stack_z_range, overshoot, stage, cam, capture, metadata_getter, ) if result == "success": success = True break # If a stack is not successful, move to the start and autofocus self.reset_stack( heights, autofocus_dz, stage, sharpness_monitor, ) # Save the sharpest image, and images either side of focus self.save_stack( sharpest_index, captures, images_to_test, images_to_capture, logger, capture, ) # Return the z position of the sharpest image, for path planning and tracking return heights[-images_to_test:][sharpest_index] def reset_stack( self, heights: list[int], autofocus_dz: int, stage: Stage, sharpness_monitor: SharpnessMonitorDep, ) -> None: """Return to the initial height of the current stack, and run a looping autofocus. Clears all previous captures, heights and sharpnesses. Arguments: heights: a list of the z positions of previous captures autofocus_dz: the range in steps to autofocus variables stage and sharpness_monitor are Thing dependencies passed through from the calling action """ stage.move_absolute(z=heights[0]) self.looping_autofocus( stage=stage, sharpness_monitor=sharpness_monitor, dz=autofocus_dz, ) def save_stack( self, sharpest_index: int, captures: list[list], images_to_test: int, images_to_capture: int, logger: InvocationLogger, capture: CaptureDep, ) -> int: """Save the required captures to disk. Will save the sharpest image, and any images either side of focus. Arguments: sharpest_index: the index of the sharpest image, within the "images_to_test" slice captures: a list of captures, including file name, image data and metadata images_to_test: the number of images in the stack tested for success images_to_capture: the number of images to save to disk variables logger and capture are Thing dependencies passed through from the calling action """ # Find the range of images from the stack to capture stack_extent = int((images_to_capture - 1) / 2) stack_range = range( sharpest_index - stack_extent, sharpest_index + stack_extent + 1 ) # Loop through the range, saving each capture to disk for capture_index in stack_range: capture._save_capture( jpeg_path=captures[-images_to_test:][capture_index][0], image=captures[-images_to_test:][capture_index][1], metadata=captures[-images_to_test:][capture_index][2], logger=logger, ) return sharpest_index def z_stack( self, stack_dz: int, images_to_test: int, images_dir: str, stack_z_range: int, overshoot: int, stage: Stage, cam: WrappedCamera, capture: CaptureDep, metadata_getter: GetThingStates, ) -> list: """Capture a series of images offset by stack_dz, and test whether the sharpest image is towards the centre of the stack. Returns a test result string, a list of z positions, a list of captures and the index of the sharpest image in a successful stack. Arguments: stack_dz: the distance in steps between images in the stack images_to_test: the number of images in the stack to test for focus images_dir: a string of the path to write all images stack_z_range: the height in steps of the stack to test overshoot: how far below the estimated optimal starting position to begin stacking variables stage to metadata_getter are Thing dependencies passed through from the calling action """ # 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(z=-(overshoot + BACKLASH_CORRECTION + stack_z_range / 2)) stage.move_relative(z=BACKLASH_CORRECTION) captures = [] sharpnesses = [] heights = [] # If the sharpest image isn't found within the 15 images above the estimated point, break # the loop and return "restart" while len(captures) <= images_to_test + 15: time.sleep(SETTLING_TIME) # Append a new image to the stack captures, heights, sharpnesses = self.capture_stack_image( captures, heights, sharpnesses, images_dir, cam, stage, capture, metadata_getter, ) # If the number of images is enough to test, test them if len(captures) >= images_to_test: stack_result = self.check_stack_result(sharpnesses[-images_to_test:]) if stack_result == "success": sharpest_index = np.argmax(sharpnesses[-images_to_test:]) return "success", heights, captures, sharpest_index if stack_result == "restart": return "restart", heights, None, None stage.move_relative(z=stack_dz) return "restart", heights, None, None def capture_stack_image( self, captures: list[list], heights: list[int], sharpnesses: list[int], images_dir: str, cam: WrappedCamera, stage: Stage, capture: CaptureDep, metadata_getter: GetThingStates, ) -> list: """Append a new capture to the ongoing stack. Includes appending the height, image sharpness and image data to the relevant lists. arguments: captures: a list of captures, including file name, image data and metadata heights: a list of the z positions of previous captures sharpnesses: a list of the sharpnesses of previous captures images_dir: a path to the folder to save images variables stage to metadata_getter are Thing dependencies passed through from the calling action """ stage_location = stage.position jpeg_path = os.path.join( images_dir, f"{stage_location['x']}_{stage_location['y']}_{stage_location['z']}.jpeg", ) image, metadata = capture._capture_image( cam=cam, metadata_getter=metadata_getter, ) captures.append([jpeg_path, image, metadata]) sharpnesses.append(cam.grab_jpeg_size(stream_name="lores")) heights.append(stage_location["z"]) return captures, heights, sharpnesses def validate_stack_inputs( self, images_to_test: int, images_to_capture: int ) -> None: """Check the stack settings are appropriate, and raise an error if not Arguments: images_to_test: number of images in the stack to test for focus images_to_capture: number of images to be captured around the focused image """ if images_to_test < images_to_capture: raise RuntimeError( "Can't capture more images than are tested. Please increase number to test, or decrease number to capture" ) if images_to_test % 2 == 0: raise RuntimeError("Images to test should be odd") if images_to_test <= 0 or images_to_capture <= 0: raise RuntimeError("Stack parameters need to be at least 1") def check_stack_result(self, sharpnesses: list[int]) -> str: """Test a list of sharpnesses, to decide whether the sharpest image from a stack is within them Returns a string 'success' if the sharpest image is towards the centre 'continue' if the sharpest image is in the final two images of the list 'restart' if the sharpest image is in the first two images of the list Arguments: sharpnesses: a list of the sharpnesses to test for focus """ sharpest_index = np.argmax(sharpnesses) sharpness_length = len(sharpnesses) # If only testing one image, then by definition the sharpest is central if sharpness_length == 1: return "success" # If testing three images, test if the centre is the sharpest if sharpness_length == 3: if sharpest_index == 1: return "success" if sharpest_index == 0: return "restart" return "continue" # For larger stacks, test if the best image is not within two of the edge of the stack # ie for a stack of 7 images, best image must be between 3rd and and 5th exclusion_range = 2 if sharpest_index < exclusion_range: return "restart" if sharpest_index >= sharpness_length - exclusion_range: return "continue" return "success"