diff --git a/src/openflexure_microscope_server/scan_directories.py b/src/openflexure_microscope_server/scan_directories.py index 6f235191..de4539cc 100644 --- a/src/openflexure_microscope_server/scan_directories.py +++ b/src/openflexure_microscope_server/scan_directories.py @@ -1,15 +1,18 @@ """Functionality to manage file system operations for scan directories.""" -from typing import Optional +from typing import Optional, Any, Mapping import os import re import shutil import zipfile import threading +import json +from datetime import datetime, timedelta -from pydantic import BaseModel +from pydantic import BaseModel, field_validator, field_serializer from openflexure_microscope_server.utilities import requires_lock +from openflexure_microscope_server.utilities import make_name_safe IMG_DIR_NAME = "images" SCAN_ZERO_PAD_DIGITS = 4 @@ -17,6 +20,8 @@ SCAN_ZERO_PAD_DIGITS = 4 STITCH_REGEX = re.compile(r"stitched\.jpe?g$") IMAGE_REGEX = re.compile(r"-?[0-9]+_-?[0-9]+\.jpe?g$") +SCAN_DATA_FILENAME = "scan_data.json" + class NotEnoughFreeSpaceError(IOError): """An exception raised if there is not enough free space on disk to scan.""" @@ -33,6 +38,130 @@ class ScanInfo(BaseModel): dzi: Optional[str] +class ScanData(BaseModel): + """Data about a scan to be saved to a JSON file in the directory. + + This serialises into a human readable format where possible with + + timestamps in %Y-%m-%d_%H:%M:%S format + timedeltas in %H:%M:%S format + + Properties that are not known until the end have ``None`` serialised as "Unknown" + """ + + model_config = {"extra": "forbid"} + + scan_name: str + """The name of the scan i.e. scan_0001""" + + starting_position: Mapping[str, int] + """The starting position in dictionary format.""" + + overlap: float + """The overlap between adjacent images as a fraction of the image size.""" + + max_dist: int + """The maximum distance the scan could move (in steps) from the starting position.""" + + dx: int + """The number of steps between adjacent images in x.""" + + dy: int + """The number of steps between adjacent images in y.""" + + autofocus_dz: int + """The z range used for autofocus (in steps).""" + + autofocus_on: bool + """Whether autofocus is on.""" + + start_time: datetime + """The time the scan started.""" + + skip_background: bool + """Whether automatic background detection is on, skipping locations with no sample.""" + + stitch_automatically: bool + """Whether the scan is set to automatically stitch when complete.""" + + correlation_resize: float + """The resize factor applied to images when the stitching program is correlating.""" + + save_resolution: tuple[int, int] + """The resolution that scan images are saved at.""" + + image_count: int = 0 + """The number of images taken.""" + + duration: Optional[timedelta] = None + """The duration of the scan. + + This is automatically set when ``set_final_data()`` is run. + """ + + scan_result: Optional[str] = None + """The result of the scan. + + This should be set with ``set_final_data()`` to ensure duration is set. + """ + + def set_final_data(self, result: str): + """Set the final data for the scan, scan duration is automatically calculated. + + :param result: A string describing the result. + """ + self.duration = datetime.now() - self.start_time + self.scan_result = result + + @field_validator("start_time", mode="before") + @classmethod + def parse_timestamp(cls, value: str | datetime) -> datetime: + """Validate a timestamp that may be a string in Year-Month-Day_Hrs:Min:Sec format.""" + if isinstance(value, str): + return datetime.strptime(value, "%Y-%m-%d_%H:%M:%S") + return value + + @field_serializer("start_time") + def serialize_timestamp(self, value: datetime) -> str: + """Serialise timestamp to Year-Month-Day_Hrs:Min:Sec format.""" + return value.strftime("%Y-%m-%d_%H:%M:%S") + + @field_validator("duration", mode="before") + @classmethod + def parse_timedelta(cls, value: Optional[str | timedelta]) -> Optional[timedelta]: + """Validate a timedelta that may be a string in Hrs:Min:Sec format or "Unknown".""" + if isinstance(value, str): + if value == "Unknown": + return None + hrs, mins, secs = map(int, value.split(":")) + return timedelta(hours=hrs, minutes=mins, seconds=secs) + return value + + @field_serializer("duration") + def serialize_timedelta(self, value: Optional[timedelta]) -> str: + """Serialise timedelta to Hrs:Min:Sec (or "Unknown" if None).""" + if value is None: + return "Unknown" + # Calculate the string manually rather than use str(), as this may convert to + # days for a very very long scan, and then it is much harder to parse. + total_secs = value.total_seconds() + hrs = int(total_secs // 3600) + mins = int((total_secs % 3600) // 60) + secs = int((total_secs % 60)) + return f"{hrs}:{mins:02}:{secs:02}" + + @field_validator("scan_result", mode="before") + @classmethod + def parse_unknown_as_none(cls, value: Optional[str | int]) -> Optional[str | int]: + """Validate the string "Unknown" as None.""" + return None if value == "Unknown" else value + + @field_serializer("scan_result") + def serialize_none_as_unknown(self, value: Optional[str | int]) -> str | int: + """Serialise None as "Unknown" for a more human readable result.""" + return "Unknown" if value is None else value + + class ScanDirectoryManager: """A class for managing interactions with scan directories.""" @@ -111,6 +240,34 @@ class ScanDirectoryManager: return None return self.get_file_path_from_img_dir(scan_name, stitch_fname) + @requires_lock + def get_scan_data_path(self, scan_name: str) -> Optional[str]: + """Return the file full path scan data JSON file. + + If no scan data JSON file is found, return None + """ + scan_data_path = ScanDirectory(scan_name, self.base_dir).scan_data_path + if scan_data_path is None: + return None + if not os.path.isfile(scan_data_path): + 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 model 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]: @@ -131,6 +288,7 @@ class ScanDirectoryManager: For more explanation on the scan naming see `new_scan_dir` """ + scan_name = make_name_safe(scan_name) # A regex with the scan name and a group for the numbers scan_regex = re.compile( "^" + scan_name + "_([0-9]{" + str(SCAN_ZERO_PAD_DIGITS) + "})$" @@ -250,6 +408,16 @@ class ScanDirectory: return im_path return None + @property + def scan_data_path(self) -> Optional[str]: + """The path to the scan data json file for this directory. + + Returns None if there is no images dir to write to. + """ + if self.images_dir is not None: + return os.path.join(self.images_dir, SCAN_DATA_FILENAME) + return None + @property def created_time(self) -> float: """The time the directory was created on disk.""" @@ -341,6 +509,16 @@ class ScanDirectory: files.append(os.path.relpath(full_path, self.dir_path)) return files + def save_scan_data(self, scan_data: ScanData): + """Save the scan data for this scan to disk.""" + if self.scan_data_path is None: + raise FileNotFoundError( + "There is no images directory to save scan data into." + ) + + with open(self.scan_data_path, "w", encoding="utf-8") as f: + f.write(scan_data.model_dump_json(indent=4)) + def zip_files(self, final_version: bool = False) -> str: """Zips any images from the scan not yet zipped, return full path to zip. diff --git a/src/openflexure_microscope_server/stitching.py b/src/openflexure_microscope_server/stitching.py new file mode 100644 index 00000000..b88c4036 --- /dev/null +++ b/src/openflexure_microscope_server/stitching.py @@ -0,0 +1,322 @@ +"""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 due to the +CPU intensity of stitching causing scanning problems due to the Python Global +Interpreter Lock (GIL). May be possible to shift to multiprocessing in the future. +""" + +from typing import Optional, Any +import threading +import subprocess +import os +from io import TextIOWrapper +import shlex + +import labthings_fastapi as lt + +from openflexure_microscope_server.utilities import make_path_safe + +STITCHING_CMD = "openflexure-stitch" +STITCHING_RESOLUTION = (820, 616) + +DEFAULT_OVERLAP = 0.1 +DEFAULT_RESIZE = 0.5 + + +class StitcherValidationError(RuntimeError): + """The stitcher received values that it deems unsafe to create a command from.""" + + +def validate_command(cmd: list[str]): + """Validate that the command only characters that are allowed in a path. + + The values in the commands should be numbers, commandline flags, paths, and + executables. All of these should be allowed by ``make_path_safe``. + + :raises StitcherValidationError: if any element in the command is not safe. + """ + for element in cmd: + if element != make_path_safe(element): + raise StitcherValidationError( + "Invalid stiching command: Contains unsafe characters." + ) + + +class BaseStitcher: + """A base stitching class for all stitchers. Don't initialise this directly. + + The base class has no way to run the command. Child classes should either implement + ``start``, ``running``, and ``wait`` methods if return after starting the + subprocess and can be polled or waited on like a thread; or ``run`` if the + the function blocks while the stitching subprocess is ongoing and return once + complete. + """ + + 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 + try: + overlap = float(overlap) + correlation_resize = float(correlation_resize) + except ValueError as e: + raise StitcherValidationError( + "Stitching inputs overlap or correlation_resize were not floats as " + "expected." + ) from e + self.validate_path() + + self.min_overlap = round(overlap * 0.9, 2) + self.correlation_resize = float(correlation_resize) + self._mode = "all" + self._extra_args = [] + + @property + def command(self) -> list[str]: + """The command to run with subprocess.Popen.""" + # Revalidate for good measure, + self.validate_path() + + # The command, and the mode + initial_args = shlex.split(STITCHING_CMD) + ["--stitching_mode", self._mode] + # Use float() just to really ensure that anything input is a float + setting_args = [ + "--minimum_overlap", + f"{float(self.min_overlap)}", + "--resize", + f"{float(self.correlation_resize)}", + ] + full_cmd = initial_args + self._extra_args + setting_args + [self.images_dir] + validate_command(full_cmd) + return full_cmd + + def validate_path(self) -> None: + """Check path is safe before making a command to run with subprocess. + + This is essential for stopping arbitrary code execution. + + :raises RuntimeError: if inputs are unsafe. + """ + if self.images_dir != make_path_safe(self.images_dir): + raise StitcherValidationError( + "Invalid directory path: Contains unsafe characters." + ) + + +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" + + 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) -> None: + """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 stitch 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 run( + 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) + + self._log_ongoing(process, cancel=cancel) + + if process.poll() == 0: + self.logger.info("Stitching complete") + else: + raise ChildProcessError(f"Subprocess {STITCHING_CMD} 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: TextIOWrapper): + """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 caad1acc..c4206bcc 100644 --- a/src/openflexure_microscope_server/things/smart_scan.py +++ b/src/openflexure_microscope_server/things/smart_scan.py @@ -7,13 +7,12 @@ It also controls external processes for live stitching composite images, and the creation of the final stitched images. """ -from typing import Optional, Mapping +from typing import Optional import threading import os import time -import json from datetime import datetime -from subprocess import CompletedProcess, Popen, PIPE, SubprocessError, STDOUT +from subprocess import SubprocessError from fastapi import HTTPException from fastapi.responses import FileResponse @@ -22,9 +21,9 @@ from PIL import Image 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 +39,6 @@ AutofocusDep = lt.deps.direct_thing_client_dependency(AutofocusThing, "/autofocu JPEGBlob = lt.blob.blob_type("image/jpeg") ZipBlob = lt.blob.blob_type("application/zip") -SCAN_DATA_FILENAME = "scan_data.json" -STITCHING_CMD = "openflexure-stitch" - -STITCHING_RESOLUTION = (820, 616) - class ScanNotRunningError(RuntimeError): """Exception called when scan not running that requires a scan to be running.""" @@ -86,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 @@ -97,21 +89,18 @@ class SmartScanThing(lt.Thing): # when the `sample_scan` lt.thing_action is called. It is saved as # private class variable along with many others here. # Access to these variables requires a scan to be running, - # any method that calls these should be decorrected with + # any method that calls these should be decorated with # @_scan_running self._scan_logger: Optional[lt.deps.InvocationLogger] = None self._cancel: Optional[lt.deps.CancelHook] = None self._autofocus: Optional[AutofocusDep] = None self._stage: Optional[StageDep] = None self._cam: Optional[CamDep] = None - self._metadata_getter: Optional[lt.deps.GetThingStates] = None self._csm: Optional[CSMDep] = None + self._ongoing_scan: Optional[scan_directories.ScanDirectory] = None - 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[dict] = None + self._scan_data: Optional[scan_directories.ScanData] = None + self._preview_stitcher: Optional[stitching.PreviewStitcher] = None @lt.thing_action def sample_scan( @@ -121,7 +110,6 @@ class SmartScanThing(lt.Thing): autofocus: AutofocusDep, stage: StageDep, cam: CamDep, - metadata_getter: lt.deps.GetThingStates, csm: CSMDep, scan_name: str = "", ): @@ -141,22 +129,16 @@ class SmartScanThing(lt.Thing): self._autofocus = autofocus self._stage = stage self._cam = cam - self._metadata_getter = metadata_getter self._csm = csm - self._capture_thread = None - self._scan_images_taken = 0 - - # Set _scan_data to None. This is needed just in case an exception is raised - # before _run_scan (which sets the real data). As we check this in the `except` + # `scan_data` should already be None. This is added as a precaution as + # the presence of `scan_data` is used during error handling to + # determine whether the scan started. self._scan_data = None - try: self._check_background_and_csm_set() self._ongoing_scan = self._scan_dir_manager.new_scan_dir(scan_name) self._latest_scan_name = self._ongoing_scan.name self._autofocus.looping_autofocus(dz=self.autofocus_dz, start="centre") - # record starting position so we can return there - self._starting_position = self._stage.position self._run_scan() except Exception as e: # If _scan_data is set then scan started @@ -164,6 +146,9 @@ class SmartScanThing(lt.Thing): self._return_to_starting_position() if not isinstance(e, scan_directories.NotEnoughFreeSpaceError): # Don't stitch if drive is full (already logged) + self._scan_logger.info( + "Attempting to stitch and archive the images acquired so far." + ) self._perform_final_stitch() # Error must be raised so UI gives correct output raise e @@ -174,13 +159,12 @@ class SmartScanThing(lt.Thing): self._autofocus = None self._stage = None self._cam = None - self._metadata_getter = None self._csm = None - self._capture_thread = None self._ongoing_scan = None - 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): @@ -280,17 +264,16 @@ class SmartScanThing(lt.Thing): return dx, dy @_scan_running - def _set_scan_data(self): - """Set date for this scan to the ``self._scan_data`` variable. - - This needs to become a dataclass at some point. - """ + def _collect_scan_data(self) -> scan_directories.ScanData: + """Collect and return the data for this scan so it cannot be changed mid-scan.""" + # Record starting position so it can be returned to at end of scan. + starting_position = self._stage.position overlap = self.overlap dx, dy = self._calc_displacement_from_test_image(overlap) - stitch_resize = STITCHING_RESOLUTION[0] / self.save_resolution[0] + correlation_resize = stitching.STITCHING_RESOLUTION[0] / self.save_resolution[0] self._scan_logger.debug( - f"Resizing images when stitching by a factor of {stitch_resize}" + f"Resizing images when correlating by a factor of {correlation_resize}" ) self._scan_logger.info( @@ -308,87 +291,40 @@ class SmartScanThing(lt.Thing): autofocus_dz = 0 # Fix scan parameters in case UI is updated during scan. - self._scan_data = { - "scan_name": self._ongoing_scan.name, - "overlap": overlap, - "max_dist": self.max_range, - "dx": dx, - "dy": dy, - "autofocus_dz": autofocus_dz, - "autofocus_on": bool(autofocus_dz), - "start_time": time.strftime("%H_%M_%S-%d_%m_%Y"), - "skip_background": self.skip_background, - "stitch_automatically": self.stitch_automatically, - "stitch_resize": stitch_resize, - "save_resolution": self.save_resolution, - } - - @_scan_running - def _save_scan_inputs_json(self): - """Save scan inputs as a JSON file in the scan folder. - - This file allows the user to review the settings used in the scan. - """ - # Should this be a method of the scan_data dataclass? - - data = { - "scan_name": self._ongoing_scan.name, - "overlap": self._scan_data["overlap"], - "autofocus range": self._scan_data["autofocus_dz"], - "dx": self._scan_data["dx"], - "dy": self._scan_data["dy"], - "start time": self._scan_data["start_time"], - "skipping background": self._scan_data["skip_background"], - "capture resolution": self._scan_data["save_resolution"], - } - - scan_inputs_fname = os.path.join( - self._ongoing_scan.images_dir, SCAN_DATA_FILENAME + return scan_directories.ScanData( + scan_name=self._ongoing_scan.name, + starting_position=starting_position, + overlap=overlap, + max_dist=self.max_range, + dx=dx, + dy=dy, + autofocus_dz=autofocus_dz, + autofocus_on=bool(autofocus_dz), + start_time=datetime.now(), + skip_background=self.skip_background, + stitch_automatically=self.stitch_automatically, + correlation_resize=correlation_resize, + save_resolution=self.save_resolution, ) - with open(scan_inputs_fname, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=4) @_scan_running - def _update_scan_data_json(self, scan_result: str): + def _save_final_scan_data(self, scan_result: str): """Update scan data JSON file with data only known at the end of the scan. Takes scan_result, a string that is either "success", "cancelled by user", or the error that ended the scan. """ - # Should this be a method of the scan_data dataclass? - current_time = datetime.now().replace(microsecond=0) - start_time = datetime.strptime( - self._scan_data["start_time"], "%H_%M_%S-%d_%m_%Y" - ).replace(microsecond=0) - - duration = current_time - start_time - - outputs = { - "image_count": self._scan_images_taken, - "duration": str(duration), - "scan_result": scan_result, - } - - scan_data_fname = os.path.join( - self._ongoing_scan.images_dir, SCAN_DATA_FILENAME - ) - - with open(scan_data_fname, encoding="utf-8") as f: - data = json.load(f) - - data.update(outputs) - - with open(scan_data_fname, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=4) + self._scan_data.set_final_data(result=scan_result) + self._ongoing_scan.save_scan_data(self._scan_data) @_scan_running def _manage_stitching_threads(self): """Manage the stitching threads, starting them if needed and not already running.""" # 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 self._scan_data.image_count > 3: + if not self._preview_stitcher.running: + self._preview_stitcher.start() @_scan_running def _run_scan(self): @@ -398,62 +334,41 @@ class SmartScanThing(lt.Thing): scan should be stitched and whether the microscope should return to the starting x,y,z position. """ - # Used to check if finally was reached via exception (except - # cancel by user) - scan_successful = True - try: self._cam.start_streaming(main_resolution=(3280, 2464)) - self._set_scan_data() - self._save_scan_inputs_json() - if self._scan_images_taken != 0: - msg = "_scan_images_taken should be zero before starting scanning" - raise RuntimeError(msg) + 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.correlation_resize, + ) # This is the main loop of the scan! self._main_scan_loop() - self._update_scan_data_json(scan_result="success") + self._save_final_scan_data(scan_result="success") except lt.exceptions.InvocationCancelledError: - scan_successful = False # Reset the cancel event so it can be thrown again self._cancel.clear() self._scan_logger.info("Stopping scan because it was cancelled.") - self._update_scan_data_json(scan_result="cancelled by user") - except scan_directories.NotEnoughFreeSpaceError as e: - scan_successful = False - self._update_scan_data_json(scan_result=str(e)) - self._scan_logger.error( - f"Stopping scan to avoid filling up the disk: {e}", - exc_info=e, - ) - raise e + self._save_final_scan_data(scan_result="cancelled by user") except Exception as e: - scan_successful = False + err_name = type(e).__name__ + if self._scan_data is not None: + self._save_final_scan_data(scan_result=f"{err_name}: {e}") self._scan_logger.error( - f"The scan stopped because of an error: {e} " - "Attempting to stitch and archive the images acquired so far.", + f"The scan stopped because of an error: {e}", exc_info=e, ) raise e finally: + # Don't set Preview Stitcher to None yet. It is used by + # _perform_final_stitch, which may also be run after this function completes + # if it ended due to an exception. + # Start streaming in the default resolution again as soon as possible self._cam.start_streaming() - if self._capture_thread: - # If the capture thread had an error, we capture it here - try: - self._capture_thread.join() - except Exception as e: - # If the scan has already ended due to an exception, - # ignore any exceptions. If it appeared to be successful, - # log the error. - if scan_successful: - self._scan_logger.error( - "The scan appears to have completed successfully, however " - f"the final capture raised the following error: {e}." - "Attempting to stitch and archive images.", - exc_info=e, - ) # This is what happens if the scan completes successfully or the # user cancels it. @@ -475,9 +390,9 @@ class SmartScanThing(lt.Thing): # have multiple starting positions, each of which will be visited before the # scan can end. planner_settings = { - "dx": self._scan_data["dx"], - "dy": self._scan_data["dy"], - "max_dist": self._scan_data["max_dist"], + "dx": self._scan_data.dx, + "dy": self._scan_data.dy, + "max_dist": self._scan_data.max_dist, } route_planner = scan_planners.SmartSpiral( initial_position=(self._stage.position["x"], self._stage.position["y"]), @@ -502,7 +417,7 @@ class SmartScanThing(lt.Thing): capture_image = True # If skipping background, take an image to check if current field of view is background - if self._scan_data["skip_background"]: + if self._scan_data.skip_background: capture_image, bg_message = self._cam.image_is_sample() if not capture_image: @@ -515,8 +430,8 @@ class SmartScanThing(lt.Thing): focused, focused_height = self._autofocus.run_smart_stack( images_dir=self._ongoing_scan.images_dir, - autofocus_dz=self._scan_data["autofocus_dz"], - save_resolution=self._scan_data["save_resolution"], + autofocus_dz=self._scan_data.autofocus_dz, + save_resolution=self._scan_data.save_resolution, ) current_pos_xyz = (new_pos_xyz[0], new_pos_xyz[1], focused_height) @@ -526,7 +441,7 @@ class SmartScanThing(lt.Thing): ) # increment capture counter as thread has completed - self._scan_images_taken += 1 + self._scan_data.image_count += 1 # Add it to the incremental zip self._ongoing_scan.zip_files() @@ -534,15 +449,15 @@ class SmartScanThing(lt.Thing): def _return_to_starting_position(self): """Return to the initial scan position, if set.""" self._scan_logger.info("Returning to starting position.") - if self._starting_position is not None: + if self._scan_data is not None: self._stage.move_absolute( - **self._starting_position, block_cancellation=True + **self._scan_data.starting_position, block_cancellation=True ) @_scan_running def _perform_final_stitch(self): """Update the scan zip and perform final stitch of the data.""" - if self._scan_images_taken <= 3: + if self._scan_data.image_count <= 3: self._scan_logger.info("Not performing a stitch as 3 or fewer images taken") return @@ -550,19 +465,18 @@ 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) + if self._preview_stitcher is not None: + 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, + correlation_resize=self._scan_data.correlation_resize, + overlap=self._scan_data.overlap, + ) @lt.fastapi_endpoint( "get", @@ -774,190 +688,39 @@ 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, logger: lt.deps.InvocationLogger, cancel: lt.deps.CancelHook, scan_name: str, - stitch_resize: Optional[float] = None, - overlap: float = 0.0, + correlation_resize: Optional[float] = None, + overlap: Optional[float] = None, ) -> None: """Generate a stitched image based on stage position metadata. Note that as this is a lt.thing_action it needs the logger passed as a variable if called from another thing action """ - json_fpath = self._scan_dir_manager.get_file_path_from_img_dir( - scan_name=scan_name, filename=SCAN_DATA_FILENAME + 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=correlation_resize, + stitch_tiff=self.stitch_tiff, + scan_data_dict=scan_data_dict, ) - - 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 {SCAN_DATA_FILENAME} 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) - save_resolution = data_loaded["capture resolution"] - 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 {SCAN_DATA_FILENAME} missing or corrupt? " - "Attempting stitch with resize value of 0.5" - ) - stitch_resize = 0.5 - except KeyError: - logger.warning( - "Value for capture resolution not found in scan data. " - "Attempting stitch with resize value of 0.5" - ) - stitch_resize = 0.5 - 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), - ], - ) + # start the final stitch, providing the cancel hook to allow aborting + final_stitcher.run(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( diff --git a/src/openflexure_microscope_server/utilities.py b/src/openflexure_microscope_server/utilities.py index 21099216..bbad56f9 100644 --- a/src/openflexure_microscope_server/utilities.py +++ b/src/openflexure_microscope_server/utilities.py @@ -3,6 +3,7 @@ from typing import TypeVar, Callable, ParamSpec import os import re +import sys from threading import Thread import logging from importlib.metadata import version @@ -126,6 +127,50 @@ def _wrap_and_catch_errors(target, error_buffer, *args, **kwargs): error_buffer.append(e) +# Compiled regular expressions for unsafe characters +# Matches anything that isn't a-z, A-Z, 0-9, _, ., -, :, /, \ +_WINDOWS_UNSAFE_PATTERN = re.compile(r"[^a-zA-Z0-9_.\-:/\\]") +# Matches anything that isn't a-z, A-Z, 0-9, _, ., -, \ +_POSIX_UNSAFE_PATTERN = re.compile(r"[^a-zA-Z0-9_.\-/]") +# Matches anything that isn't a-z, A-Z, 0-9, _, ., - +_NAME_UNSAFE_PATTERN = re.compile(r"[^a-zA-Z0-9_.\-]") + + +def make_path_safe(unsafe_path_string: str) -> str: + """Replace unsafe characters in a file path with underscores, preserving separators. + + This can be used to check that user inputs have not been concatenated into an + unsafe filepath. + + This ensures compatibility across platforms by preserving only valid characters, + including path separators (forward slash on POSIX, and forward + slash/backslash/colon on Windows). + + :param unsafe_path_string: The original path string to sanitise. + + :returns: A version of the input string safe to use as a file path. + """ + unsafe_character_pattern = ( + _WINDOWS_UNSAFE_PATTERN + if sys.platform.startswith("win") + else _POSIX_UNSAFE_PATTERN + ) + return unsafe_character_pattern.sub("_", unsafe_path_string) + + +def make_name_safe(unsafe_name_string: str) -> str: + """Replace unsafe characters in a filename or identifier with underscores. + + This excludes all path separators, ensuring the result is safe to use as a + standalone filename or identifier component. + + :param unsafe_name_string: The original name string to sanitise. + + :returns: A version of the input string safe to use as a file name or identifier. + """ + return _NAME_UNSAFE_PATTERN.sub("_", unsafe_name_string) + + class VersionData(BaseModel): """A BaseModel containing the information about the server version. diff --git a/tests/mock_stitching/mock-stitch.py b/tests/mock_stitching/mock-stitch.py new file mode 100644 index 00000000..b4c9d16a --- /dev/null +++ b/tests/mock_stitching/mock-stitch.py @@ -0,0 +1,19 @@ +"""CLI script used for testing subprocess calls that should go to openflexure-stitch.""" + +import sys +import time + + +def main(): + """Echo command-line arguments with a short delay between each.""" + input_arguments = sys.argv[1:] + for arg in input_arguments: + # This is used to check we catch errors correctly. + if arg == "ERROR": + raise RuntimeError("I was told to do this.") + print(arg, flush=True) + time.sleep(0.2) + + +if __name__ == "__main__": + main() diff --git a/tests/mock_things/mock_camera.py b/tests/mock_things/mock_camera.py index 81104dde..a817a47d 100644 --- a/tests/mock_things/mock_camera.py +++ b/tests/mock_things/mock_camera.py @@ -11,18 +11,25 @@ from openflexure_microscope_server.background_detect import ( BackgroundDetectorStatus, ColourChannelDetectSettings, ) +from unittest.mock import Mock, PropertyMock -class MockCameraThing: +class MockCameraThing(Mock): """A mock camera Thing that imports no code from ``BaseCamera``. The class needs functionality added to it over time as more complex - mocking is needed. It imports no code from ``BaseCamera or any other + mocking is needed. It imports no code from ``BaseCamera`` or any other camera Thing, so that coverage is not artificially inflated. """ - background_detector_status = BackgroundDetectorStatus( - ready=True, - settings=ColourChannelDetectSettings().model_dump(), - settings_schema={"fake": "schema"}, - ) + def __init__(self, *args, **kwargs): + """Initialise the mock camera, args and kwargs are the mock args and kwargs.""" + super().__init__(*args, **kwargs) + + self.background_detector_status = PropertyMock( + return_value=BackgroundDetectorStatus( + ready=True, + settings=ColourChannelDetectSettings().model_dump(), + settings_schema={"fake": "schema"}, + ) + ) diff --git a/tests/test_scan_data.py b/tests/test_scan_data.py new file mode 100644 index 00000000..a1ea4f1c --- /dev/null +++ b/tests/test_scan_data.py @@ -0,0 +1,125 @@ +"""Test the functionality in the scan_directories module.""" + +from datetime import datetime, timedelta +from copy import copy +import json +from math import floor + +from openflexure_microscope_server.scan_directories import ScanData + +MOCK_START_TIME = datetime( + year=2024, + month=12, + day=25, + hour=11, + minute=0, + second=0, + microsecond=1234, +) + +MOCK_END_TIME = datetime( + year=2024, + month=12, + day=25, + hour=12, + minute=29, + second=3, + microsecond=5678, +) + + +def _fake_scan_data(**kwargs) -> ScanData: + """Make fake scan data, the start time is now. Final properties are not added. + + :param **kwargs: Key word arguments can be used to override other values. + """ + data_dict = { + "scan_name": "fake_scan_0001", + "starting_position": {"x": 123, "y": 456, "z": 789}, + "overlap": 0.1, + "max_dist": 100000, + "dx": 100, + "dy": 100, + "autofocus_dz": 1000, + "autofocus_on": True, + "start_time": copy(MOCK_START_TIME), + "skip_background": True, + "stitch_automatically": True, + "correlation_resize": 0.25, + "save_resolution": (1000, 1000), + } + for key, value in kwargs.items(): + data_dict[key] = value + return ScanData(**data_dict) + + +def test_set_final_data(): + """Check that adding final data to a ScanData object works as expected.""" + scan_data = _fake_scan_data() + + assert scan_data.image_count == 0 + assert scan_data.duration is None + assert scan_data.scan_result is None + + # Quickly take 123 images! + scan_data.image_count += 123 + # Should set duration based on finishing at datetime.now() + scan_data.set_final_data(result="Success") + expected_duration = datetime.now() - MOCK_START_TIME + expected_duration_s = expected_duration.total_seconds() + scan_duration_s = scan_data.duration.total_seconds() + + assert scan_data.image_count == 123 + assert expected_duration_s - 1 < scan_duration_s < expected_duration_s + 1 + assert scan_data.scan_result == "Success" + + +def test_custom_serialisation(): + """Check that the custom serialisation in ScanData works as expected.""" + scan_data = _fake_scan_data() + # Serialise to string then load directly as json + scan_data_dict = json.loads(scan_data.model_dump_json()) + assert scan_data_dict["start_time"] == "2024-12-25_11:00:00" + assert scan_data_dict["image_count"] == 0 + assert scan_data_dict["duration"] == "Unknown" + assert scan_data_dict["scan_result"] == "Unknown" + + scan_data.image_count += 123 + scan_data.set_final_data(result="Success") + # Can't mock datetime.now as datetime is immutable. So just replace duration + scan_data.duration = MOCK_END_TIME - scan_data.start_time + scan_data_dict = json.loads(scan_data.model_dump_json()) + assert scan_data_dict["image_count"] == 123 + assert scan_data_dict["duration"] == "1:29:03" + assert scan_data_dict["scan_result"] == "Success" + + +def test_round_trip_not_finalised(): + """Check that ScanData without final data can be serialised and deserialised.""" + scan_data = _fake_scan_data() + scan_data_dict = json.loads(scan_data.model_dump_json()) + scan_data_reloaded = ScanData(**scan_data_dict) + + # For the round trip to be equal we must remove microseconds from the start + # time as they are not saved + scan_data.start_time = scan_data.start_time.replace(microsecond=0) + assert scan_data == scan_data_reloaded + + +def test_round_trip_finalised(): + """Check that finalised ScanData can be serialised and deserialised.""" + scan_data = _fake_scan_data() + # Finalise the data. + scan_data.image_count += 123 + scan_data.set_final_data(result="Success") + + scan_data_dict = json.loads(scan_data.model_dump_json()) + scan_data_reloaded = ScanData(**scan_data_dict) + + # For the round trip to be equal we must remove microseconds from the start + # time and duration as they are not saved + scan_data.model_copy(update={}) + scan_data.start_time = scan_data.start_time.replace(microsecond=0) + scan_data.duration = timedelta(seconds=floor(scan_data.duration.total_seconds())) + + assert scan_data == scan_data_reloaded diff --git a/tests/test_scan_directories.py b/tests/test_scan_directories.py index 1f6c0609..f589c8a7 100644 --- a/tests/test_scan_directories.py +++ b/tests/test_scan_directories.py @@ -7,6 +7,7 @@ import shutil import logging import random import time +import json from collections import namedtuple import pytest @@ -17,8 +18,10 @@ from openflexure_microscope_server.scan_directories import ( ScanInfo, get_files_in_zip, NotEnoughFreeSpaceError, + SCAN_DATA_FILENAME, ) +from .test_scan_data import _fake_scan_data from .utilities import assert_unique_of_length # A global logger to pass in as an Invocation Logger @@ -157,6 +160,16 @@ def test_basic_directory_operations(): assert not os.path.isdir(scan_path) +def test_bad_scan_names(): + """Check scan names with spaces, or worse BASH commands are not allowed.""" + _clear_scan_dir() + scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR) + scan_dir = scan_dir_manager.new_scan_dir("fake scan") + assert scan_dir.name == "fake_scan_0001" + scan_dir = scan_dir_manager.new_scan_dir("fake scan;rm -rf /;") + assert scan_dir.name == "fake_scan_rm_-rf____0001" + + def test_scan_sequence_and_listing(): """Check created scans are added in order and listed correctly.""" _clear_scan_dir() @@ -281,6 +294,49 @@ def test_get_final_stitch(): assert scan_dir.get_final_stitch_name() == fake_scan_name +def test_get_scan_data_path(): + """Check that a scan data path behaves as expected.""" + _clear_scan_dir() + scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR) + scan_dir = scan_dir_manager.new_scan_dir("fake_scan") + scan_name = scan_dir.name + expected_path = os.path.join(scan_dir.images_dir, SCAN_DATA_FILENAME) + # Asking the scan manager for the path will return None as there is no file. + assert scan_dir_manager.get_scan_data_path(scan_name) is None + # Asking the scan directly will return the location the file should be located + # this is so it can be created. + assert scan_dir.scan_data_path == expected_path + + # Remove the images directory. + shutil.rmtree(scan_dir.images_dir) + # When the images directory is deleted, internally `scan_dir.images_dir` + # will return `None`. Check that `get_scan_data_path` copes with the `None` + # return from `scan_dir.images_dir`. + assert scan_dir_manager.get_scan_data_path(scan_name) is None + + +def test_get_scan_data_dict(): + """Check that the dictionary for the scan data is returned, or None if doesn't exist.""" + _clear_scan_dir() + scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR) + scan_dir = scan_dir_manager.new_scan_dir("fake_scan") + scan_name = scan_dir.name + # Doesn't yet exist + assert scan_dir_manager.get_scan_data_dict(scan_name) is None + + fake_data = {"foo": 1, "bar": "foobar"} + with open(scan_dir.scan_data_path, "w", encoding="utf-8") as json_file: + json.dump(fake_data, json_file) + + # Should now be able to load this fake data from disk + assert scan_dir_manager.get_scan_data_dict(scan_name) == fake_data + + # Check None is returned if the data cannot be read. + with open(scan_dir.scan_data_path, "w", encoding="utf-8") as json_file: + json_file.write("this is not json") + assert scan_dir_manager.get_scan_data_dict(scan_name) is None + + def test_empty_scan_info(): """Test the scan info is correct even if the scan is empty.""" _clear_scan_dir() @@ -339,6 +395,42 @@ def test_zipping_scan_data(): assert not file.endswith(".dzi") +def test_saving_and_loading_scan_data(): + """Test that scan data is saved and loaded as expected.""" + _clear_scan_dir() + scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR) + scan_dir = scan_dir_manager.new_scan_dir("fake_scan") + scan_name = scan_dir.name + + # Should start without a scan data file. + assert not os.path.isfile(scan_dir.scan_data_path) + # Create + scan_data_obj = _fake_scan_data() + scan_dir.save_scan_data(scan_data_obj) + # File should now exist + assert os.path.isfile(scan_dir.scan_data_path) + + # Dump the scan json to a string an reload it + # Note that more detailed checking of the dumping and loading of ScanData is in + # tests/test_scan_data.py + scan_data_dict = json.loads(scan_data_obj.model_dump_json()) + # What is loaded from file should be the same as from dumping and loading. + assert scan_dir_manager.get_scan_data_dict(scan_name) == scan_data_dict + + +def test_saving_scan_data_error(): + """Test that saving scan data if there is no images directory raises FileNotFoundError.""" + _clear_scan_dir() + scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR) + scan_dir = scan_dir_manager.new_scan_dir("fake_scan") + + # Remove the images directory. + shutil.rmtree(scan_dir.images_dir) + # Should raise FileNotFoundError. + with pytest.raises(FileNotFoundError): + scan_dir.save_scan_data(_fake_scan_data()) + + def test_all_files(): """Test all_files returns the path, and respects skipped directories.""" _clear_scan_dir() diff --git a/tests/test_smart_scan.py b/tests/test_smart_scan.py index 1a0be750..ab619e8a 100644 --- a/tests/test_smart_scan.py +++ b/tests/test_smart_scan.py @@ -18,14 +18,21 @@ import tempfile import os import shutil import logging +from datetime import datetime from fastapi import HTTPException import pytest +from labthings_fastapi.exceptions import InvocationCancelledError + from openflexure_microscope_server.things.smart_scan import ( SmartScanThing, ScanNotRunningError, ) +from openflexure_microscope_server.scan_directories import ( + ScanData, + NotEnoughFreeSpaceError, +) from .mock_things.mock_csm import MockCSMThing from .mock_things.mock_autofocus import MockAutoFocusThing @@ -170,7 +177,6 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None): af_mock = MockAutoFocusThing() stage_mock = MockStageThing() cam_mock = MockCameraThing() - meta_mock = 5 # not called csm_mock = MockCSMThing() class MockedSmartScanThing(SmartScanThing): @@ -189,10 +195,7 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None): assert self._autofocus is af_mock assert self._stage is stage_mock assert self._cam is cam_mock - assert self._metadata_getter is meta_mock assert self._csm is csm_mock - assert self._capture_thread is None - assert self._scan_images_taken == 0 # mock smart scan thing mock_ss_thing = MockedSmartScanThing(SCAN_DIR) @@ -208,7 +211,6 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None): autofocus=af_mock, stage=stage_mock, cam=cam_mock, - metadata_getter=meta_mock, csm=csm_mock, scan_name="FooBar", ) @@ -222,10 +224,7 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None): assert mock_ss_thing._autofocus is None assert mock_ss_thing._stage is None assert mock_ss_thing._cam is None - assert mock_ss_thing._metadata_getter is None assert mock_ss_thing._csm is None - assert mock_ss_thing._capture_thread is None - assert mock_ss_thing._scan_images_taken is None # Return the mock thing for further state testing, and the # exec_info of any uncaught exceptions that were raised @@ -253,3 +252,224 @@ def test_outer_scan_wo_sample_skip(): assert exec_info is None # Checked the mocked _run_scan was run exactly once assert mock_ss_thing.mock_call_count["_run_scan"] == 1 + + +MOCK_SCAN_NAME = "test_name_0001" +MOCK_SCAN_DIR = "scans/test_name_0001/images/" +MOCK_START_POS = {"x": 123, "y": 456, "z": 789} + + +def _expected_scan_data(): + """Return the expected ScanData object for a SmartScan with default properties.""" + expected_dict = { + "scan_name": MOCK_SCAN_NAME, + "starting_position": MOCK_START_POS, + "overlap": 0.45, + "max_dist": 45000, + "dx": 100, + "dy": 100, + "autofocus_dz": 1000, + "autofocus_on": True, + "skip_background": True, + "stitch_automatically": True, + "correlation_resize": 0.5, + "save_resolution": (1640, 1232), + } + return ScanData(start_time=datetime.now(), **expected_dict) + + +@pytest.fixture +def scan_thing_mocked_for_scan_data(smart_scan_thing, mocker): + """Return a scan thing that is mocked so that _collect_scan_data will run.""" + # Give the scan thing a scan invocation logger so it thinks a scan is running. + smart_scan_thing._scan_logger = LOGGER + mocker.patch.object( + smart_scan_thing, "_calc_displacement_from_test_image", return_value=[100, 100] + ) + mock_ongoing_scan = mocker.Mock() + type(mock_ongoing_scan).name = mocker.PropertyMock(return_value=MOCK_SCAN_NAME) + type(mock_ongoing_scan).images_dir = mocker.PropertyMock(return_value=MOCK_SCAN_DIR) + mock_stage = mocker.Mock() + type(mock_stage).position = mocker.PropertyMock(return_value=MOCK_START_POS) + + smart_scan_thing._ongoing_scan = mock_ongoing_scan + smart_scan_thing._stage = mock_stage + return smart_scan_thing + + +def test_collect_scan_data(scan_thing_mocked_for_scan_data): + """Run _collect_scan_data, and check the ScanData object has the expected values.""" + scan_thing = scan_thing_mocked_for_scan_data + + data = scan_thing._collect_scan_data() + expected_data = _expected_scan_data() + time_diff = expected_data.start_time - data.start_time + assert abs(time_diff.total_seconds()) < 1 + # Set times to exactly the same before final comparison + expected_data.start_time = data.start_time + assert data == expected_data + + +def test_save_final_scan_data(scan_thing_mocked_for_scan_data): + """Run _save_final_scan_data, check save is called with final results in ScanData.""" + scan_thing = scan_thing_mocked_for_scan_data + + scan_thing._scan_data = scan_thing._collect_scan_data() + scan_thing._scan_data.image_count = 44 + scan_thing._save_final_scan_data("Mocked!") + # _ongoing_scan is a mock so we can check that save_scan data was called and get + # the value + scan_thing._ongoing_scan.save_scan_data.assert_called() + final_data = scan_thing._ongoing_scan.save_scan_data.call_args[0][0] + assert isinstance(final_data, ScanData) + assert final_data.scan_result == "Mocked!" + assert final_data.image_count == 44 + assert final_data.duration.total_seconds() < 1 + + +@pytest.fixture +def scan_thing_mocked_for_run_scan(scan_thing_mocked_for_scan_data, mocker): + """Return a scan_thing mocked so _run_scan works. + + main_scan_loop(), _return_to_starting_position(), _perform_final_stitch(), and + purge_empty_scans() are Mocks + _cam is a MockCameraThing + """ + scan_thing = scan_thing_mocked_for_scan_data + scan_thing._cam = MockCameraThing() + mocker.patch.object(scan_thing, "_cancel") + mocker.patch.object(scan_thing, "_main_scan_loop") + mocker.patch.object(scan_thing, "_return_to_starting_position") + mocker.patch.object(scan_thing, "_perform_final_stitch") + mocker.patch.object(scan_thing, "purge_empty_scans") + + return scan_thing + + +def check_run_scan(scan_thing, caplog, expected_exception=None): + """Run _run_scan with a mocked scan_thing, capture logs and call counts. + + This can be used to check how run_scan executes in different exit modes. + A separate tests is above for scan completing with no exception. + + :returns: a tuple of: + + * The final result as reported in scan_data + * The logging records + * a dictionary of call counts + """ + if expected_exception is None: + with caplog.at_level(logging.WARNING): + scan_thing._scan_data = scan_thing._run_scan() + else: + with pytest.raises(expected_exception), caplog.at_level(logging.WARNING): + scan_thing._scan_data = scan_thing._run_scan() + # The preview stitcher object should still exist. And images dir should be set. + assert scan_thing._preview_stitcher.images_dir == MOCK_SCAN_DIR + + final_scan_data = scan_thing._ongoing_scan.save_scan_data.call_args[0][0] + calls = { + "cam_start_streaming_calls": scan_thing._cam.start_streaming.call_count, + "main_scan_loop_calls": scan_thing._main_scan_loop.call_count, + "return_to_start_calls": scan_thing._return_to_starting_position.call_count, + "perform_final_stitch_calls": scan_thing._perform_final_stitch.call_count, + "purge_empty_scans_calls": scan_thing.purge_empty_scans.call_count, + "save_scan_data_calls": scan_thing._ongoing_scan.save_scan_data.call_count, + } + return final_scan_data.scan_result, caplog.records, calls + + +def test_run_scan(scan_thing_mocked_for_run_scan, caplog): + """Run _save_final_scan_data, check save is called with final results in ScanData.""" + result, logs, calls = check_run_scan(scan_thing_mocked_for_run_scan, caplog) + + assert result == "success" + assert len(logs) == 0 + + expected_calls_numbers = { + "cam_start_streaming_calls": 2, + "main_scan_loop_calls": 1, + "return_to_start_calls": 1, + "perform_final_stitch_calls": 1, + "purge_empty_scans_calls": 1, + "save_scan_data_calls": 2, + } + assert calls == expected_calls_numbers + + +def test_run_scan_err_in_main_loop(scan_thing_mocked_for_run_scan, caplog, mocker): + """Check correct methods called if main_loop errors.""" + scan_thing = scan_thing_mocked_for_run_scan + mocker.patch.object( + scan_thing, "_main_scan_loop", side_effect=FileNotFoundError("mocked") + ) + + result, logs, calls = check_run_scan(scan_thing, caplog, FileNotFoundError) + + assert result.startswith("FileNotFoundError:") + assert len(logs) == 1 + assert logs[0].levelno == logging.ERROR + + # Main loop not run, nor are return to start, final stitch, or purging of empty + # scans. Save scan data is still called twice + expected_calls_numbers = { + "cam_start_streaming_calls": 2, + "main_scan_loop_calls": 1, + "return_to_start_calls": 0, + "perform_final_stitch_calls": 0, + "purge_empty_scans_calls": 0, + "save_scan_data_calls": 2, + } + assert calls == expected_calls_numbers + + +def test_run_scan_cancelled(scan_thing_mocked_for_run_scan, caplog, mocker): + """Check correct methods called scan is cancelled.""" + scan_thing = scan_thing_mocked_for_run_scan + mocker.patch.object( + scan_thing, "_main_scan_loop", side_effect=InvocationCancelledError() + ) + + result, logs, calls = check_run_scan(scan_thing, caplog) + + assert result == "cancelled by user" + # No logs at warning level. + assert len(logs) == 0 + + # Main loop not run, nor are return to start, final stitch, or purging of empty + # scans. Save scan data is still called twice + expected_calls_numbers = { + "cam_start_streaming_calls": 2, + "main_scan_loop_calls": 1, + "return_to_start_calls": 1, + "perform_final_stitch_calls": 1, + "purge_empty_scans_calls": 1, + "save_scan_data_calls": 2, + } + assert calls == expected_calls_numbers + + +def test_run_scan_fill_disk(scan_thing_mocked_for_run_scan, caplog, mocker): + """Check correct methods called if disk fills up.""" + scan_thing = scan_thing_mocked_for_run_scan + mocker.patch.object( + scan_thing, "_main_scan_loop", side_effect=NotEnoughFreeSpaceError() + ) + + result, logs, calls = check_run_scan(scan_thing, caplog, NotEnoughFreeSpaceError) + + assert result.startswith("NotEnoughFreeSpaceError:") + assert len(logs) == 1 + assert logs[0].levelno == logging.ERROR + + # Main loop not run, nor are return to start, final stitch, or purging of empty + # scans. Save scan data is still called twice + expected_calls_numbers = { + "cam_start_streaming_calls": 2, + "main_scan_loop_calls": 1, + "return_to_start_calls": 0, + "perform_final_stitch_calls": 0, + "purge_empty_scans_calls": 0, + "save_scan_data_calls": 2, + } + assert calls == expected_calls_numbers diff --git a/tests/test_stitching.py b/tests/test_stitching.py new file mode 100644 index 00000000..ef248b82 --- /dev/null +++ b/tests/test_stitching.py @@ -0,0 +1,273 @@ +"""Test that the code that talks to the external stitching process acts as expected. + +This does not actually run stitching. Instead it checks that the expected commands are +generated, and the subprocess calling works as expected. +""" + +import logging +import os +from copy import copy +import uuid +import threading +from time import sleep + +import pytest + +import labthings_fastapi as lt + +from openflexure_microscope_server.stitching import ( + BaseStitcher, + PreviewStitcher, + FinalStitcher, + STITCHING_RESOLUTION, + StitcherValidationError, +) + +# A global logger pretending to be an Invocation Logger +LOGGER = logging.getLogger("mock-invocation_logger") +FAKE_DIR: list[str] = os.path.join("a", "dir", "that", "is", "fake") +THIS_DIR: str = os.path.dirname(os.path.realpath(__file__)) +MOCK_STITCHER: str = os.path.join(THIS_DIR, "mock_stitching", "mock-stitch.py") + + +def test_base_stitcher(): + """Test the logic in BaseStitcher. + + Pretty much the only logic in base stitcher is forming a command, and calculating + the min_overlap from overlap. + + The BaseStitcher can't start as the start method is explicitly NotImplemented. + """ + # Overlaps and expected minimum overlap to be in command line argument. + overlaps = [(0.1, "0.09"), (0.4, "0.36"), (0.8, "0.72")] + + for overlap, min_overlap in overlaps: + expected_command = [ + "openflexure-stitch", + "--stitching_mode", + "all", + "--minimum_overlap", + min_overlap, + "--resize", + "0.5", + FAKE_DIR, + ] + stitcher = BaseStitcher(FAKE_DIR, overlap=overlap, correlation_resize=0.5) + assert stitcher.command == expected_command + + +def test_preview_stitcher_command(): + """Check preview stitcher command for a specific example.""" + expected_command = [ + "openflexure-stitch", + "--stitching_mode", + "preview_stitch", + "--minimum_overlap", + "0.09", + "--resize", + "0.5", + FAKE_DIR, + ] + stitcher = PreviewStitcher(FAKE_DIR, overlap=0.1, correlation_resize=0.5) + assert stitcher.command == expected_command + + +FINAL_EXPECTED_COMMAND = [ + "openflexure-stitch", + "--stitching_mode", + "all", + "--stitch_dzi", + "--no-stitch_tiff", + "--minimum_overlap", + "0.09", + "--resize", + "0.5", + FAKE_DIR, +] + + +def test_final_stitcher_command_defaults(caplog): + """Check the FinalStitcher stitches with expected default values. + + It should warn when default values are used as they are a fallback. + """ + n_logs = 0 + # Test with no dictionary data and with irrelevant dictionary data. + with caplog.at_level(logging.WARNING): + for data_dict in [None, {"irrelevant": "data"}]: + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER, scan_data_dict=data_dict) + # Should log for overlap being None and correlation_resize being None + n_logs += 2 + assert len(caplog.records) == n_logs + assert stitcher.command == FINAL_EXPECTED_COMMAND + + +def test_final_stitcher_command_tiff(caplog): + """Check that the tiff can be requested.""" + # Modify defaults + expected_command = copy(FINAL_EXPECTED_COMMAND) + expected_command[4] = "--stitch_tiff" + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER, stitch_tiff=True) + # Should log for overlap being None and correlation_resize being None + assert len(caplog.records) == 2 + assert stitcher.command == expected_command + + +def test_final_stitcher_command_set_val_directly(): + """Check that values are set as expected when directly input.""" + # Modify defaults + expected_command = copy(FINAL_EXPECTED_COMMAND) + expected_command[6] = "0.36" + expected_command[8] = "0.25" + # Test with no data dictionary, irrelevant data, and also the wrong data + # When wrong data is submitted, it should take the directly input data. + dict_vals = [ + None, + {"irrelevant": "data"}, + {"overlap": 0.2, "save_resolution": [5, 5]}, + ] + for data_dict in dict_vals: + stitcher = FinalStitcher( + FAKE_DIR, + logger=LOGGER, + overlap=0.4, + correlation_resize=0.25, + scan_data_dict=data_dict, + ) + assert stitcher.command == expected_command + + +def test_final_stitcher_command_set_with_dict(): + """Check that values are set as expected when set from a ScanData dictionary.""" + # Modify defaults + expected_command = copy(FINAL_EXPECTED_COMMAND) + expected_command[6] = "0.36" + expected_command[8] = "0.25" + # Check same thing works with a dictionary, resize is calculated from the saved image + # resolution. Make 4x bigger than STITCHING_RESOLUTION to get 0.25 + resolution = [dim * 4 for dim in STITCHING_RESOLUTION] + # Check legacy key as well as current one: + for resolution_key in ["save_resolution", "capture resolution"]: + stitcher = FinalStitcher( + FAKE_DIR, + logger=LOGGER, + scan_data_dict={"overlap": 0.4, resolution_key: resolution}, + ) + assert stitcher.command == expected_command + + +def _validation_error_tester(scan_path, **kwargs): + """Check each type of stitcher throws a validation error for the given init args.""" + # If scan_data_dict is in the kwargs only test the scan_data_dict + if "scan_data_dict" not in kwargs: + with pytest.raises(StitcherValidationError): + BaseStitcher(scan_path, **kwargs).command + with pytest.raises(StitcherValidationError): + PreviewStitcher(scan_path, **kwargs).command + with pytest.raises(StitcherValidationError): + FinalStitcher(scan_path, logger=LOGGER, **kwargs).command + + +def test_validation_error(): + """Test a number of ways to try to inject malicious arguments into the stitcher. + + The stitcher should throw a validation error each attempt. + """ + _validation_error_tester("/dir;rm -rf /;", overlap=".2", correlation_resize=".25") + _validation_error_tester(FAKE_DIR, overlap=".2", correlation_resize=".25;rm -rf /;") + _validation_error_tester(FAKE_DIR, overlap=".2;rm -rf /;", correlation_resize=".25") + _validation_error_tester(FAKE_DIR, scan_data_dict={"overlap": ".2;rm -rf /;"}) + + +def test_extra_arg_validation(): + """Test that malicious arguments in extra_args also throw validation error. + + Currently extra args do not come from user input. But this makes checks more + future-proof. + """ + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER) + stitcher._extra_args = ["&&rm -rf /&&"] + with pytest.raises(StitcherValidationError): + stitcher.command + + +def test_preview_stitching_command(caplog, mocker): + """Check the preview process runs in a background thread and doesn't log.""" + mock_cmd = "python -m mock_command.py" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + stitcher = PreviewStitcher(FAKE_DIR, overlap=0.1, correlation_resize=0.5) + stitcher.start() + # Should take a second or so to run so will still be running + assert stitcher.running + # Can't start another time, instead get a runtime error + with pytest.raises(RuntimeError): + stitcher.start() + # Wait for it to complete + stitcher.wait() + # It is now not running + assert not stitcher.running + assert len(caplog.records) == 0 + + +def test_final_stitching_command(caplog, mocker): + """Check the final stitch runs until completion, and print statements are logged.""" + mock_cmd = f"python {MOCK_STITCHER}" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + # Input values to prevent logging + stitcher = FinalStitcher( + FAKE_DIR, logger=LOGGER, overlap=0.1, correlation_resize=0.5 + ) + # For the final stitcher it will always complete before returning. + stitcher.run(cancel=lt.deps.CancelHook(id=uuid.uuid4())) + # The mock command logs the inputs (but not the initial command) and the + # stitcher logs # "Stitching complete" when it ends. + assert len(caplog.records) == len(FINAL_EXPECTED_COMMAND) + for i, record in enumerate(caplog.records): + msg = record.message.strip() + if i == len(FINAL_EXPECTED_COMMAND) - 1: + assert msg == "Stitching complete" + else: + assert msg == FINAL_EXPECTED_COMMAND[i + 1] + + +def test_final_stitching_command_cancelled(caplog, mocker): + """Check that final stitch can be cancelled.""" + mock_cmd = f"python {MOCK_STITCHER}" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER) + cancel_hook = lt.deps.CancelHook(id=uuid.uuid4()) + + # Start stitching in a thread. + thread = threading.Thread(target=stitcher.run, kwargs={"cancel": cancel_hook}) + thread.start() + # Sleep long enough for at least 1 log. + sleep(0.5) + # use set() to cancel! + cancel_hook.set() + thread.join() + assert len(caplog.records) < len(FINAL_EXPECTED_COMMAND) + 1 + assert caplog.records[-1].message == "Stitching cancelled by user" + + +def test_final_stitching_command_error(caplog, mocker): + """Check that ChildProcessError is raised if the final stitch errors.""" + mock_cmd = f"python {MOCK_STITCHER}" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER) + # Send in the argument ERROR to mock-stitch and it will raise an error rather + # than echo. + stitcher._extra_args = ["ERROR"] + with pytest.raises(ChildProcessError): + stitcher.run(cancel=lt.deps.CancelHook(id=uuid.uuid4()))