Convert milling GIF API endpoint into workflow instead (#829)

* Converted 'make_milling_gif' endpoint function into a workflow instead to avoid endpoint timing out
* Removed intermediate NumPy array and convert frame-by-frame straight to PIL Image objects to reduce memory footprint
* Updated test for 'make_gif' endpoint function
* Added test module for 'make_milling_gif' workflow

Author: tieneupin

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"

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"),
+        },
+    )

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

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()

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")