Pull stitching out of SmartScanThing

src/openflexure_microscope_server/scan_directories.py

@@ -1,11 +1,12 @@
 """Functionality to manage file system operations for scan directories."""
 
-from typing import Optional
+from typing import Optional, Any
 import os
 import re
 import shutil
 import zipfile
 import threading
+import json
 from datetime import datetime, timedelta
 
 from pydantic import BaseModel, field_validator, field_serializer
@@ -59,7 +60,9 @@ class ScanData(BaseModel):
     start_time: datetime
     skip_background: str
     stitch_automatically: str
-    stitch_resize: int
+    # TODO: Think about changing stitch_resize name as it is NOT the resize for
+    # just for correlation stitching
+    stitch_resize: float
     save_resolution: tuple[int, int]
     final_image_count: Optional[int] = None
     duration: Optional[timedelta] = None
@@ -215,6 +218,21 @@ class ScanDirectoryManager:
             return None
         return scan_data_path
 
+    def get_scan_data_dict(self, scan_name: str) -> Optional[dict[str, Any]]:
+        """Return the scan data read from a JSON file as a dict.
+
+        This is a dictionary not a base models as the data format has changed
+        somewhat over time.
+        """
+        json_fpath = self.get_scan_data_path(scan_name)
+        if json_fpath is None:
+            return None
+        try:
+            with open(json_fpath, "r", encoding="utf-8") as data_file:
+                return json.load(data_file)
+        except (json.decoder.JSONDecodeError, IOError):
+            return None
+
     @property
     @requires_lock
     def all_scans(self) -> list[str]:

src/openflexure_microscope_server/stitching.py

@@ -0,0 +1,295 @@
+"""Communicate with OpenFlexure Stitching to perform stitches for scans.
+
+This includes both live stitching and final stitching. This is done via subprocess
+to call openflexure-stitching over CLI. This cannot be done via Threading dut to the
+CPU intensity of stitching causing scanning problems due to the Python Global
+Interpreter Lock (GIL). May be possible to shift to multiprocessing int the future.
+"""
+
+from typing import Optional, Any
+import threading
+import subprocess
+import os
+import time
+
+import labthings_fastapi as lt
+
+STITCHING_CMD = "openflexure-stitch"
+STITCHING_RESOLUTION = (820, 616)
+
+DEFAULT_OVERLAP = 0.5
+DEFAULT_RESIZE = 0.5
+
+
+class BaseStitcher:
+    """A base stitching class for all stitchers. Don't initialise this directly."""
+
+    def __init__(self, images_dir: str, *, overlap: float, correlation_resize: float):
+        """Initialise a stitcher.
+
+        All args except images_dir are positional only.
+
+        :param images_dir: The images directory of the scan to stitch.
+        :param overlap: The scan overlap.
+        :param correlation_resize: The fraction to resize images by when correlating.
+        """
+        # Set minimum overlap to 90% of the scan overlap to catch only images
+        # directly adjacent, not images with overlapping corners.
+        self.images_dir = images_dir
+        self.min_overlap = round(overlap * 0.9, 2)
+        self.correlation_resize = correlation_resize
+        self._mode = "all"
+        self._extra_args = []
+
+    @property
+    def command(self) -> list[str]:
+        """The command to run with subprocess.Popen."""
+        # The command, and the mode
+        initial_args = [STITCHING_CMD, "--stitching_mode", self._mode]
+        setting_args = [
+            "--minimum_overlap",
+            f"{self.min_overlap}",
+            "--resize",
+            f"{self.correlation_resize}",
+        ]
+        return initial_args + self._extra_args + setting_args + [self.images_dir]
+
+    def start(self) -> None:
+        """Start stitching a stitching process.
+
+        This should be overridden by any child class.
+        """
+        raise NotImplementedError(
+            "Child stitchers should implement their own ``start`` method."
+        )
+
+
+class PreviewStitcher(BaseStitcher):
+    """A stitcher for stitching an ongoing scan in preview mode.
+
+    Use ``start()`` to start a scan, and ``running`` to check if it is complete, or
+    ``wait()`` to wait for it to complete.
+
+    The same stitcher object can be run multiple times to update the preview. However,
+    one preview must finish before another can be started.
+    """
+
+    def __init__(self, images_dir: str, *, overlap: float, correlation_resize: float):
+        """Initialise a preview stitcher.
+
+        All args except images_dir are positional only.
+
+        :param images_dir: The images directory of the scan to stitch.
+        :param overlap: The scan overlap.
+        :param correlation_resize: The fraction to resize images by when correlating.
+        """
+        super().__init__(
+            images_dir, overlap=overlap, correlation_resize=correlation_resize
+        )
+        self._popen_lock = threading.Lock()
+        self._popen_obj: Optional[subprocess.Popen] = None
+
+        self._mode = "preview_stitch"
+
+    @property
+    def command(self) -> list[str]:
+        """The command to run with subprocess.Popen."""
+        return [
+            STITCHING_CMD,
+            "--stitching_mode",
+            "preview_stitch",
+            "--minimum_overlap",
+            f"{self.min_overlap}",
+            "--resize",
+            f"{self.correlation_resize}",
+            self.images_dir,
+        ]
+
+    def start(self) -> None:
+        """Start stitching a preview of the scan in a background subprocess.
+
+        This uses popen and returns immediately.
+        """
+        if self.running:
+            raise RuntimeError("Cannot start stitch. It is already running.")
+        with self._popen_lock:
+            self._popen_obj = subprocess.Popen(self.command)
+
+    @property
+    def running(self) -> bool:
+        """Whether the preview stitch is running in a subprocess."""
+        with self._popen_lock:
+            if self._popen_obj is None:
+                return False
+            if self._popen_obj.poll() is None:
+                return True
+            return False
+
+    def wait(self):
+        """Wait for this preview stitch to return."""
+        if self.running:
+            with self._popen_lock:
+                self._popen_obj.wait()
+
+
+class FinalStitcher(BaseStitcher):
+    """A class to handle the final stich for a scan."""
+
+    def __init__(
+        self,
+        images_dir,
+        *,
+        logger: lt.deps.InvocationLogger,
+        overlap: Optional[float] = None,
+        correlation_resize: Optional[float] = None,
+        stitch_tiff: bool = False,
+        scan_data_dict: Optional[dict[str, Any]] = None,
+    ):
+        """Initialise a final stitcher, this has more args than the base class.
+
+        All args except images_dir are positional only.
+
+        :param images_dir: The images directory of the scan to stitch.
+        :param logger: The invocation logger for this stitch process.
+        :param overlap: The scan overlap, if not known enter None. A value will be
+            chosen from the scan_data_dict, or set to a default value if no value is
+            available in the scan_data.
+        :param correlation_resize: The fraction to resize images by when correlating,
+            if not known enter None. A value will be chosen from the scan_data_dict, or
+            set to a default value if no value is available in the scan_data.
+        :param stitch_tiff: Whether to stitch a pyramidal TIFF.
+        :param scan_data_dict: The ScanData for this scan as a dictionary. This is used
+            to read/calculate overlap and correlation_resize if they are not provided.
+        """
+        self.logger = logger
+        overlap, correlation_resize = self._process_inputs(
+            overlap=overlap,
+            correlation_resize=correlation_resize,
+            scan_data_dict=scan_data_dict,
+        )
+        super().__init__(
+            images_dir, overlap=overlap, correlation_resize=correlation_resize
+        )
+        self._mode = "all"
+
+        tiff_arg = "--stitch_tiff" if stitch_tiff else "--no-stitch_tiff"
+        self._extra_args = ["--stitch-dzi", tiff_arg]
+
+    def _process_inputs(
+        self,
+        overlap: float,
+        correlation_resize: float,
+        scan_data_dict: Optional[dict[str, Any]],
+    ) -> tuple[float, float]:
+        """Process inputs to ensure ``overlap`` and ``correlation_resize`` have values.
+
+        First the scan_data_dict is inspected for values to allow ``overlap`` and
+        ``correlation_resize`` to be set correctly, if these values are not available
+        then default values are used, and a warning is logged to the invocation logger.
+
+        :param overlap: overlap as input to __init__
+        :param correlation_resize: correlation_resize as input to __init__
+        :param scan_data_dict: scan_data_dict as input to __init__
+
+        :returns: overlap and correlation_resize as floats.
+        """
+        if overlap is None:
+            if scan_data_dict is not None and "overlap" in scan_data_dict:
+                overlap = scan_data_dict["overlap"]
+
+            # Warn if still None and set to default.
+            if overlap is None:
+                overlap = DEFAULT_OVERLAP
+                self.logger.warning(
+                    "No value set for overlap. Attempting stitch with overlap "
+                    f"value of {DEFAULT_OVERLAP}"
+                )
+
+        if correlation_resize is None:
+            if scan_data_dict is not None:
+                # Handle "capture resolution" being used to store the save resolution
+                # in old scans.
+                key = (
+                    "capture resolution"
+                    if "capture resolution" in scan_data_dict
+                    else "save_resolution"
+                )
+                if key in scan_data_dict:
+                    save_resolution = scan_data_dict[key]
+                    correlation_resize = STITCHING_RESOLUTION[0] / save_resolution[0]
+            # Warn if still None and set to default.
+            if correlation_resize is None:
+                correlation_resize = DEFAULT_RESIZE
+                self.logger.warning(
+                    "No information available to calculate stitch resize. Attempting "
+                    f"stitch with resize value of {DEFAULT_RESIZE}"
+                )
+        return overlap, correlation_resize
+
+    def start(
+        self,
+        cancel: lt.deps.CancelHook,
+    ) -> None:
+        """Run the final stitch logging any output.
+
+        Raises:
+            ChildProcessError if exit code is not zero
+            InvocationCancelledError if the action is cancelled.
+
+        """
+        cmd = self.command
+        self.logger.debug(f"Running command in subprocess: `{' '.join(cmd)}`")
+
+        # Run the command piping stdout into the process for reading and
+        # forwarding the stdrerr to stdout
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            bufsize=1,
+            universal_newlines=True,
+        )
+        # Stop opening pipe blocking writing to it
+        os.set_blocking(process.stdout.fileno(), False)
+
+        # TODO check if we still want this? Is it for debugging, or should it have
+        # more explanation.
+        self.logger.info(
+            time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time()))
+        )
+
+        self._log_ongoing(process, cancel=cancel)
+
+        if process.poll() == 0:
+            self.logger.info("Stitching complete")
+        else:
+            raise ChildProcessError(f"Subprocess {cmd[0]} exited with an error.")
+
+    def _log_ongoing(
+        self,
+        process: subprocess.Popen,
+        cancel: lt.deps.CancelHook,
+    ) -> None:
+        """Log the ongoing process unless it is cancelled."""
+        # Poll returns None while running, will return the error code when finished
+        while process.poll() is None:
+            self.log_buffer(process.stdout)
+            # Once buffer is clear sleep for 0.2s before trying again.
+
+            try:
+                # Note that using cancel.sleep allows the InvocationCancelledError to
+                # be thrown
+                cancel.sleep(0.2)
+            except lt.exceptions.InvocationCancelledError as e:
+                self.logger.info("Stitching cancelled by user")
+                process.kill()
+                raise e
+
+        # Print everything in the buffer when program finishes
+        self.log_buffer(process.stdout)
+
+    def log_buffer(self, buffer):
+        # TODO work out type
+        """Log everything in the buffer at INFO level."""
+        while line := buffer.readline():
+            self.logger.info(line)

src/openflexure_microscope_server/things/smart_scan.py

@@ -11,8 +11,7 @@ from typing import Optional, Mapping
 import threading
 import os
 import time
-import json
-from subprocess import CompletedProcess, Popen, PIPE, SubprocessError, STDOUT
+from subprocess import SubprocessError
 
 from fastapi import HTTPException
 from fastapi.responses import FileResponse
@@ -24,6 +23,7 @@ import labthings_fastapi as lt
 from openflexure_microscope_server.utilities import ErrorCapturingThread
 from openflexure_microscope_server import scan_directories
 from openflexure_microscope_server import scan_planners
+from openflexure_microscope_server import stitching
 
 # Things
 from .autofocus import AutofocusThing
@@ -40,11 +40,6 @@ JPEGBlob = lt.blob.blob_type("image/jpeg")
 ZipBlob = lt.blob.blob_type("application/zip")
 
 
-STITCHING_CMD = "openflexure-stitch"
-
-STITCHING_RESOLUTION = (820, 616)
-
-
 class ScanNotRunningError(RuntimeError):
     """Exception called when scan not running that requires a scan to be running."""
 
@@ -85,8 +80,6 @@ class SmartScanThing(lt.Thing):
             HTTP interface.
         """
         self._scan_dir_manager = scan_directories.ScanDirectoryManager(scans_folder)
-        self._preview_stitch_popen = None
-        self._preview_stitch_popen_lock = threading.Lock()
         self._scan_lock = threading.Lock()
 
         # Variables set by the scan
@@ -106,11 +99,13 @@ class SmartScanThing(lt.Thing):
         self._metadata_getter: Optional[lt.deps.GetThingStates] = None
         self._csm: Optional[CSMDep] = None
         self._ongoing_scan: Optional[scan_directories.ScanDirectory] = None
+        # TODO see if starting position can go into ScanData
         self._starting_position: Optional[Mapping[str, int]] = 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
         self._scan_data: Optional[scan_directories.ScanData] = None
+        self._preview_stitcher: Optional[stitching.PreviewStitcher] = None
 
     @lt.thing_action
     def sample_scan(
@@ -181,6 +176,8 @@ class SmartScanThing(lt.Thing):
             self._scan_images_taken = None
             self._scan_data = None
             self._scan_lock.release()
+            # Ensure any PreviewStitcher created cannot be reused.
+            self._preview_stitcher = None
 
     @_scan_running
     def _check_background_and_csm_set(self):
@@ -284,7 +281,7 @@ class SmartScanThing(lt.Thing):
         """Collect and return the data for this scan so it cannot be changed mid-scan."""
         overlap = self.overlap
         dx, dy = self._calc_displacement_from_test_image(overlap)
-        stitch_resize = STITCHING_RESOLUTION[0] // self.save_resolution[0]
+        stitch_resize = stitching.STITCHING_RESOLUTION[0] / self.save_resolution[0]
 
         self._scan_logger.debug(
             f"Resizing images when stitching by a factor of {stitch_resize}"
@@ -339,8 +336,8 @@ class SmartScanThing(lt.Thing):
         # Assume 4 images means at least one offset in x and y, making the stitching
         # well constrained.
         if self._scan_images_taken > 3:
-            if not self._preview_stitch_running():
-                self._preview_stitch_start(overlap=self._scan_data.overlap)
+            if not self._preview_stitcher.running:
+                self._preview_stitcher.start()
 
     @_scan_running
     def _run_scan(self):
@@ -358,6 +355,12 @@ class SmartScanThing(lt.Thing):
             self._cam.start_streaming(main_resolution=(3280, 2464))
             self._scan_data = self._collect_scan_data()
             self._ongoing_scan.save_scan_data(self._scan_data)
+            self._preview_stitcher = stitching.PreviewStitcher(
+                self._ongoing_scan.images_dir,
+                overlap=self._scan_data.overlap,
+                correlation_resize=self._scan_data.stitch_resize,
+            )
+
             if self._scan_images_taken != 0:
                 msg = "_scan_images_taken should be zero before starting scanning"
                 raise RuntimeError(msg)
@@ -390,6 +393,8 @@ class SmartScanThing(lt.Thing):
             )
             raise e
         finally:
+            # Don't set Preview Stitcher to None yet. It is used by _perform_final_scan.
+
             # Start streaming in the default resolution again as soon as possible
             self._cam.start_streaming()
             if self._capture_thread:
@@ -503,19 +508,17 @@ class SmartScanThing(lt.Thing):
 
         self._scan_logger.info("Waiting for background processes to finish...")
 
-        self._preview_stitch_wait()
-        try:
-            if self._scan_data.stitch_automatically:
-                self._scan_logger.info("Stitching final image (may take some time)...")
-                self.stitch_scan(
-                    logger=self._scan_logger,
-                    cancel=self._cancel,
-                    scan_name=self._ongoing_scan.name,
-                    stitch_resize=self._scan_data.stitch_resize,
-                    overlap=self._scan_data.overlap,
-                )
-        except SubprocessError as e:
-            self._scan_logger.error(f"Stitching failed: {e}", exc_info=e)
+        self._preview_stitcher.wait()
+
+        if self._scan_data.stitch_automatically:
+            self._scan_logger.info("Stitching final image (may take some time)...")
+            self.stitch_scan(
+                logger=self._scan_logger,
+                cancel=self._cancel,
+                scan_name=self._ongoing_scan.name,
+                stitch_resize=self._scan_data.stitch_resize,
+                overlap=self._scan_data.overlap,
+            )
 
     @lt.fastapi_endpoint(
         "get",
@@ -727,103 +730,6 @@ class SmartScanThing(lt.Thing):
             raise HTTPException(404, "File not found")
         return FileResponse(preview_path)
 
-    @_scan_running
-    def _preview_stitch_start(self, overlap: float) -> None:
-        """Start stitching a preview of the scan in a background subprocess.
-
-        This uses popen and returns immediately
-
-        - self._preview_stitch_popen holds the popen for polling
-        - self._preview_stitch_popen_lock is a lock acquired while interacting
-              with Popen
-        """
-        # Set minimum overlap to 90% of the scan overlap to catch only images directly adjacent,
-        # not images with overlapping corners.
-        min_overlap = round(overlap * 0.9, 2)
-        if self._preview_stitch_running():
-            raise RuntimeError("Only one subprocess is allowed at a time")
-        with self._preview_stitch_popen_lock:
-            self._preview_stitch_popen = Popen(
-                [
-                    STITCHING_CMD,
-                    "--stitching_mode",
-                    "preview_stitch",
-                    "--minimum_overlap",
-                    f"{min_overlap}",
-                    "--resize",
-                    f"{self._scan_data.stitch_resize}",
-                    self._ongoing_scan.images_dir,
-                ]
-            )
-
-    @_scan_running
-    def _preview_stitch_running(self) -> bool:
-        """Whether there is a preview stitch running in a subprocess."""
-        with self._preview_stitch_popen_lock:
-            if self._preview_stitch_popen is None:
-                return False
-            if self._preview_stitch_popen.poll() is None:
-                return True
-            return False
-
-    @_scan_running
-    def _preview_stitch_wait(self):
-        """Wait for an ongoing preview stitch to return."""
-        if self._preview_stitch_running():
-            with self._preview_stitch_popen_lock:
-                self._preview_stitch_popen.wait()
-
-    def run_subprocess(
-        self,
-        logger: lt.deps.InvocationLogger,
-        cancel: lt.deps.CancelHook,
-        cmd: list[str],
-    ) -> CompletedProcess:
-        """Run a  subprocess and log any output.
-
-        Raises:
-            ChildProcessError if exit code is not zero
-            InvocationCancelledError if the action is cancelled.
-
-        """
-        logger.info(f"Running command in subprocess: `{' '.join(cmd)}`")
-
-        def log_buffer(buffer):
-            """Log everything in the buffer at INFO level."""
-            while line := buffer.readline():
-                logger.info(line)
-
-        # Run the command piping stdout into the process for reading and
-        # forwarding the stdrerr to stdout
-        process = Popen(
-            cmd, stdout=PIPE, stderr=STDOUT, bufsize=1, universal_newlines=True
-        )
-        # Stop opening pipe blocking writing to it
-        os.set_blocking(process.stdout.fileno(), False)
-        logger.info(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time())))
-
-        # Poll returns None while running, will return the error code when finished
-        while process.poll() is None:
-            log_buffer(process.stdout)
-            # Once buffer is clear sleep for 0.2s before trying again.
-
-            try:
-                # Note that using cancel.sleep allows the InvocationCancelledError to
-                # be thrown
-                cancel.sleep(0.2)
-            except lt.exceptions.InvocationCancelledError as e:
-                logger.info("Stitching cancelled by user")
-                process.kill()
-                raise e
-
-        # Print everything in the buffer when program finishes
-        log_buffer(process.stdout)
-
-        if process.poll() == 0:
-            logger.info("Stitching complete")
-        else:
-            raise ChildProcessError(f"Subprocess {cmd[0]} exited with an error.")
-
     @lt.thing_action
     def stitch_scan(
         self,
@@ -831,7 +737,7 @@ class SmartScanThing(lt.Thing):
         cancel: lt.deps.CancelHook,
         scan_name: str,
         stitch_resize: Optional[float] = None,
-        overlap: float = 0.0,
+        overlap: Optional[float] = None,
     ) -> None:
         """Generate a stitched image based on stage position metadata.
 
@@ -840,85 +746,24 @@ class SmartScanThing(lt.Thing):
         """
         # TODO tidy this function
         # TODO handle if json_path is None (easier once tidied)
-        json_fpath = self._scan_dir_manager.get_scan_data_path()
-
-        if self.stitch_tiff:
-            tiff_arg = "--stitch_tiff"
-        else:
-            tiff_arg = "--no-stitch_tiff"
-
-        if overlap == 0.0:
-            try:
-                with open(json_fpath, "r", encoding="utf-8") as data_file:
-                    data_loaded = json.load(data_file)
-                overlap = data_loaded["overlap"]
-            except (json.decoder.JSONDecodeError, FileNotFoundError, TypeError):
-                # As there is no schema or pydantic model this should handle
-                # the file not being there, it not being json in the file,
-                # or the imported data not being indexable
-                logger.warning(
-                    f"Couldn't read scan data, is {json_fpath} missing or corrupt? "
-                    "Attempting stitch with overlap value of 0.1"
-                )
-                overlap = 0.1
-            except KeyError:
-                logger.warning(
-                    "Value for overlap not found in scan data. "
-                    "Attempting stitch with overlap value of 0.1"
-                )
-                overlap = 0.1
-
-        if stitch_resize is None:
-            try:
-                with open(json_fpath, "r", encoding="utf-8") as data_file:
-                    data_loaded = json.load(data_file)
-
-                # Handle "capture resolution" being used to store the save resolution
-                # in old scans.
-                key = (
-                    "capture resolution"
-                    if "capture resolution" in data_loaded
-                    else "save_resolution"
-                )
-                save_resolution = data_loaded[key]
-                stitch_resize = STITCHING_RESOLUTION[0] / save_resolution[0]
-            except (json.decoder.JSONDecodeError, FileNotFoundError, TypeError):
-                # As there is no schema or pydantic model this should handle
-                # the file not being there, it not being json in the file,
-                # or the imported data not being indexable
-                logger.warning(
-                    f"Couldn't read scan data, is {json_fpath} missing or corrupt? "
-                    "Attempting stitch with resize value of 0.5"
-                )
-                stitch_resize = 0.5
-            except KeyError:
-                logger.warning(
-                    "Value for save resolution not found in scan data. "
-                    "Attempting stitch with resize value of 0.5"
-                )
-                stitch_resize = 0.5
-
+        scan_data_dict = self._scan_dir_manager.get_scan_data_dict(scan_name)
+        if scan_data_dict is None:
+            logger.warning("Couldn't read scan data it may be missing or corrupt.")
+        final_stitcher = stitching.FinalStitcher(
+            self._scan_dir_manager.img_dir_for(scan_name),
+            logger=logger,
+            overlap=overlap,
+            correlation_resize=stitch_resize,
+            stitch_tiff=self.stitch_tiff,
+            scan_data_dict=scan_data_dict,
+        )
         try:
-            self.run_subprocess(
-                logger=logger,
-                cancel=cancel,
-                cmd=[
-                    STITCHING_CMD,
-                    "--stitching_mode",
-                    "all",
-                    f"{tiff_arg}",
-                    "--stitch_dzi",
-                    "--minimum_overlap",
-                    f"{round(overlap * 0.9, 2)}",
-                    "--resize",
-                    f"{stitch_resize}",
-                    self._scan_dir_manager.img_dir_for(scan_name),
-                ],
-            )
+            final_stitcher.start(cancel)
         except lt.exceptions.InvocationCancelledError:
             # Sleep for 1 second just to allow invocation logs to pass to user.
             time.sleep(1)
-            pass
+        except SubprocessError as e:
+            self._scan_logger.error(f"Stitching failed: {e}", exc_info=e)
 
     @lt.thing_action
     def download_zip(