diff --git a/pyproject.toml b/pyproject.toml index e7e200657..d16e87965 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,7 @@ TomographyMetadataContext = "murfey.client.contexts.tomo_metadata:TomographyMeta "data_collection" = "murfey.workflows.register_data_collection:run" "data_collection_group" = "murfey.workflows.register_data_collection_group:run" "experiment_type_update" = "murfey.workflows.register_experiment_type_update:run" +"fib.make_milling_gif" = "murfey.workflows.fib.make_milling_gif:run" "fib.register_atlas" = "murfey.workflows.fib.register_atlas:run" "fib.register_milling_progress" = "murfey.workflows.fib.register_milling_progress:run" "pato" = "murfey.workflows.notifications:notification_setup" diff --git a/src/murfey/server/api/workflow_fib.py b/src/murfey/server/api/workflow_fib.py index 94c3de8c1..1b02e6e7a 100644 --- a/src/murfey/server/api/workflow_fib.py +++ b/src/murfey/server/api/workflow_fib.py @@ -1,20 +1,12 @@ import logging -import os from pathlib import Path -import numpy as np -import PIL.Image from fastapi import APIRouter, Depends from pydantic import BaseModel -from sqlmodel import select -import murfey.util.db as MurfeyDB from murfey.server import _transport_object from murfey.server.api.auth import validate_instrument_token -from murfey.server.murfey_db import murfey_db -from murfey.util import sanitise_path -from murfey.util.config import get_machine_config -from murfey.util.models import LamellaSiteInfo +from murfey.util.models import FIBGIFParameters, LamellaSiteInfo logger = logging.getLogger("murfey.server.api.workflow_fib") @@ -65,76 +57,19 @@ def register_fib_milling_progress( ) -class FIBGIFParameters(BaseModel): - lamella_number: int - images: list[Path] - output_file: Path - - @router.post("/sessions/{session_id}/make_gif") async def make_gif( session_id: int, gif_params: FIBGIFParameters, - db=murfey_db, ): - # Load machine config and session info - session_entry = db.exec( - select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) - ).one() - instrument_name = session_entry.instrument_name - visit_name = session_entry.visit - machine_config = get_machine_config(instrument_name=instrument_name)[ - instrument_name - ] - rsync_basepath = machine_config.rsync_basepath or Path(".").resolve() - - # Sanitise and verify that the output directory is relative to rsync basepath - output_file = sanitise_path(gif_params.output_file) - if not output_file.is_relative_to(rsync_basepath): - logger.error("Output file path is not permitted") - raise ValueError - - # Create folders in the visit directory and onwards and change permissions - visit_index = output_file.parts.index(visit_name) - for current_path in list(reversed(output_file.parents))[visit_index + 1 :]: - if not current_path.exists(): - current_path.mkdir(parents=True) - logger.debug(f"Created output directory {current_path}") - try: - os.chmod(current_path, mode=machine_config.mkdir_chmod) - except PermissionError: - logger.warning( - f"Insufficient permissions to modify directory {current_path}" - ) - continue - - # Load the images as PIL Image objects - arr: list[np.ndarray] = [] - for f in gif_params.images: - with PIL.Image.open(f) as im: - im.thumbnail((512, 512)) - frame = np.array(im, dtype=np.float32) - vmin, vmax = np.percentile(frame, (0.5, 99.5)) - scale = 255 / ((vmax - vmin) or 1) - np.clip(frame, a_min=vmin, a_max=vmax, out=frame) - np.subtract(frame, vmin, out=frame) - np.multiply(frame, scale, out=frame) - arr.append(frame.astype(np.uint8)) - arr = np.array(arr).astype(np.uint8) - - # Convert back to PIL.Image objects and save as GIF - try: - converted = [PIL.Image.fromarray(a, mode="L") for a in arr] - converted[0].save( - output_file, - format="GIF", - append_images=converted[1:], - save_all=True, - duration=30, - loop=0, - ) - logger.info(f"Created GIF file {output_file}") - return {"output_gif": str(output_file)} - finally: - for im in converted: - im.close() + if _transport_object is None: + logger.error("No TransportManager object was set up") + return None + _transport_object.send( + _transport_object.feedback_queue, + { + "register": "fib.make_milling_gif", + "session_id": session_id, + "gif_params": gif_params.model_dump(mode="json"), + }, + ) diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py index a13118e4d..c8b7af7ac 100644 --- a/src/murfey/util/models.py +++ b/src/murfey/util/models.py @@ -233,6 +233,12 @@ class LamellaSiteInfo(BaseModel): steps: MillingSteps | None = None +class FIBGIFParameters(BaseModel): + lamella_number: int + images: list[Path] + output_file: Path + + """ ======================================================================================= Single Particle Analysis diff --git a/src/murfey/workflows/fib/make_milling_gif.py b/src/murfey/workflows/fib/make_milling_gif.py new file mode 100644 index 000000000..109e14ec9 --- /dev/null +++ b/src/murfey/workflows/fib/make_milling_gif.py @@ -0,0 +1,93 @@ +import logging +import os +from pathlib import Path +from typing import Any + +import numpy as np +import PIL.Image +from sqlmodel import Session as SQLModelSession, select + +import murfey.util.db as MurfeyDB +from murfey.util import sanitise_path +from murfey.util.config import get_machine_config +from murfey.util.models import FIBGIFParameters + +logger = logging.getLogger(__name__) + + +def run(message: dict[str, Any], murfey_db: SQLModelSession): + # Outer try-finally block to close Murfey DB with + try: + try: + # Parse and unpack incoming message + session_id = int(message["session_id"]) + gif_params = FIBGIFParameters(**message["gif_params"]) + except Exception: + logger.error("Error parsing contents of message", exc_info=True) + return {"success": False, "requeue": False} + + # Load machine config and session info + session_entry = murfey_db.exec( + select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) + ).one() + instrument_name = session_entry.instrument_name + visit_name = session_entry.visit + machine_config = get_machine_config(instrument_name=instrument_name)[ + instrument_name + ] + rsync_basepath = machine_config.rsync_basepath or Path(".").resolve() + + # Sanitise and verify that the output directory is relative to rsync basepath + output_file = sanitise_path(gif_params.output_file) + if not output_file.is_relative_to(rsync_basepath): + raise ValueError("Output file path is not permitted") + + # Create folders in the visit directory and onwards and change permissions + visit_index = output_file.parts.index(visit_name) + for current_path in list(reversed(output_file.parents))[visit_index + 1 :]: + if not current_path.exists(): + current_path.mkdir(parents=True) + logger.debug(f"Created output directory {current_path}") + try: + os.chmod(current_path, mode=machine_config.mkdir_chmod) + except PermissionError: + logger.warning( + f"Insufficient permissions to modify directory {current_path}" + ) + continue + + # Load the images as PIL Image objects + converted: list[PIL.Image.Image] = [] + for f in gif_params.images: + with PIL.Image.open(f) as im: + im.thumbnail((512, 512)) + frame = np.array(im, dtype=np.float32) + # Normalise to 8-bit + vmin, vmax = np.percentile(frame, (0.5, 99.5)) + scale = 255 / ((vmax - vmin) or 1) + np.clip(frame, a_min=vmin, a_max=vmax, out=frame) + np.subtract(frame, vmin, out=frame) + np.multiply(frame, scale, out=frame) + # Convert back to PIL Image + converted.append( + PIL.Image.fromarray(frame.astype(np.uint8), mode="L").copy() + ) + del frame # Explicitly remove frame from memory + # Save stack as a GIF + if not converted: + raise ValueError("No images were provided or loaded") + converted[0].save( + output_file, + format="GIF", + append_images=converted[1:], + save_all=True, + duration=30, + loop=0, + ) + logger.info(f"Created GIF file {output_file}") + return {"success": True} + except Exception: + logger.error("Error creating FIB milling GIF", exc_info=True) + return {"success": False, "requeue": False} + finally: + murfey_db.close() diff --git a/tests/server/api/test_workflow_fib.py b/tests/server/api/test_workflow_fib.py index 1fd400c4d..3a8f997b0 100644 --- a/tests/server/api/test_workflow_fib.py +++ b/tests/server/api/test_workflow_fib.py @@ -1,17 +1,15 @@ from pathlib import Path from unittest.mock import MagicMock -import numpy as np -import PIL.Image import pytest from pytest_mock import MockerFixture from murfey.server.api.workflow_fib import ( FIBAtlasFile, - FIBGIFParameters, make_gif, register_fib_atlas, ) +from murfey.util.models import FIBGIFParameters @pytest.mark.parametrize( @@ -33,7 +31,7 @@ def test_register_fib_atlas( # Mock the logger mock_logger = mocker.patch("murfey.server.api.workflow_fib.logger") - # Mock the tranposrt object + # Mock the transport object if has_transport_object: mock_transport_object = MagicMock() mock_transport_object.feedback_queue = "dummy" @@ -67,74 +65,61 @@ def test_register_fib_atlas( mock_logger.error.assert_called_with("No TransportManager object was set up") +@pytest.mark.parametrize( + "has_transport_object", + ( + True, + False, + ), +) @pytest.mark.asyncio async def test_make_gif( mocker: MockerFixture, tmp_path: Path, + has_transport_object: bool, ): - # Set up test variables - session_id = 10 - instrument_name = "test_instrument" - rsync_basepath = tmp_path / "data" - visit_name = "cm12345-6" - year = 2020 - visit_dir = rsync_basepath / str(year) / visit_name - lamella_num = 12 - lamella_folder = "Lamella" - if lamella_num > 1: - lamella_folder += f" ({lamella_num})" - output_file = ( - visit_dir - / "processed" - / "project_name" - / "grid_1" - / "drift_correction" - / f"lamella_{lamella_num}.gif" - ) - - # Create a list of test image file paths - raw_images = [ - visit_dir - / "autotem" - / visit_name - / "Sites" - / lamella_folder - / "DCImages/DCM_asdfjkl/asdfjkl-Polishing-dc_rescan-image-.png" - ] * 5 - # Mock the output of PIL.Image.open to always return a NumPY array - mocker.patch( - "murfey.server.api.workflow_fib.PIL.Image.open", - return_value=PIL.Image.fromarray(np.ones((512, 512), dtype=np.uint16)), - ) - - # Create the Pydantic model - params = FIBGIFParameters( - lamella_number=lamella_num, - images=[str(f) for f in raw_images], - output_file=output_file, - ) + # Set up the variables + session_id = 1 + gif_params_dict = { + "lamella_number": 1, + "images": [ + str(tmp_path / "some_file.png"), + ], + "output_file": str(tmp_path / "target_file.gif"), + } + gif_params = FIBGIFParameters(**gif_params_dict) - # Mock the database query - mock_db = MagicMock() - mock_db.exec.return_value.one.return_value.instrument_name = instrument_name - mock_db.exec.return_value.one.return_value.visit = visit_name + # Mock the logger + mock_logger = mocker.patch("murfey.server.api.workflow_fib.logger") - # Mock the machine config and 'get_machine_config' - mock_machine_config = MagicMock() - mock_machine_config.mkdir_chmod = 0o2775 - mock_machine_config.rsync_basepath = rsync_basepath - mocker.patch( - "murfey.server.api.workflow_fib.get_machine_config", - return_value={ - instrument_name: mock_machine_config, - }, - ) + # Mock the transport object + if has_transport_object: + mock_transport_object = MagicMock() + mock_transport_object.feedback_queue = "dummy" + mocker.patch( + "murfey.server.api.workflow_fib._transport_object", + mock_transport_object, + ) + else: + mocker.patch( + "murfey.server.api.workflow_fib._transport_object", + None, + ) - # Run the function and check that the expected outputs are there - result = await make_gif( + # Run the function and check that the expected calls were made + await make_gif( session_id=session_id, - gif_params=params, - db=mock_db, + gif_params=gif_params, ) - assert output_file.exists() - assert result.get("output_gif") == str(output_file) + + if has_transport_object: + mock_transport_object.send.assert_called_with( + "dummy", + { + "register": "fib.make_milling_gif", + "session_id": session_id, + "gif_params": gif_params_dict, + }, + ) + else: + mock_logger.error.assert_called_with("No TransportManager object was set up") diff --git a/tests/workflows/fib/test_make_milling_gif.py b/tests/workflows/fib/test_make_milling_gif.py new file mode 100644 index 000000000..a402273b6 --- /dev/null +++ b/tests/workflows/fib/test_make_milling_gif.py @@ -0,0 +1,84 @@ +from pathlib import Path +from unittest.mock import MagicMock + +import numpy as np +import PIL.Image +from pytest_mock import MockerFixture + +from murfey.util.models import FIBGIFParameters +from murfey.workflows.fib.make_milling_gif import run + + +def test_make_gif( + mocker: MockerFixture, + tmp_path: Path, +): + # Set up test variables + session_id = 10 + instrument_name = "test_instrument" + rsync_basepath = tmp_path / "data" + visit_name = "cm12345-6" + year = 2020 + visit_dir = rsync_basepath / str(year) / visit_name + lamella_num = 12 + lamella_folder = "Lamella" + if lamella_num > 1: + lamella_folder += f" ({lamella_num})" + output_file = ( + visit_dir + / "processed" + / "project_name" + / "grid_1" + / "drift_correction" + / f"lamella_{lamella_num}.gif" + ) + + # Create a list of test image file paths + raw_images = [ + visit_dir + / "autotem" + / visit_name + / "Sites" + / lamella_folder + / "DCImages/DCM_asdfjkl/asdfjkl-Polishing-dc_rescan-image-.png" + ] * 5 + # Mock the output of PIL.Image.open to always return a NumPY array + mocker.patch( + "murfey.workflows.fib.make_milling_gif.PIL.Image.open", + return_value=PIL.Image.fromarray(np.ones((512, 512), dtype=np.uint16)), + ) + + # Create the Pydantic model + params = FIBGIFParameters( + lamella_number=lamella_num, + images=[str(f) for f in raw_images], + output_file=output_file, + ) + + # Mock the database query + mock_db = MagicMock() + mock_db.exec.return_value.one.return_value.instrument_name = instrument_name + mock_db.exec.return_value.one.return_value.visit = visit_name + + # Mock the machine config and 'get_machine_config' + mock_machine_config = MagicMock() + mock_machine_config.mkdir_chmod = 0o2775 + mock_machine_config.rsync_basepath = rsync_basepath + mocker.patch( + "murfey.workflows.fib.make_milling_gif.get_machine_config", + return_value={ + instrument_name: mock_machine_config, + }, + ) + + # Run the function and check that the expected outputs are there + result = run( + message={ + "register": "fib.make_milling_gif", + "session_id": session_id, + "gif_params": params.model_dump(mode="json"), + }, + murfey_db=mock_db, + ) + assert output_file.exists() + assert result.get("success", False)