StackInfo dataclass

This commit is contained in:
jaknapper 2025-06-23 12:56:31 +01:00
parent 8fea35a663
commit 7cff0f8f44

View file

@ -29,9 +29,35 @@ from .camera import CameraDependency as WrappedCamera
from .stage import StageDependency as Stage from .stage import StageDependency as Stage
from .capture import RawCaptureDependency as CaptureDep from .capture import RawCaptureDependency as CaptureDep
SETTLING_TIME = 0.3 class StackInfo(BaseModel):
BACKLASH_CORRECTION = 250 """Dataclass with default scan parameters"""
# These are generic settings that seem to apply well to a standard scan
# and are not altered by default when running a scan
# time (in seconds) between moving and capturing an image
settling_time: float = 0.3
# distance (in steps) to overshoot a move and then undo, to account for backlash
backlash_correction: int = 250
# how many times the minimum distance between images to include as a "nearby" image
# default 1.4 includes images offset in x or y, but not diagonally
neighbour_cutoff: float = 1.4
# how many images can be appended to the stack after the predicted peak to test for focus
# before assuming the focus was passed, and restarting the stack
stack_height_limit: int = 15
# how far below (in factors of stack_dz) the estimated optimal starting position to
# begin the stack. Better to start slightly too low and require many images, rather than
# too high and needing to autofocus and restart the stack
undershoot: int = 5
# These we expect to be overwritten by AutofocusThing properties
stack_dz: int = 50
images_to_capture: int = 1
images_to_test: int = 5
autofocus_dz: int = 2000
class JPEGSharpnessMonitor: class JPEGSharpnessMonitor:
__globals__ = globals() # Required for FastAPI dependency __globals__ = globals() # Required for FastAPI dependency
@ -308,29 +334,31 @@ class AutofocusThing(Thing):
""" """
# Set the variables to prevent changes from the GUI or other windows # Set the variables to prevent changes from the GUI or other windows
stack_dz = self.stack_dz stack_parameters = StackInfo(
images_to_capture = self.stack_images_to_capture stack_dz = self.stack_dz,
images_to_test = self.stack_images_to_test images_to_capture = self.stack_images_to_capture,
images_to_test = self.stack_images_to_test,
autofocus_dz = autofocus_dz
)
stack_z_range = stack_dz * (images_to_test - 1) stack_z_range = stack_parameters.stack_dz * (stack_parameters.images_to_test - 1)
# Starting too low by "overshoot" makes smart stacking faster. # Starting too low by "undershoot" makes smart stacking faster.
# Starting a stack too high requires it to move to the start, # Starting a stack too high requires it to move to the start,
# autofocus and then re-stack. Starting slightly too low only # autofocus and then re-stack. Starting slightly too low only
# requires extra +z movements and captures. # requires extra +z movements and captures.
overshoot = stack_dz * 5 undershoot = stack_parameters.stack_dz * stack_parameters.undershoot
# Ensure the stack settings are appropriate # Ensure the stack settings are appropriate
self.validate_stack_inputs(images_to_test, images_to_capture) self.validate_stack_inputs(stack_parameters)
success = False success = False
# Loop until a stack is successful # Loop until a stack is successful
while not success: while not success:
result, heights, captures, sharpest_index = self.z_stack( result, heights, captures, sharpest_index = self.z_stack(
stack_dz,
images_to_test,
images_dir, images_dir,
stack_z_range, stack_z_range,
overshoot, undershoot,
stack_parameters,
stage, stage,
cam, cam,
capture, capture,
@ -344,7 +372,7 @@ class AutofocusThing(Thing):
# If a stack is not successful, move to the start and autofocus # If a stack is not successful, move to the start and autofocus
self.reset_stack( self.reset_stack(
heights, heights,
autofocus_dz, stack_parameters.autofocus_dz,
stage, stage,
sharpness_monitor, sharpness_monitor,
) )
@ -353,14 +381,13 @@ class AutofocusThing(Thing):
self.save_stack( self.save_stack(
sharpest_index, sharpest_index,
captures, captures,
images_to_test, stack_parameters,
images_to_capture,
logger, logger,
capture, capture,
) )
# Return the z position of the sharpest image, for path planning and tracking # Return the z position of the sharpest image, for path planning and tracking
return heights[-images_to_test:][sharpest_index] return heights[-stack_parameters.images_to_test:][sharpest_index]
def reset_stack( def reset_stack(
self, self,
@ -388,8 +415,7 @@ class AutofocusThing(Thing):
self, self,
sharpest_index: int, sharpest_index: int,
captures: list[list], captures: list[list],
images_to_test: int, stack_parameters: StackInfo,
images_to_capture: int,
logger: InvocationLogger, logger: InvocationLogger,
capture: CaptureDep, capture: CaptureDep,
) -> int: ) -> int:
@ -399,12 +425,11 @@ class AutofocusThing(Thing):
Arguments: Arguments:
sharpest_index: the index of the sharpest image, within the "images_to_test" slice 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 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 stack_parameters: a StackInfo dataclass with stack settings
images_to_capture: the number of images to save to disk
variables logger and capture are Thing dependencies passed through from the calling action variables logger and capture are Thing dependencies passed through from the calling action
""" """
# Find the range of images from the stack to capture # Find the range of images from the stack to capture
stack_extent = int((images_to_capture - 1) / 2) stack_extent = int((stack_parameters.images_to_capture - 1) / 2)
stack_range = range( stack_range = range(
sharpest_index - stack_extent, sharpest_index + stack_extent + 1 sharpest_index - stack_extent, sharpest_index + stack_extent + 1
) )
@ -412,43 +437,41 @@ class AutofocusThing(Thing):
# Loop through the range, saving each capture to disk # Loop through the range, saving each capture to disk
for capture_index in stack_range: for capture_index in stack_range:
capture._save_capture( capture._save_capture(
jpeg_path=captures[-images_to_test:][capture_index][0], jpeg_path=captures[-stack_parameters.images_to_test:][capture_index][0],
image=captures[-images_to_test:][capture_index][1], image=captures[-stack_parameters.images_to_test:][capture_index][1],
metadata=captures[-images_to_test:][capture_index][2], metadata=captures[-stack_parameters.images_to_test:][capture_index][2],
logger=logger, logger=logger,
) )
return sharpest_index return sharpest_index
def z_stack( def z_stack(
self, self,
stack_dz: int,
images_to_test: int,
images_dir: str, images_dir: str,
stack_z_range: int, stack_z_range: int,
overshoot: int, overshoot: int,
stack_parameters: StackInfo,
stage: Stage, stage: Stage,
cam: WrappedCamera, cam: WrappedCamera,
capture: CaptureDep, capture: CaptureDep,
metadata_getter: GetThingStates, metadata_getter: GetThingStates,
) -> list: ) -> list:
"""Capture a series of images offset by stack_dz, and test whether """Capture a series of images offset by stack_parameters.stack_dz, and test whether
the sharpest image is towards the centre of the stack. 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 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. index of the sharpest image in a successful stack.
Arguments: 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 images_dir: a string of the path to write all images
stack_z_range: the height in steps of the stack to test stack_z_range: the height in steps of the stack to test
overshoot: how far below the estimated optimal starting position to begin stacking overshoot: how far below the estimated optimal starting position to begin stacking
stack_parameters: a StackInfo dataclass with stack settings
variables stage to metadata_getter are Thing dependencies passed through from the calling action 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 # 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 # 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=-(overshoot + stack_parameters.backlash_correction + stack_z_range / 2))
stage.move_relative(z=BACKLASH_CORRECTION) stage.move_relative(z=stack_parameters.backlash_correction)
captures = [] captures = []
sharpnesses = [] sharpnesses = []
@ -456,8 +479,8 @@ class AutofocusThing(Thing):
# If the sharpest image isn't found within the 15 images above the estimated point, break # If the sharpest image isn't found within the 15 images above the estimated point, break
# the loop and return "restart" # the loop and return "restart"
while len(captures) <= images_to_test + 15: while len(captures) <= stack_parameters.images_to_test + 15:
time.sleep(SETTLING_TIME) time.sleep(stack_parameters.settling_time)
# Append a new image to the stack # Append a new image to the stack
captures, heights, sharpnesses = self.capture_stack_image( captures, heights, sharpnesses = self.capture_stack_image(
@ -472,17 +495,17 @@ class AutofocusThing(Thing):
) )
# If the number of images is enough to test, test them # If the number of images is enough to test, test them
if len(captures) >= images_to_test: if len(captures) >= stack_parameters.images_to_test:
stack_result = self.check_stack_result(sharpnesses[-images_to_test:]) stack_result = self.check_stack_result(sharpnesses[-stack_parameters.images_to_test:])
if stack_result == "success": if stack_result == "success":
sharpest_index = np.argmax(sharpnesses[-images_to_test:]) sharpest_index = np.argmax(sharpnesses[-stack_parameters.images_to_test:])
return "success", heights, captures, sharpest_index return "success", heights, captures, sharpest_index
if stack_result == "restart": if stack_result == "restart":
return "restart", heights, None, None return "restart", heights, None, None
stage.move_relative(z=stack_dz) stage.move_relative(z=stack_parameters.stack_dz)
return "restart", heights, None, None return "restart", heights, None, None
def capture_stack_image( def capture_stack_image(
@ -523,7 +546,7 @@ class AutofocusThing(Thing):
return captures, heights, sharpnesses return captures, heights, sharpnesses
def validate_stack_inputs( def validate_stack_inputs(
self, images_to_test: int, images_to_capture: int self, stack_parameters
) -> None: ) -> None:
"""Check the stack settings are appropriate, and raise an error if not """Check the stack settings are appropriate, and raise an error if not
@ -531,13 +554,13 @@ class AutofocusThing(Thing):
images_to_test: number of images in the stack to test for focus 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 images_to_capture: number of images to be captured around the focused image
""" """
if images_to_test < images_to_capture: if stack_parameters.images_to_test < stack_parameters.images_to_capture:
raise RuntimeError( raise RuntimeError(
"Can't capture more images than are tested. Please increase number to test, or decrease number to capture" "Can't capture more images than are tested. Please increase number to test, or decrease number to capture"
) )
if images_to_test % 2 == 0: if stack_parameters.images_to_test % 2 == 0:
raise RuntimeError("Images to test should be odd") raise RuntimeError("Images to test should be odd")
if images_to_test <= 0 or images_to_capture <= 0: if stack_parameters.images_to_test <= 0 or stack_parameters.images_to_capture <= 0:
raise RuntimeError("Stack parameters need to be at least 1") raise RuntimeError("Stack parameters need to be at least 1")
def check_stack_result(self, sharpnesses: list[int]) -> str: def check_stack_result(self, sharpnesses: list[int]) -> str: