Split up z stack function

src/openflexure_microscope_server/things/autofocus.py

@@ -288,7 +288,7 @@ class AutofocusThing(Thing):
         self.thing_settings["stack_dz"] = value
 
     @thing_action
-    def run_z_stack(
+    def run_smart_stack(
         self,
         cam: WrappedCamera,
         stage: Stage,
@@ -305,72 +305,75 @@ class AutofocusThing(Thing):
         images_to_capture = self.stack_images_to_capture
         images_to_test = self.stack_images_to_test
 
-        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")
-
         stack_z_range = stack_dz * (images_to_test - 1)
         overshoot = stack_dz * 5
 
-        # Move down by the height of the z stack, plus an overshoot
+        self.validate_scan_inputs(images_to_test, images_to_capture)
-        # 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)
 
-        continue_stack = True
+        success = False
-        captures = []
+        while not success:
-        sharpnesses = []
+            result, heights, captures, sharpest_index = self.z_stack(
-        heights = []
+                stack_dz,
-        capture_count = 0
+                images_to_test,
-        starting_height = stage.position["z"]
-
-        while continue_stack:
-            time.sleep(SETTLING_TIME)
-            stage_location = stage.position
-
-            jpeg_path = os.path.join(
                 images_dir,
-                f"{stage_location['x']}_{stage_location['y']}_{stage_location['z']}.jpeg",
+                stack_z_range,
+                overshoot,
+                stage,
+                cam,
+                capture,
+                metadata_getter,
+                logger,
             )
-            image, metadata = capture._capture_image(
+
-                cam=cam,
+            if result == "success":
-                metadata_getter=metadata_getter,
+                success = True
+                break
+
+            # If result remains as "restart", move to the start and autofocus
+            self.reset_stack(
+                heights,
+                stack_z_range,
+                overshoot,
+                autofocus_dz,
+                stage,
+                sharpness_monitor,
             )
-            captures.append([jpeg_path, image, metadata])
-            sharpnesses.append(cam.grab_jpeg_size(stream_name="lores"))
-            heights.append(stage.position["z"])
-            capture_count += 1
 
-            # If the stack isn't complete yet, move
+        self.save_stack(
-            if capture_count >= images_to_test:
+            sharpest_index,
-                stack_result = self.test_stack(sharpnesses[-images_to_test:])
+            captures,
+            images_to_test,
+            images_to_capture,
+            logger,
+            capture,
+        )
 
-                if stack_result == "success":
+        return heights[-images_to_test:][sharpest_index]
-                    continue_stack = False
 
-                # These are signs of overshooting. Autofocus and retry
+    def reset_stack(
-                elif stack_result == "restart" or capture_count > 20:
+        self,
-                    captures = []
+        heights,
-                    sharpnesses = []
+        autofocus_dz,
-                    heights = []
+        stage,
-                    capture_count = 0
+        sharpness_monitor,
-                    stage.move_absolute(z=starting_height - stack_z_range)
+    ):
-                    self.looping_autofocus(
+        stage.move_absolute(z=heights[0])
-                        stage=stage,
+        self.looping_autofocus(
-                        sharpness_monitor=sharpness_monitor,
+            stage=stage,
-                        dz=autofocus_dz,
+            sharpness_monitor=sharpness_monitor,
-                    )
+            dz=autofocus_dz,
-                    stage.move_relative(
+        )
-                        z=-(overshoot + BACKLASH_CORRECTION + stack_z_range / 2)
-                    )
-                    stage.move_relative(z=BACKLASH_CORRECTION)
-            stage.move_relative(z=stack_dz)
 
+    def save_stack(
+        self,
+        sharpest_index,
+        captures,
+        images_to_test,
+        images_to_capture,
+        logger,
+        capture,
+    ):
         # Find the range of images from the stack to capture
-        sharpest_index = np.argmax(sharpnesses[-images_to_test:])
         stack_extent = int((images_to_capture - 1) / 2)
         stack_range = range(
             sharpest_index - stack_extent, sharpest_index + stack_extent + 1
@@ -384,8 +387,91 @@ class AutofocusThing(Thing):
                 metadata=captures[-images_to_test:][capture_index][2],
                 logger=logger,
             )
+        return sharpest_index
 
-        return heights[-images_to_test:][sharpest_index]
+    def z_stack(
+        self,
+        stack_dz,
+        images_to_test,
+        images_dir,
+        stack_z_range,
+        overshoot,
+        stage,
+        cam,
+        capture,
+        metadata_getter,
+    ):
+        # 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 = []
+
+        while len(captures) <= 20:
+            time.sleep(SETTLING_TIME)
+
+            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.test_stack(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,
+        heights,
+        sharpnesses,
+        images_dir,
+        cam,
+        stage,
+        capture,
+        metadata_getter,
+    ):
+        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_scan_inputs(self, images_to_test, images_to_capture):
+        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")
 
     def test_stack(self, sharpnesses: list):
         """Test a list of sharpnesses, to decide whether the sharpest image from a stack is within them

src/openflexure_microscope_server/things/smart_scan.py

@@ -642,7 +642,7 @@ class SmartScanThing(Thing):
                 self._scan_logger.info(msg)
                 continue
 
-            focused_height = self._autofocus.run_z_stack(
+            focused_height = self._autofocus.run_smart_stack(
                 images_dir=self._ongoing_scan_images_dir,
                 autofocus_dz=self._scan_data["autofocus_dz"],
             )