Compete first (untested) pass of trying to make fast stack compatible with new memory captures

src/openflexure_microscope_server/things/autofocus.py

@@ -9,8 +9,9 @@ See repository root for licensing information.
 from contextlib import contextmanager
 import logging
 import time
-from typing import Annotated, Mapping, Optional, Sequence
+from typing import Annotated, Mapping, Optional, Sequence, Literal
 import os
+from dataclasses import dataclass
 
 from fastapi import Depends
 import numpy as np
@@ -26,7 +27,8 @@ 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
+
+# TODO Capture reolution should be save resolution for consistency
 
 
 class StackParams(BaseModel):
@@ -39,7 +41,9 @@ class StackParams(BaseModel):
     is always evaluated with the same number of images. i.e. if min_images_to_test=9,
     then 9 images are captured, if the stack is not well focussed, a 10th image is
     captured and images 2 to 10 are evaluated for focus
-    :param autofocus_dz: The number of steps in a full autofocus (if required)
+    :param autofocus_dz: The number of steps in a full autofocus (when required)
+    :param images_dir: The directory to save images to disk
+    :param capture_resolution: The resolution to save the captures to disk with
 
     """
 
@@ -48,13 +52,15 @@ class StackParams(BaseModel):
     images_to_save: int
     min_images_to_test: int
     autofocus_dz: int
+    images_dir: str
+    capture_resolution: tuple[int, int]
 
     # The following parameters have values that 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
 
@@ -76,7 +82,7 @@ class StackParams(BaseModel):
     @property
     def stack_z_range(self) -> int:
         """The range of the z stack, in steps
-        
+
         Note that this is the range of the minimum number of image captured,
         which is also the range of the images sored in memory that can be
         saved."""
@@ -97,13 +103,58 @@ class StackParams(BaseModel):
 
     @computed_field
     @property
-    def max_images_to_test(self) -> int
+    def max_images_to_test(self) -> int:
         """The maximum number of images that will be captured and tested in a stack
 
         This is 15 images more then the minimum number that are captured.
         """
         return self.min_images_to_test
-    
+
+    def slice_to_save(self, sharpest_index):
+        """Return the slice of images to save given the index of the sharpest image"""
+        images_each_side = (self.images_to_save - 1) // 2
+        return slice(
+            sharpest_index - images_each_side, sharpest_index + images_each_side + 1
+        )
+
+
+@dataclass
+class CaptureInfo:
+    """
+    The information from a capture in a z_stack
+    """
+
+    buffer_id: int
+    position: tuple[int, int, int]
+    sharpness: int
+
+    @property
+    def filename(self) -> str:
+        """The filename for this image generated from the poistion"""
+        return f"{self.position[0]}_{self.position[1]}_{self.position[2]}.jpeg"
+
+
+def _get_capture_by_id(captures: list[CaptureInfo], buffer_id: int) -> CaptureInfo:
+    """Return the capture from a list of CaptureInfo objects with the matching id.
+
+    :param captures: A list of capture objects
+    :param buffer_id: The buffer id of the image to return
+    """
+    return captures[_get_capture_index_by_id(captures, buffer_id)]
+
+
+def _get_capture_index_by_id(captures: list[CaptureInfo], buffer_id: int) -> int:
+    """Return the index of the capture with the matching id from a list of CaptureInfo
+    objects
+
+    :param captures: A list of capture objects
+    :param buffer_id: The buffer id of the image to return
+    """
+    ids = [capture.buffer_id for capture in captures]
+    if buffer_id not in ids:
+        raise ValueError(f"No capture has a buffer id of {buffer_id}")
+    return ids.index(buffer_id)
+
 
 class SharpnessDataArrays(BaseModel):
     jpeg_times: NDArray
@@ -372,7 +423,11 @@ class AutofocusThing(Thing):
         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
+        variables cam to sharpness_monitor are Thing dependencies injected automatically by
+        LabThings FastAPI
+
+        Returns:
+        The z position of the sharpest image
         """
 
         # Set the variables to prevent changes from the GUI or other windows
@@ -381,6 +436,8 @@ class AutofocusThing(Thing):
             images_to_save=self.stack_images_to_save,
             min_images_to_test=self.stack_min_images_to_test,
             autofocus_dz=autofocus_dz,
+            images_dir=images_dir,
+            capture_resolution=capture_resolution,
         )
 
         # Ensure the stack settings are appropriate
@@ -388,22 +445,23 @@ class AutofocusThing(Thing):
 
         success = False
         # Loop until a stack is successful
-        while not success:
+        while not success:  # TODO add a maximum number?
-            result, heights, captures, sharpest_index = self.z_stack(
+            success, captures, sharpest_id = self.z_stack(
-                images_dir,
+                stack_parameters=stack_parameters,
-                stack_parameters,
+                stage=stage,
-                stage,
+                cam=cam,
-                cam,
+                logger=logger,
-                metadata_getter,
+                metadata_getter=metadata_getter,
             )
 
-            if result == "success":
+            if success:
-                success = True
                 break
 
+            # The z position of the first images in the previous attempt.
+            initial_z_pos = captures[0].position[2]
             # If a stack is not successful, move to the start and autofocus
             self.reset_stack(
-                heights,
+                initial_z_pos,
                 stack_parameters.autofocus_dz,
                 stage,
                 sharpness_monitor,
@@ -411,31 +469,33 @@ class AutofocusThing(Thing):
 
         # Save the sharpest image, and images either side of focus
         self.save_stack(
-            sharpest_index,
+            sharpest_id=sharpest_id,
-            captures,
+            captures=captures,
-            stack_parameters,
+            stack_parameters=stack_parameters,
-            logger,
+            cam=cam,
+            logger=logger,
         )
 
         # Return the z position of the sharpest image, for path planning and tracking
-        return heights[-stack_parameters.min_images_to_test :][sharpest_index]
+        return _get_capture_by_id(captures, sharpest_id).position[2]
 
     def reset_stack(
         self,
-        heights: list[int],
+        initial_z_pos: 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.
+        a looping autofocus.
 
         Arguments:
-        heights: a list of the z positions of previous captures
+        initial_z_pos: The initial 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
+        variables stage and sharpness_monitor are Thing dependencies passed through from
+        the calling action
         """
-        stage.move_absolute(z=heights[0])
+        stage.move_absolute(z=initial_z_pos)
         self.looping_autofocus(
             stage=stage,
             sharpness_monitor=sharpness_monitor,
@@ -444,48 +504,52 @@ class AutofocusThing(Thing):
 
     def save_stack(
         self,
-        sharpest_index: int,
+        sharpest_id: int,
         captures: list[list],
         stack_parameters: StackParams,
+        cam: WrappedCamera,
         logger: InvocationLogger,
     ) -> 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 "min_images_to_test" slice
+        sharpest_id: the buffer id index of the sharpest image
         captures: a list of captures, including file name, image data and metadata
         stack_parameters: a StackParams Pydantic model with stack settings
-        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
         """
-        
+
+        sharpest_index = _get_capture_index_by_id(sharpest_id)
+        slice_to_save = stack_parameters.slice_to_save(sharpest_index)
 
         # Loop through the range, saving each capture to disk
-        for capture_index in stack_range:
+        for capture in captures[slice_to_save]:
-            capture._save_capture(
+            cam.save_from_memory(
-                jpeg_path=captures[-stack_parameters.min_images_to_test :][capture_index][
+                jpeg_path=os.path.join(stack_parameters.images_dir, capture.filename),
-                    0
-                ],
-                image=captures[-stack_parameters.min_images_to_test :][capture_index][1],
-                metadata=captures[-stack_parameters.min_images_to_test :][capture_index][2],
                 logger=logger,
+                save_resolution=stack_parameters.capture_resolution,
+                buffer_id=capture.buffer_id,
             )
+        cam.clear_buffers()
         return sharpest_index
 
     def z_stack(
         self,
-        images_dir: str,
         stack_parameters: StackParams,
         stage: Stage,
         cam: WrappedCamera,
-        capture: CaptureDep,
+        logger: InvocationLogger,
         metadata_getter: GetThingStates,
-    ) -> list:
+    ) -> tuple[bool, list[CaptureInfo], Optional[int]]:
         """Capture a series of images offset by stack_parameters.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
+        Returns:
-        index of the sharpest image in a successful stack.
+        - the stack result (True for successful stack, False for failed stack),
+        - a list of CaptureInfo objects,
+        - the buffer_id of the shapest image (or None if the stack failed)
 
         Arguments:
         images_dir: a string of the path to write all images
@@ -504,80 +568,60 @@ class AutofocusThing(Thing):
         stage.move_relative(z=stack_parameters.backlash_correction)
 
         captures = []
-        sharpnesses = []
+        # Always check for focus using the the last `min_images_to_test` in the
-        heights = []
+        # stack so check is fair
+        ims_to_check = slice(-stack_parameters.min_images_to_test, None)
 
-        # 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 maximum number of images
-        # the loop and return "restart"
+        # end the loop and return "restart"
-        while len(captures) <= stack_parameters.min_images_to_test + 15:
+        while len(captures) <= stack_parameters.max_images_to_test:
             time.sleep(stack_parameters.settling_time)
 
             # Append a new image to the stack
-            captures, heights, sharpnesses = self.capture_stack_image(
+            captures.append(
-                captures,
+                self.capture_stack_image(
-                heights,
+                    cam,
-                sharpnesses,
+                    stage,
-                images_dir,
+                    logger,
-                cam,
+                    metadata_getter,
-                stage,
+                    buffer_max=stack_parameters.min_images_to_test,
-                capture,
+                )
-                metadata_getter,
             )
 
             # If the number of images is enough to test, test them
             if len(captures) >= stack_parameters.min_images_to_test:
-                stack_result = self.check_stack_result(
+                result, capture_id = self.check_stack_result(captures[ims_to_check])
-                    sharpnesses[-stack_parameters.min_images_to_test :]
-                )
 
-                if stack_result == "success":
+                if result == "success":
-                    sharpest_index = np.argmax(
+                    return True, captures, capture_id
-                        sharpnesses[-stack_parameters.min_images_to_test :]
-                    )
-                    return "success", heights, captures, sharpest_index
-
-                if stack_result == "restart":
-                    return "restart", heights, None, None
 
+                if result == "restart":
+                    return False, captures, None
+                # If reached here the result was "continue"
             stage.move_relative(z=stack_parameters.stack_dz)
-        return "restart", heights, None, None
+        return False, captures, None
 
     def capture_stack_image(
         self,
-        captures: list[list],
-        heights: list[int],
-        sharpnesses: list[int],
-        images_dir: str,
         cam: WrappedCamera,
         stage: Stage,
-        capture: CaptureDep,
+        logger: InvocationLogger,
         metadata_getter: GetThingStates,
+        buffer_max: int,
     ) -> list:
-        """Append a new capture to the ongoing stack.
+        """Capture another image and return the capture information.
-        Includes appending the height, image sharpness and image
-        data to the relevant lists.
 
-        arguments:
+        The capture is stored by the camera Thing, and can be saved by ID.
-        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(
+        buffer_id = cam.capture_to_memory(
-            images_dir,
+            logger, metadata_getter, buffer_max=buffer_max
-            f"{stage_location['x']}_{stage_location['y']}_{stage_location['z']}.jpeg",
         )
-        image, metadata = capture._capture_image(
+        return CaptureInfo(
-            cam=cam,
+            buffer_id=buffer_id,
-            metadata_getter=metadata_getter,
+            position=stage_location,
+            sharpness=cam.grab_jpeg_size(stream_name="lores"),
         )
-        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, stack_parameters) -> None:
         """Check the stack settings are appropriate, and raise an error if not
@@ -598,38 +642,44 @@ class AutofocusThing(Thing):
         ):
             raise RuntimeError("Stack parameters need to be at least 1")
 
-    def check_stack_result(self, sharpnesses: list[int]) -> str:
+    def check_stack_result(
-        """Test a list of sharpnesses, to decide whether the sharpest image from a stack is within them
+        self, captures: list[CaptureInfo]
+    ) -> tuple[Literal["success", "continue", "restart"], int]:
+        """Test a list of captures, to decide whether the sharpest image from a
+        stack is within them
 
-        Returns a string
+        Returns two values,
-        'success' if the sharpest image is towards the centre
+        result - which is one of three literal values:
-        'continue' if the sharpest image is in the final two images of the list
+          'success' if the sharpest image is towards the centre
-        'restart' if the sharpest image is in the first two images of the list
+          '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
+        capture_id - the buffer id of the sharpest image
 
         Arguments:
         sharpnesses: a list of the sharpnesses to test for focus
         """
-
+        sharpest_index = np.argmax([capture.sharpness for capture in captures])
-        sharpest_index = np.argmax(sharpnesses)
+        # The buffer id of the sharpest image
-        sharpness_length = len(sharpnesses)
+        capture_id = captures[sharpest_index].buffer_id
+        sharpness_length = len(captures)
 
         # If only testing one image, then by definition the sharpest is central
         if sharpness_length == 1:
-            return "success"
+            return "success", capture_id
         # If testing three images, test if the centre is the sharpest
         if sharpness_length == 3:
             if sharpest_index == 1:
-                return "success"
+                return "success", capture_id
             if sharpest_index == 0:
-                return "restart"
+                return "restart", capture_id
-            return "continue"
+            return "continue", capture_id
 
         # 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"
+            return "restart", capture_id
         if sharpest_index >= sharpness_length - exclusion_range:
-            return "continue"
+            return "continue", capture_id
-        return "success"
+        return "success", capture_id

src/openflexure_microscope_server/things/camera/__init__.py

@@ -7,7 +7,7 @@ See repository root for licensing information.
 """
 
 from __future__ import annotations
-from typing import Literal, Optional, Tuple
+from typing import Literal, Optional, Tuple, Any
 import json
 
 from pydantic import RootModel
@@ -50,13 +50,78 @@ class NoImageInMemoryError(RuntimeError):
     """An error called if no image in in memory when an method is called to use that image"""
 
 
+class CameraMemoryBuffer:
+    """
+    A class that holds images in memory. The images are by default PIL images.
+
+    However subclasses of BaseCamera can use this class to store other object types
+    """
+
+    _storage: dict[int, tuple[Any, Optional[dict]]]
+
+    def __init__(self):
+        self._storage = {}
+        self._latest_id: int = 0
+
+    def add_image(
+        self, image: Any, metadata: Optional[dict], buffer_max: int = 1
+    ) -> None:
+        self._latest_id += 1
+        self._create_space(buffer_max)
+        self._storage[self._latest_id] = (image, metadata)
+        return self._latest_id
+
+    def get_image(
+        self, buffer_id: Optional[int] = None, remove: bool = True
+    ) -> tuple[Any, Optional[dict]]:
+        """
+        If the buffer ID is not set, always clear whole buffer
+        """
+
+        # No id given
+        if buffer_id is None:
+            # Get the latest image and metadata tuple from storage
+            try:
+                image_tuple = list(self._storage.value())[-1]
+            except IndexError as e:
+                raise NoImageInMemoryError("No image in memory to save.") from e
+            # Clear the storage so images don't get retrieved out of order
+            self._storage.clear()
+            return image_tuple
+
+        if remove:
+            return self._storage.pop(buffer_id)
+        return self._storage[buffer_id]
+
+    def clear(self):
+        self._storage.clear()
+
+    def _create_space(self, buffer_max: int) -> None:
+        """
+        Create space to add an image.
+        """
+        # If only one image to be stored just clear the storage and return
+        if buffer_max <= 1:
+            self._storage.clear()
+            return
+
+        # Number to remove to get the storage down to 1 less than the buffer length
+        to_remove = len(self._storage) - (buffer_max - 1)
+        # If if there is space. Nothing to do, just return
+        if to_remove < 1:
+            return
+
+        keys_to_remove = list(self._storage.keys())[:to_remove]
+        for key in keys_to_remove:
+            del self._storage[key]
+
+
 class BaseCamera(Thing):
     """The base class for all cameras. All cameras must directly inherit from this class"""
 
-    _memory_image: Optional[Image] = None
-    _memory_metadata: Optional[dict] = None
     mjpeg_stream = MJPEGStreamDescriptor()
     lores_mjpeg_stream = MJPEGStreamDescriptor()
+    _memory_buffer = CameraMemoryBuffer()
 
     def __enter__(self) -> None:
         raise NotImplementedError("CameraThings must define their own __enter__ method")
@@ -64,17 +129,6 @@ class BaseCamera(Thing):
     def __exit__(self, _exc_type, _exc_value, _traceback) -> None:
         raise NotImplementedError("CameraThings must define their own __exit__ method")
 
-    @thing_property
-    def image_in_memory(self) -> bool:
-        """True if an image is in memory ready to be saved"""
-        return self._memory_image is not None and self._memory_metadata is not None
-
-    @thing_action
-    def clear_image_memory(self) -> None:
-        """Clear any image in memory"""
-        self._memory_image = None
-        self._memory_metadata = None
-
     @thing_action
     def start_streaming(
         self, main_resolution: tuple[int, int], buffer_count: int
@@ -133,6 +187,18 @@ class BaseCamera(Thing):
         frame = portal.call(stream.grab_frame)
         return JPEGBlob.from_bytes(frame)
 
+    @thing_action
+    def grab_jpeg_size(
+        self,
+        portal: BlockingPortal,
+        stream_name: Literal["main", "lores"] = "main",
+    ) -> int:
+        """Acquire one image from the preview stream and return its size"""
+        stream = (
+            self.lores_mjpeg_stream if stream_name == "lores" else self.mjpeg_stream
+        )
+        return portal.call(stream.next_frame_size)
+
     @thing_action
     def capture_image(
         self,
@@ -173,20 +239,18 @@ class BaseCamera(Thing):
 
     @thing_action
     def capture_to_memory(
-        self,
+        self, logger: InvocationLogger, metadata_getter: GetThingStates, buffer_max
-        logger: InvocationLogger,
-        metadata_getter: GetThingStates,
     ) -> None:
         """
         Capture an image to memory. This can be saved later with `save_from_memory`
 
         Note that only one image is held in memory so this will overwrite any image
         in memory.
+
+        Return the buffer id of the image captured
         """
-        self._memory_image, self._memory_metadata = self._robust_image_capture(
+        image, metadata = self._robust_image_capture(metadata_getter, logger)
-            metadata_getter,
+        return self._memory_buffer.add_image(image, metadata, buffer_max=buffer_max)
-            logger=logger,
-        )
 
     @thing_action
     def save_from_memory(
@@ -194,21 +258,21 @@ class BaseCamera(Thing):
         jpeg_path: str,
         logger: InvocationLogger,
         save_resolution: Optional[Tuple[int, int]] = None,
+        buffer_id: Optional[int] = None,
     ) -> None:
         """
         Save an image that has been captured to memory.
         """
-        if not self.image_in_memory:
+        image, metadata = self._memory_buffer.get_image(buffer_id)
-            raise NoImageInMemoryError("No image in memory to save.")
 
         self._save_capture(
             jpeg_path=jpeg_path,
-            image=self._memory_image,
+            image=image,
-            metadata=self._memory_metadata,
+            metadata=metadata,
             logger=logger,
             save_resolution=save_resolution,
         )
-        self.clear_image_memory()
+        self._memory_buffer.clear()
 
     def _robust_image_capture(
         self,
@@ -248,10 +312,13 @@ class BaseCamera(Thing):
         if save_resolution is not None and image.size != save_resolution:
             image = image.resize(save_resolution, Image.BOX)
         try:
-            # Per PIL documentation, (https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#jpeg)
+            # Per PIL documentation,
-            # there are two factors when saving a JPEG. subsampling affects the colour, quality the pixels
+            # (https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#jpeg)
+            # there are two factors when saving a JPEG. Subsampling affects the colour,
+            # quality affects the pixels.
             # subsampling = 0 disables subsampling of colour
-            # quality = 95 is the maximum recommended - above this, JPEG compression is disabled, file size increases and quality is barely or not affected
+            # quality = 95 is the maximum recommended - above this, JPEG compression is
+            # disabled, file size increases and quality is barely or not affected
             image.save(jpeg_path, quality=95, subsampling=0)
             try:
                 # Load EXIF metadata from image so it can be added to.

src/openflexure_microscope_server/things/smart_scan.py

@@ -665,7 +665,7 @@ class SmartScanThing(Thing):
             )
 
             route_planner.mark_location_visited(
-                current_pos_xyz, imaged=True, focused=focused
+                current_pos_xyz, imaged=True, focused=True
             )
 
             site_folder = os.path.join(