"""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 contextlib import contextmanager import logging import time from typing import Annotated, Mapping, Optional, Sequence import os import shutil import glob 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 STACK_OVERSHOOT = 200 SETTLING_TIME = 0.3 class SharpnessDataArrays(BaseModel): jpeg_times: NDArray jpeg_sizes: NDArray stage_times: NDArray stage_positions: list[dict[str, int]] class JPEGSharpnessMonitor: 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 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_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_z_stack( self, cam: WrappedCamera, stage: Stage, logger: InvocationLogger, metadata_getter: GetThingStates, images_dir: str, stack_dir: str, capture_resolution: tuple[int, int], ) -> None: """Run a z stack, saving all images to stack_dir and copying the central image to stack_dir""" stack_dz = self.stack_dz images_to_capture = self.stack_images_to_capture stack_z_range = stack_dz * (images_to_capture - 1) if stack_z_range > 0: # Perform backlash corrected move. See issue #420 stage.move_relative(z=-(STACK_OVERSHOOT + stack_z_range / 2)) stage.move_relative(z=STACK_OVERSHOOT) time.sleep(0.3) for capture_count in range(images_to_capture): time.sleep(SETTLING_TIME) jpeg_path = os.path.join(stack_dir, f"{capture_count}.jpeg") start_time = time.time() cam.capture_to_memory(logger=logger, metadata_getter=metadata_getter) captured_time = time.time() if capture_count + 1 < images_to_capture: stage.move_relative(z=stack_dz) moved_time = time.time() cam.save_from_memory( jpeg_path=jpeg_path, logger=logger, save_resolution=capture_resolution, ) saved_time = time.time() time_remaining = SETTLING_TIME - (saved_time - start_time) if time_remaining > 0: time.sleep(time_remaining) logger.info(f"Settled for an extra {round(time_remaining, 3)} seconds") logger.debug(f"Capturing took {round(captured_time - start_time, 2)} s") logger.debug(f"Saving took {round(saved_time - moved_time, 2)} s") logger.debug( f"Effective settling time was {round(time.time() - moved_time, 2)} s" ) self.copy_sharpest_image_from_stack(images_dir, stack_dir, logger) def copy_central_image_from_stack( self, images_dir: str, stack_dir: str, ): """Gets a list of images in a folder (stack_dir), sorts them, and copies the central image to images dir.""" image_list = glob.glob(os.path.join(stack_dir, "*")) image_list.sort() central_index = (len(image_list) - 1) // 2 central_image = image_list[central_index] xy_location = os.path.basename(stack_dir) shutil.copy(central_image, os.path.join(images_dir, f"{xy_location}.jpeg")) def copy_sharpest_image_from_stack( self, images_dir: str, stack_dir: str, logger: InvocationLogger, ) -> None: """Gets a list of images in a folder (stack_dir), sorts them by filesize, and copies the sharpest image to images dir.""" image_list = glob.glob(os.path.join(stack_dir, "*")) image_list = sorted(image_list, key=os.path.getsize) sharpest_image = image_list[-1] xy_location = os.path.basename(stack_dir) logger.info(sharpest_image) shutil.copy(sharpest_image, os.path.join(images_dir, f"{xy_location}.jpeg"))