diff --git a/openflexure_microscope/api/default_extensions/scan.py b/openflexure_microscope/api/default_extensions/scan.py index 4253bb6d..363838dc 100644 --- a/openflexure_microscope/api/default_extensions/scan.py +++ b/openflexure_microscope/api/default_extensions/scan.py @@ -2,7 +2,7 @@ import datetime import logging import time import uuid -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple, Callable import marshmallow import numpy as np @@ -29,6 +29,137 @@ XyzCoordinate = Tuple[int, int, int] ### Grid construction +class FocusManager: + """Manage axial motion during a series of XY moves + + This class keeps track of the focus position as we move around. + It currently uses the regular autofocus method, and has support + for background detection. + """ + initial_position: XyzCoordinate = None + focused_positions: List[XyzCoordinate] = None + current_image_is_background: Optional[Callable] = None + autofocus: Optional[Callable] = None + microscope: Microscope = None + axial_jump_threshold: Optional[float] = None + + def __init__( + self, + microscope: Microscope, + initial_position: XyzCoordinate, + autofocus_function: Optional[Callable], + current_image_is_background_function: Optional[Callable] = None, + axial_jump_threshold: Optional[float] = 0.4 + ): + """Set up management of axial motion. + + The `FocusManager` keeps track of previous positions where the + microscope was in focus, and will estimate the best Z value for + future XY positions. + + Arguments: + * microscope: The microscope object. + * initial_position: the XYZ position of the start of the scan + * autofocus: A function that performs an autofocus. + * current_image_is_background: a function that returns `True` + if the microscope is currently looking at an empty field. + * axial_jump_threshold: the maximum ratio between axial and + lateral moves. If this is not None, it will not count the + autofocus routine as successful if the focus moves more than + this ratio times the lateral move between two points. This can + help avoid focus drift due to accidentally focusing on the coverslip. + + If either of the `autofocus` or `current_image_is_background` + functions are None, we will neither perform an autofocus, nor + use background estimation. + """ + self.initial_position = initial_position + self.microscope = microscope + self.autofocus_function = autofocus_function + if not autofocus_function: + logging.info("Setting up FocusManager with autofocus disabled.") + self.current_image_is_background_function = current_image_is_background_function + self.axial_jump_threshold = axial_jump_threshold + self.focused_positions = [] + + def record_focused_point(self, position: XyzCoordinate): + """Add a position to the list of successfully-focused points""" + self.focused_positions.append(position) + + def closest_focused_point(self, position: XyCoordinate) -> XyzCoordinate: + """The closest point in our list of focused points to a given XY position.""" + return closest_point_in_xy(position, self.focused_positions) + + def estimate_z(self, position: XyCoordinate) -> int: + """Estimate the z position most likely to be in focus at an XY point + + The next z position is estimated based on the closest point that was in focus. + For a snake/spiral scan, this should always be the last point, unless + it's skipped for some reason. In a raster scan, this should be the last + point, except when we're at the start of a row when it will be the first + point of the preceding row. + + It is possible that for some scan geometries, we won't be using the most + recent point (e.g. if X and Y spacing in a raster scan are very different). + This does not happen with the default settings used for raster scanning + clinical samples. + """ + closest_focused_point = self.closest_focused_point(position) + if closest_focused_point: + return closest_focused_point[2] + else: + return self.initial_position[2] + + def check_for_axial_jumps(self, position: XyzCoordinate) -> bool: + """Check if a position is inconsistent with previous positions. + + This function returns `True` if the specified position is not consistent + with the list of previously-visited positions, i.e. it's made a big axial + move but a small lateral move. + + The threshold is set by `self.axial_jump_threshold`, and if that is `None` + no check is performed. Sensible values are probably between 0.1 and 0.5. + """ + if not self.axial_jump_threshold: + return False + closest_focused_point = self.closest_focused_point(position) + move = np.array(position) - np.array(closest_focused_point) + lateral_move = np.sqrt(np.sum(move[:2]**2)) + axial_move = np.abs(move[2]) + return axial_move > lateral_move * self.axial_jump_threshold + + def autofocus(self): + """Perform an autofocus routine. + + If autofocus is disabled, nothing happens here. If it is enabled, we optionally + check whether there's anything in the image to focus on, and then run the autofocus + routine. + """ + if not self.autofocus_function: + logging.debug("Autofocus is disabled, skipping autofocus.") + return + if self.current_image_is_background_function: + # If it's been set, call the background detect function and skip + # autofocus if appropriate + if self.current_image_is_background_function(): + here = self.microscope.stage.position + logging.info(f"Detected an empty field at {here}, skipping autofocus.") + return + # Assuming it's not disabled, and we're not skipping it, actually run the autofocus now + self.autofocus_function() + # Now, check for big jumps and record the new position if we've not made a big jump + here = self.microscope.stage.position + if self.check_for_axial_jumps(here): + logging.warning( + f"During a scan, there was a large axial jump from " + f"{self.closest_focused_point(here[:2])}" + f" to {here}. This may mean autofocus has failed." + ) + else: + # If there has not been a jump in focus, record the point as successful + self.record_focused_point(here) + + def construct_grid( initial: XyCoordinate, @@ -94,6 +225,23 @@ def construct_grid( return arr +def construct_grid_1d( + initial: XyCoordinate, + step_sizes: XyCoordinate, + n_steps: XyCoordinate, + style: Literal["raster", "snake", "spiral"] = "raster", +) -> List[XyCoordinate]: + """Construct coordinates for a scan, returning a 1D list. + + This is the same set of coordinates returned by `construct_grid` + but the list-of-lists is flattened to a simple list. + """ + path = [] + grid = construct_grid(initial=initial, step_sizes=step_sizes, n_steps=n_steps, style=style) + for line in grid: + path += line # concatenate the lists + return path + def closest_point_in_xy( current_position: XyCoordinate, points: List[XyzCoordinate] @@ -181,6 +329,51 @@ class ScanExtension(BaseExtension): logging.info(progress) return progress + def get_autofocus_function(self, dz: int, use_fast_autofocus: bool) -> Callable: + """Return a function that will perform an autofocus routine. + + This should be called at the start of a scan. It will check that + the necessary hardware and software are present, and raise a helpful + error if they are not. + """ + microscope = find_component("org.openflexure.microscope") + # Locate the autofocus extension + autofocus_extension = find_extension("org.openflexure.autofocus") + if not autofocus_extension: + raise RuntimeError("The Autofocus extension is missing: select 'None' as your autofocus type to scan without autofocusing.") + if not (microscope.has_real_stage() and microscope.has_real_camera()): + raise RuntimeError("A real stage and camera are needed in order to autofocus. You can still run a scan without autofocus.") + + if use_fast_autofocus: + def autofocus(): + # Run fast autofocus. Client should provide dz ~ 2000 + autofocus_extension.fast_autofocus(microscope, dz=dz) + time.sleep(0.5) + return autofocus + else: + def autofocus(): + # Run slow autofocus. Client should provide dz ~ 50 + autofocus_extension.autofocus( + microscope, + range(-3 * dz, 4 * dz, dz), + ) + time.sleep(0.5) + return autofocus + + def get_background_detect_function(self): + """Return a function that returns true if we are looking at background""" + # Check for the background detect extension, raise an error now if it's missing. + background_detect_extension = find_extension("org.openflexure.background-detect") + if not background_detect_extension: + raise RuntimeError( + "Detecting background fields requires the background detect extension and it was not found." + ) + def current_image_is_background(): + verdict = background_detect_extension.grab_and_classify_image() + logging.debug(f"Background detection verdict: {verdict}") + return verdict["classification"] == "background" + return current_image_is_background + ### Scanning def tile( self, @@ -209,15 +402,13 @@ class ScanExtension(BaseExtension): # Store initial position initial_position = microscope.stage.position - # Construct an x-y grid (worry about z later) - x_y_grid = construct_grid( + # Construct an x-y scan path (list of 2D coordinates) + path = construct_grid_1d( initial_position[:2], stride_size[:2], grid[:2], style=style - ) # This is a list of lists. + ) # Keep task progress - # NB the number of points is found from the number of elements in - # x_y_grid; this is not guaranteed to be the same in every row. - self._images_to_be_captured = sum([len(line) for line in x_y_grid]) + self._images_to_be_captured = len(path) self._images_captured_so_far = 0 # Generate a basename if none given @@ -236,134 +427,68 @@ class ScanExtension(BaseExtension): "autofocusDz": autofocus_dz, } - # Check if autofocus is enabled - autofocus_extension = find_extension("org.openflexure.autofocus") - if ( - autofocus_dz - and autofocus_extension - and microscope.has_real_stage() - and microscope.has_real_camera() - ): - autofocus_enabled = True + # Perform set-up to be able to autofocus, if needed + if autofocus_dz: + autofocus = self.get_autofocus_function(dz=autofocus_dz, use_fast_autofocus=fast_autofocus) else: - autofocus_enabled = False - + autofocus = None if detect_empty_fields_and_skip_autofocus: - # Check for the background detect extension if we need it, raise an error now if it's missing. - background_detect_extension = find_extension( - "org.openflexure.background-detect" - ) - if not background_detect_extension: - raise RuntimeError( - "Detecting background fields requires the background detect extension and it was not found." - ) - - focused_positions: List[ - XyzCoordinate - ] = [] # Positions where we found sharp images + current_image_is_background = self.get_background_detect_function() + else: + current_image_is_background = None + focus_manager = FocusManager( + microscope, + initial_position, + autofocus_function = autofocus, + current_image_is_background_function = current_image_is_background, + axial_jump_threshold = 0.4 if detect_empty_fields_and_skip_autofocus else None + ) # Now step through each point in the x-y coordinate array - for line in x_y_grid: - for x_y in line: - # Set the next z position based on the closest point that was in focus. - # For a snake/spiral scan, this should always be the last point, unless - # it's skipped for some reason. In a raster scan, this should be the last - # point, except when we're at the start of a row when it will be the first - # point of the preceding row. - closest_focused_point = closest_point_in_xy(x_y, focused_positions) - next_z = ( - closest_focused_point[2] - if closest_focused_point - else initial_position[2] + for x_y in path: + next_z = focus_manager.estimate_z(x_y) + # Move to new grid position + logging.debug("Moving to step %s", ([x_y[0], x_y[1], next_z])) + microscope.stage.move_abs((x_y[0], x_y[1], next_z)) + # Autofocus (if requested) + focus_manager.autofocus() + + # If we're not doing a z-stack, just capture + if grid[2] <= 1: + self.capture( + microscope, + basename, + namemode=namemode, + temporary=temporary, + use_video_port=use_video_port, + resize=resize, + bayer=bayer, + dataset=dataset_d, + annotations=annotations, + tags=tags, ) - # Move to new grid position - logging.debug("Moving to step %s", ([x_y[0], x_y[1], next_z])) - microscope.stage.move_abs((x_y[0], x_y[1], next_z)) - - # Check if the current field looks empty, and skip autofocus if it does - skip_autofocus = False - if detect_empty_fields_and_skip_autofocus: - verdict = background_detect_extension.grab_and_classify_image() - logging.debug(f"Background detection verdict: {verdict}") - if verdict["classification"] == "background": - skip_autofocus = True - logging.info( - f"Detected an empty field at {x_y}, skipping autofocus." - ) - - if autofocus_enabled and not skip_autofocus: - if fast_autofocus: - # Run fast autofocus. Client should provide dz ~ 2000 - autofocus_extension.fast_autofocus(microscope, dz=autofocus_dz) - else: - # Run slow autofocus. Client should provide dz ~ 50 - autofocus_extension.autofocus( - microscope, - range(-3 * autofocus_dz, 4 * autofocus_dz, autofocus_dz), - ) - logging.debug("Finished autofocus") - time.sleep(1) - autofocus_accepted = True - here = microscope.stage.position - # Check if we've moved worryingly far in Z - if closest_focused_point: - lateral_move = np.sqrt( - np.sum( - ( - np.array(here)[:2] - - np.array(closest_focused_point)[:2] - ) - ** 2 - ) - ) - axial_move = np.abs(here[2] - closest_focused_point[2]) - if axial_move > lateral_move * 0.4: - autofocus_accepted = False - logging.warning( - f"During a scan, there was a large axial jump from {closest_focused_point}" - f" to {here}. This may mean autofocus has failed." - ) - # Append the (current, i.e. focused) position to the list - if autofocus_accepted: - focused_positions.append(here) - - # If we're not doing a z-stack, just capture - if grid[2] <= 1: - self.capture( - microscope, - basename, - namemode=namemode, - temporary=temporary, - use_video_port=use_video_port, - resize=resize, - bayer=bayer, - dataset=dataset_d, - annotations=annotations, - tags=tags, - ) - # Update task progress - self._images_captured_so_far += 1 - update_action_progress(self.progress()) - else: - logging.debug("Entering z-stack") - self.stack( - microscope=microscope, - basename=basename, - namemode=namemode, - temporary=temporary, - step_size=stride_size[2], - steps=grid[2], - use_video_port=use_video_port, - resize=resize, - bayer=bayer, - dataset=dataset_d, - annotations=annotations, - tags=tags, - ) - if current_action() and current_action().stopped: - return - # Make sure we use our current best estimate of focus (i.e. the current position) next point - next_z = microscope.stage.position[2] + # Update task progress + self._images_captured_so_far += 1 + update_action_progress(self.progress()) + else: + logging.debug("Entering z-stack") + self.stack( + microscope=microscope, + basename=basename, + namemode=namemode, + temporary=temporary, + step_size=stride_size[2], + steps=grid[2], + use_video_port=use_video_port, + resize=resize, + bayer=bayer, + dataset=dataset_d, + annotations=annotations, + tags=tags, + ) + # Gracefully shut down if we have been requested to stop + if current_action() and current_action().stopped: + return logging.debug("Returning to %s", (initial_position)) microscope.stage.move_abs(initial_position)