Register lamella evaluation image

src/murfey/workflows/fib/register_atlas.py

@@ -251,11 +251,13 @@ def run(
                 fib_info.session_id, metadata, murfey_db
             )
             logger.info(
-                f"Registered FIB atlas image {fib_info.atlas_file} for slot {metadata.slot_number} in Murfey database"
+                f"Registered FIB atlas image {fib_info.atlas_file} "
+                f"for slot {metadata.slot_number} in Murfey database"
             )
         except Exception:
             logger.error(
-                f"Error registering FIB atlas image {fib_info.atlas_file} in Murfey database",
+                "Error registering FIB atlas image "
+                f"{fib_info.atlas_file} in Murfey database",
                 exc_info=True,
             )
             return {"success": False, "requeue": False}
@@ -273,7 +275,7 @@ def run(
         except Exception:
             # Log error but allow workflow to proceed
             logger.error(
-                "Exception encountered when registering data collection group for FIB workflow "
+                "Error registering data collection group for FIB workflow "
                 f"for {metadata.site_name!r}: \n"
                 f"{traceback.format_exc()}"
             )

src/murfey/workflows/fib/register_lamella_evaluation_image.py

@@ -1,27 +1,94 @@
 import json
 import logging
+import re
+from datetime import datetime
 from pathlib import Path
 from typing import Any, cast
 
 from pydantic import BaseModel
-from sqlmodel import Session, select
+from sqlmodel import Session as SQLModelSession, select
 
 import murfey.util.db as MurfeyDB
 from murfey.util.config import get_machine_config
 from murfey.util.models import FIBImageMetadata
-from murfey.workflows.fib.shared import parse_image_metadata
+from murfey.workflows.fib.shared import (
+    parse_image_metadata,
+    populate_fib_imaging_site_entry,
+)
 
 logger = logging.getLogger(__name__)
 
 
+# The timestamp in the lamella evaluation image follows the pattern
+# yyyy-mm-dd-HH-MM-SS
+# E.g.
+#   2026-03-09-18-24-51_drift_corrected_image_Finer Milling - Electron Image.png
+#   2026-03-10-16-06-25_drift_corrected_image_Polishing 2 - Electron Image.png
+# This can be searched for using regex
+# (?<!\d) --> Character prior to pattern CANNOT be a digit
+# (?!\d)  --> Character after pattern CANNOT be a digit
+pattern = re.compile(r"(?<!\d)\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2}(?!\d)")
+
+
+def _get_timestamp(name: str):
+    """
+    Helper function to extract the datetime information from the lamella evaluation
+    image file name.
+    """
+    if (match := pattern.search(name)) is not None:
+        return datetime.strptime(match.group(), "%Y-%m-%d-%H-%M-%S")
+    raise ValueError(f"No datetime match found in {name}")
+
+
+def _register_fib_imaging_site(
+    session_id: int,
+    metadata: FIBImageMetadata,
+    murfey_db: SQLModelSession,
+):
+    """
+    Register FIB atlas in Murfey database or update existing entry.
+    """
+    if (
+        fib_imaging_site := murfey_db.exec(
+            select(MurfeyDB.ImagingSite)
+            .where(MurfeyDB.ImagingSite.session_id == session_id)
+            .where(MurfeyDB.ImagingSite.site_name == metadata.site_name)
+            .where(MurfeyDB.ImagingSite.data_type == "grid_square")
+        ).one_or_none()
+    ) is None:
+        # Create new entry if one doesn't already exist
+        fib_imaging_site = MurfeyDB.ImagingSite(
+            session_id=session_id,
+            site_name=metadata.site_name,
+            image_path=str(metadata.file),
+            data_type="grid_square",
+        )
+        fib_imaging_site = populate_fib_imaging_site_entry(fib_imaging_site, metadata)
+    else:
+        # Check if image was acquired after the current one
+        incoming_timestamp = _get_timestamp(metadata.file.stem)
+        # Handle empty string
+        current_timestamp = datetime.min
+        if fib_imaging_site.image_path:
+            current_timestamp = _get_timestamp(Path(fib_imaging_site.image_path).stem)
+        # Update if incoming one is newer
+        if incoming_timestamp >= current_timestamp:
+            fib_imaging_site = populate_fib_imaging_site_entry(
+                fib_imaging_site, metadata
+            )
+    murfey_db.add(fib_imaging_site)
+    murfey_db.commit()
+    return fib_imaging_site
+
+
 class FIBLamellaImageInfo(BaseModel):
     session_id: int
     lamella_image_file: Path
 
 
 def run(
     message: dict[str, Any],
-    murfey_db: Session,
+    murfey_db: SQLModelSession,
 ):
     # Outer try-finally block to ensure the database connection is closed
     logger.info(
@@ -77,6 +144,21 @@ def run(
             )
             return {"success": False, "requeue": False}
 
+        try:
+            # Register imaging site to Murfey, or update existing one
+            _ = _register_fib_imaging_site(fib_info.session_id, metadata, murfey_db)
+            logger.info(
+                f"Registered lamella evaluation image {fib_info.lamella_image_file} "
+                f"for slot {metadata.slot_number} in Murfey database"
+            )
+        except Exception:
+            logger.error(
+                "Error registering lamella evaluation image "
+                f"{fib_info.lamella_image_file} in Murfey database",
+                exc_info=True,
+            )
+            return {"success": False, "requeue": False}
+
         return {"success": True}
     finally:
         murfey_db.close()

tests/workflows/fib/test_register_lamella_evaluation_image.py

@@ -4,12 +4,17 @@
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlmodel import Session as SQLModelSession, select
 
+import murfey.util.db as MurfeyDB
 from murfey.util.config import MachineConfig
 from murfey.workflows.fib.register_lamella_evaluation_image import (
     FIBImageMetadata,
+    FIBLamellaImageInfo,
+    _register_fib_imaging_site,
     run,
 )
+from murfey.workflows.fib.shared import populate_fib_imaging_site_entry
 from tests.conftest import ExampleVisit
 
 session_id = 10
@@ -24,12 +29,123 @@ def visit_dir(tmp_path: Path):
     return visit_dir
 
 
+@pytest.mark.parametrize(
+    "test_params",
+    (  # Site exists | Newer site
+        (True, True),
+        (True, False),
+        (False, False),
+    ),
+)
+def test_register_fib_imaging_site_with_db(
+    test_params: tuple[bool, bool],
+    visit_dir: Path,
+    murfey_db_session: SQLModelSession,
+):
+    # Unpack test params
+    has_existing_entry, newer_insert = test_params
+
+    # Register a Session for this test
+    murfey_session = MurfeyDB.Session(
+        id=session_id,
+        visit=visit_name,
+        name=visit_name,
+        instrument_name=instrument_name,
+        started=True,
+    )
+    murfey_db_session.add(murfey_session)
+    murfey_db_session.commit()
+
+    # Lamella directory
+    lamella_dir = (
+        visit_dir
+        / "autotem"
+        / visit_name
+        / "Sites"
+        / "Lamella"
+        / "LamellaEvaluationImages"
+    )
+    # Standard metadata to use
+    metadata_dict = {
+        "voltage": 2000,
+        "shift_x": 0,
+        "shift_y": 0,
+        "len_x": 0.003072,
+        "len_y": 0.002048,
+        "pos_x": -0.003,
+        "pos_y": 0.003,
+        "pos_z": 0.01,
+        "rotation": 1.833,
+        "slot_number": 2,
+        "tilt_alpha": 0,
+        "tilt_beta": 0,
+        "pixels_x": 3072,
+        "pixels_y": 2048,
+        "pixel_size_x": 1e-6,
+        "pixel_size_y": 1e-6,
+    }
+
+    # Create older existing entry
+    if has_existing_entry:
+        existing_timestamp = f"2026-04-16-02-39-{30 if newer_insert else 50}"
+        existing_file = (
+            lamella_dir
+            / f"{existing_timestamp}_drift_corrected_image_Polishing 2 - Electron Image.png"
+        )
+        existing_metadata = FIBImageMetadata(
+            visit_name=visit_name,
+            file=existing_file,
+            **metadata_dict,
+        )
+        existing_site = MurfeyDB.ImagingSite(
+            session_id=session_id,
+            site_name=existing_metadata.site_name,
+            data_type="grid_square",
+        )
+        existing_site = populate_fib_imaging_site_entry(
+            existing_site, existing_metadata
+        )
+        murfey_db_session.add(existing_site)
+        murfey_db_session.commit()
+
+    # Create the test image file to register
+    file = (
+        lamella_dir
+        / "2026-04-16-02-39-40_drift_corrected_image_Polishing 2 - Electron Image.png"
+    )
+    metadata = FIBImageMetadata(
+        visit_name=visit_name,
+        file=file,
+        **metadata_dict,
+    )
+
+    # Run the function and check that results are as expected
+    _register_fib_imaging_site(
+        session_id=session_id,
+        metadata=metadata,
+        murfey_db=murfey_db_session,
+    )
+
+    # Only one entry should exist
+    found_sites = murfey_db_session.exec(select(MurfeyDB.ImagingSite)).all()
+    assert len(found_sites) == 1
+
+    # Key parameters should be populated
+    registered_site = found_sites[0]
+    assert registered_site.session_id == session_id
+    assert registered_site.data_type == "grid_square"
+    assert registered_site.site_name == metadata.site_name
+    if has_existing_entry and not newer_insert:
+        assert registered_site.image_path != str(file)
+    else:
+        assert registered_site.image_path == str(file)
+
+
 def test_run(
     mocker: MockerFixture,
     visit_dir: Path,
 ):
     # Set up parameters
-    project_name = "some_project"
 
     # Mock the logger
     mock_logger = mocker.patch(
@@ -56,7 +172,7 @@ def test_run(
     file = (
         visit_dir
         / "autotem"
-        / project_name
+        / visit_name
         / "Sites"
         / "Lamella"
         / "LamellaEvaluationImages"
@@ -92,17 +208,32 @@ def test_run(
         return_value=metadata_dict,
     )
 
+    # Mock the results of '_register_fib_image_site'
+    mock_register_imaging_site = mocker.patch(
+        "murfey.workflows.fib.register_lamella_evaluation_image._register_fib_imaging_site",
+        return_value=MagicMock(),
+    )
+
     # Construct the message to pass to the function
     message = {
         "register": "fib.register_lamella_evaluation_image",
         "session_id": session_id,
         "lamella_image_file": str(file),
     }
+    fib_info = FIBLamellaImageInfo(**message)
 
     # Run function and check that expected calls were made
     result = run(message, mock_murfey_db)
-    mock_logger.info.assert_called_with(
+
+    # Metadata should have been extracted and logged
+    mock_logger.info.assert_any_call(
         "Extracted the following metadata from the image:\n"
         f"{json.dumps(metadata.model_dump(), indent=2, default=str)}"
     )
+    # Imaging site registration function should have been called
+    mock_register_imaging_site.assert_called_once_with(
+        fib_info.session_id,
+        metadata,
+        mock_murfey_db,
+    )
     assert result["success"]