From 603b9c7d1627e011418e5bbbbcdebe83aae9d5f5 Mon Sep 17 00:00:00 2001
From: jaknapper <jaknapper@hotmail.com>
Date: Wed, 16 Apr 2025 09:40:14 +0100
Subject: [PATCH] Functional z stack Thing

---
 .../things/smart_scan.py                      |   5 +
 .../things/z_stack.py                         | 137 ++++++++++++++++++
 2 files changed, 142 insertions(+)
 create mode 100644 src/openflexure_microscope_server/things/z_stack.py

diff --git a/src/openflexure_microscope_server/things/smart_scan.py b/src/openflexure_microscope_server/things/smart_scan.py
index 50a8c442..905f5968 100644
--- a/src/openflexure_microscope_server/things/smart_scan.py
+++ b/src/openflexure_microscope_server/things/smart_scan.py
@@ -32,6 +32,7 @@ from openflexure_microscope_server.utilities import ErrorCapturingThread
 from openflexure_microscope_server.things.autofocus import AutofocusThing
 from openflexure_microscope_server.things.camera_stage_mapping import CameraStageMapper
 from openflexure_microscope_server.things.background_detect import BackgroundDetectThing
+from openflexure_microscope_server.things.z_stack import ZStackThing
 from openflexure_microscope_server import scan_planners
 
 CSMDep = direct_thing_client_dependency(CameraStageMapper, "/camera_stage_mapping/")
@@ -139,6 +140,7 @@ class SmartScanThing(Thing):
         self._background_detect: Optional[BackgroundDep] = None
         self._ongoing_scan_name: Optional[str] = None
         self._starting_position: Optional[Mapping[str, int]] = None
+        self._z_stack: Optional[ZStackDep] = None
         self._capture_thread: Optional[ErrorCapturingThread] = None
         self._scan_images_taken: Optional[int] = None
         # TODO Scan data is a dict during refactoring, should become a dataclass
@@ -155,6 +157,7 @@ class SmartScanThing(Thing):
         metadata_getter: GetThingStates,
         csm: CSMDep,
         background_detect: BackgroundDep,
+        z_stack: ZStackDep,
         scan_name: str = "",
     ):
         """Move the stage to cover an area, taking images that can be tiled together.
@@ -178,6 +181,7 @@ class SmartScanThing(Thing):
         self._csm = csm
         self._background_detect = background_detect
         self._capture_thread = None
+        self._z_stack = z_stack
         self._scan_images_taken = 0
 
         # Don't set self._scan_data dictionary. This is done at the start of _run_scan
@@ -213,6 +217,7 @@ class SmartScanThing(Thing):
             self._ongoing_scan_name = None
             self._scan_images_taken = None
             self._scan_data = None
+            self._z_stack = None
             self._scan_lock.release()
 
     def promote_stitch_files(
diff --git a/src/openflexure_microscope_server/things/z_stack.py b/src/openflexure_microscope_server/things/z_stack.py
new file mode 100644
index 00000000..97cd83a7
--- /dev/null
+++ b/src/openflexure_microscope_server/things/z_stack.py
@@ -0,0 +1,137 @@
+import os
+from typing import Mapping, Optional
+import cv2
+import numpy as np
+from PIL import Image
+from pydantic import BaseModel
+from scipy.stats import norm
+import time
+import piexif
+import json
+
+from labthings_fastapi.dependencies.metadata import GetThingStates
+from labthings_fastapi.thing import Thing
+from labthings_fastapi.decorators import thing_action, thing_property
+from .camera import CameraDependency as CamDep
+from .stage import StageDependency as StageDep
+from labthings_fastapi.dependencies.invocation import (
+    CancelHook,
+    InvocationLogger,
+    InvocationCancelledError,
+)
+
+
+class CaptureError(RuntimeError):
+    """An error trying to capture from Picamera"""
+
+class ZStackThing(Thing):
+    @thing_property
+    def 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_height", 9)
+
+    @images_to_capture.setter
+    def images_to_capture(self, value: int) -> None:
+        self.thing_settings["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 smart_stack(
+        self,
+        cam: CamDep,
+        stage: StageDep,
+        images_dir: str,
+        logger: InvocationLogger,
+        metadata_getter: GetThingStates,
+    ) -> None:
+        stack_dz = self.stack_dz
+        images_to_capture = self.images_to_capture
+
+        stack_z_range = stack_dz * (images_to_capture - 1)
+        stage.move_relative(z=-stack_z_range / 2)
+
+        for capture_count in range(images_to_capture):
+            jpeg_path = os.path.join(
+                images_dir,
+                f"{capture_count}.jpeg",
+            )
+            self._capture_and_save(
+                jpeg_path, cam, logger, metadata_getter,
+            )
+            stage.move_relative(z=stack_dz)
+
+    def _capture_and_save(
+        self,
+        jpeg_path: str,
+        cam: CamDep,
+        logger: InvocationLogger,
+        metadata_getter: GetThingStates,
+    ) -> None:
+        """Capture an image and save it to disk
+
+        This will set the event `acquired` once the image has been acquired, so
+        that the stage may be moved while it's saved.
+        """
+        capture_start = time.time()
+        image, metadata = self._capture_image(cam, metadata_getter,)
+        acquisition_time = time.time()
+        self._save_capture(jpeg_path, image, metadata)
+        save_time = time.time()
+        acquisition_duration = round(acquisition_time - capture_start, 1)
+        saving_duration = round(save_time - acquisition_time, 1)
+        logger.info(
+            f"Acquired {jpeg_path} in {acquisition_duration}s then {saving_duration}s saving to disk"
+        )
+
+    def _capture_image(self, cam, metadata_getter) -> tuple[np.ndarray, dict]:
+        """Capture an image in memory and return it with metadata
+        This will set the event `acquired` once the image has been acquired, so
+        that the stage may be moved while it's saved.
+        CaptureError raised if the capture fails for any reason
+        returns tuple with numpy array of image data, and dict of metadata
+        """
+        try:
+            metadata = metadata_getter()
+            image = cam.capture_array()[..., :3]
+        except Exception as e:
+            raise CaptureError("An error occurred while capturing") from e
+        return image, metadata
+
+    def _save_capture(
+        self,
+        jpeg_path: str,
+        image: np.ndarray,
+        metadata: dict,
+    ) -> None:
+        """Saving the captured image and metadata to disk
+        logger warning (via InvocationLogger) is raised if metadata is failed to be added
+        IOError is raised if the file cannot be saved
+        nothing is returned on success"""
+        try:
+            Image.fromarray(image.astype("uint8"), "RGB").save(
+                jpeg_path, quality=95, subsampling=0
+            )
+            try:
+                exif_dict = piexif.load(jpeg_path)
+                exif_dict["Exif"][piexif.ExifIFD.UserComment] = json.dumps(
+                    metadata
+                ).encode("utf-8")
+                piexif.insert(piexif.dump(exif_dict), jpeg_path)
+            except:  # noqa: E722
+                # We need to capture any exception as there are many reasons metadata
+                # might not be added. We warn rather than log the error.
+                logger.warning(f"Failed to add metadata to {jpeg_path}")
+        except Exception as e:
+            raise IOError(f"An error occurred while saving {jpeg_path}") from e