From 6ceb1a98458b35b53ac3f08f1865c6bd83130b41 Mon Sep 17 00:00:00 2001 From: Julian Stirling Date: Sun, 3 Aug 2025 14:33:45 +0100 Subject: [PATCH] Pull stitching out of SmartScanThing --- .../scan_directories.py | 22 +- .../stitching.py | 295 ++++++++++++++++++ .../things/smart_scan.py | 241 +++----------- 3 files changed, 358 insertions(+), 200 deletions(-) create mode 100644 src/openflexure_microscope_server/stitching.py diff --git a/src/openflexure_microscope_server/scan_directories.py b/src/openflexure_microscope_server/scan_directories.py index 8104c69c..6533c589 100644 --- a/src/openflexure_microscope_server/scan_directories.py +++ b/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]: diff --git a/src/openflexure_microscope_server/stitching.py b/src/openflexure_microscope_server/stitching.py new file mode 100644 index 00000000..27461e1e --- /dev/null +++ b/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) diff --git a/src/openflexure_microscope_server/things/smart_scan.py b/src/openflexure_microscope_server/things/smart_scan.py index 6b129ab0..cb1ef83d 100644 --- a/src/openflexure_microscope_server/things/smart_scan.py +++ b/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(