Register DataCollectionGroup and Atlas placeholders using lamella images

src/murfey/workflows/fib/register_lamella_evaluation_image.py

@@ -16,6 +16,7 @@
     parse_image_metadata,
     populate_fib_imaging_site_entry,
 )
+from murfey.workflows.register_data_collection_group import register_dcg
 
 logger = logging.getLogger(__name__)
 
@@ -116,6 +117,67 @@ def _register_fib_imaging_site(
     return fib_imaging_site
 
 
+def _register_dcg(
+    session_id: int,
+    instrument_name: str,
+    visit_name: str,
+    imaging_site: MurfeyDB.ImagingSite,
+    murfey_db: SQLModelSession,
+):
+    """
+    Takes an ImagingSite entry and uses it to create and register a DataCollectionGroup
+    entry in ISPyB if one doesn't already exist, or to populate an existing entry.
+    After doing so, it will register the DataCollectionGroup ID in Murfey and add it to
+    the ImagingSite entry.
+    """
+    # Determine variables to register data collection group and atlas with
+    proposal_code = "".join(char for char in visit_name.split("-")[0] if char.isalpha())
+    proposal_number = "".join(
+        char for char in visit_name.split("-")[0] if char.isdigit()
+    )
+    visit_number = visit_name.split("-")[-1]
+
+    # Generate a name/tag for the data collection group
+    # The name will be the site name minus the "/lamella..." specifier
+    dcg_name = "/".join(imaging_site.site_name.split("/")[:-1])
+
+    # Check if a DataCollectionGroup entry with this session and tag already exists
+    dcg_entry = murfey_db.exec(
+        select(MurfeyDB.DataCollectionGroup)
+        .where(MurfeyDB.DataCollectionGroup.session_id == session_id)
+        .where(MurfeyDB.DataCollectionGroup.tag == dcg_name)
+    ).one_or_none()
+    if not dcg_entry:
+        # Create a placeholder DataCollectionGroup and Atlas if not
+        dcg_message = {
+            "microscope": instrument_name,
+            "proposal_code": proposal_code,
+            "proposal_number": proposal_number,
+            "visit_number": visit_number,
+            "session_id": session_id,
+            "tag": dcg_name,
+            "experiment_type_id": 46,
+            "atlas": "",
+            "atlas_pixel_size": 0.0,
+            "sample": None,
+        }
+        dcg_entry = register_dcg(
+            message=dcg_message,
+            murfey_db=murfey_db,
+        )
+        if not dcg_entry:
+            raise RuntimeError(
+                "Failed to create DataCollectionGroup entry for "
+                f"{imaging_site.image_path}"
+            )
+
+    # Update the ImagingSite with the DataCollectionGroup ID
+    imaging_site.dcg_id = dcg_entry.id
+    imaging_site.dcg_name = dcg_entry.tag
+    murfey_db.add(imaging_site)
+    murfey_db.commit()
+
+
 class FIBLamellaImageInfo(BaseModel):
     session_id: int
     lamella_image_file: Path
@@ -129,84 +191,58 @@ def run(
     logger.info(
         f"Received the following message:\n{json.dumps(message, indent=2, default=str)}"
     )
-    try:
-        try:
-            # Validate incoming message
-            fib_info = FIBLamellaImageInfo(**message)
-        except Exception:
-            logger.error("Could not validate incoming message", exc_info=True)
-            return {"success": False, "requeue": False}
-
-        try:
-            # Load visit information
-            murfey_session = murfey_db.exec(
-                select(MurfeyDB.Session).where(
-                    MurfeyDB.Session.id == fib_info.session_id
-                )
-            ).one()
-            visit_name = murfey_session.visit
-            instrument_name = murfey_session.instrument_name
-        except Exception:
-            logger.error(
-                "Exception encountered while querying Murfey database", exc_info=True
-            )
-            return {"success": False, "requeue": False}
 
-        try:
-            # Load the machine config
-            machine_config = get_machine_config(instrument_name)[instrument_name]
-            rotation_offset: float = cast(
-                float, machine_config.calibrations.get("rotation_offset", 0)
-            )
+    # Validate incoming message
+    fib_info = FIBLamellaImageInfo(**message)
 
-            # Extract metadata from the image
-            metadata = FIBImageMetadata(
-                visit_name=visit_name,
-                file=fib_info.lamella_image_file,
-                **parse_image_metadata(
-                    file=fib_info.lamella_image_file,
-                    rotation_offset=rotation_offset,
-                ),
-            )
-            logger.info(
-                "Extracted the following metadata from the image:\n"
-                f"{json.dumps(metadata.model_dump(), indent=2, default=str)}"
-            )
-        except Exception:
-            logger.error(
-                f"Error extracting metadata from file {fib_info.lamella_image_file}",
-                exc_info=True,
-            )
-            return {"success": False, "requeue": False}
-
-        try:
-            # Make a thumbnail of the image and update metadata accordingly
-            metadata.thumbnail_path = _make_thumbnail(
-                file=fib_info.lamella_image_file,
-                metadata=metadata,
-                visit_name=visit_name,
-            )
-        except Exception:
-            logger.warning(
-                f"Error creating thumbnail of file {fib_info.lamella_image_file}",
-                exc_info=True,
-            )
+    # Load visit information
+    murfey_session = murfey_db.exec(
+        select(MurfeyDB.Session).where(MurfeyDB.Session.id == fib_info.session_id)
+    ).one()
+    visit_name = murfey_session.visit
+    instrument_name = murfey_session.instrument_name
 
-        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}
+    # Load the machine config
+    machine_config = get_machine_config(instrument_name)[instrument_name]
+    rotation_offset: float = cast(
+        float, machine_config.calibrations.get("rotation_offset", 0)
+    )
+
+    # Extract metadata from the image
+    metadata = FIBImageMetadata(
+        visit_name=visit_name,
+        file=fib_info.lamella_image_file,
+        **parse_image_metadata(
+            file=fib_info.lamella_image_file,
+            rotation_offset=rotation_offset,
+        ),
+    )
+    logger.info(
+        "Extracted the following metadata from the image:\n"
+        f"{json.dumps(metadata.model_dump(), indent=2, default=str)}"
+    )
+
+    # Make a thumbnail of the image and update metadata accordingly
+    metadata.thumbnail_path = _make_thumbnail(
+        file=fib_info.lamella_image_file,
+        metadata=metadata,
+        visit_name=visit_name,
+    )
+
+    # Register imaging site to Murfey, or update existing one
+    fib_img_site = _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"
+    )
+
+    # Register data collection group and atlas in ISPyB
+    _register_dcg(
+        session_id=fib_info.session_id,
+        instrument_name=instrument_name,
+        visit_name=visit_name,
+        imaging_site=fib_img_site,
+        murfey_db=murfey_db,
+    )
 
-        return {"success": True}
-    finally:
-        murfey_db.close()
+    return {"success": True}

tests/workflows/fib/test_register_atlas.py

@@ -11,7 +11,9 @@
 from sqlmodel import Session as SQLModelSession, select as sm_select
 
 import murfey.util.db as MurfeyDB
-from murfey.util.fib import get_slot_number
+import murfey.workflows.fib.register_atlas
+from murfey.server.ispyb import TransportManager
+from murfey.util.fib import get_slot_number, number_from_name
 from murfey.util.models import FIBImageMetadata
 from murfey.workflows.fib.register_atlas import run
 from tests.conftest import ExampleVisit
@@ -44,16 +46,16 @@ def test_run_with_db(
 
     # Add a test visit to the database
     if not (
-        session_entry := murfey_db_session.exec(
+        murfey_session := murfey_db_session.exec(
             sm_select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id)
         ).one_or_none()
     ):
-        session_entry = MurfeyDB.Session(id=session_id)
-    session_entry.name = visit_name
-    session_entry.visit = visit_name
-    session_entry.instrument_name = instrument_name
+        murfey_session = MurfeyDB.Session(id=session_id)
+    murfey_session.name = visit_name
+    murfey_session.visit = visit_name
+    murfey_session.instrument_name = instrument_name
 
-    murfey_db_session.add(session_entry)
+    murfey_db_session.add(murfey_session)
     murfey_db_session.commit()
 
     # Mock the MachineConfig
@@ -88,15 +90,11 @@ def test_run_with_db(
     )
 
     # Patch the TransportManager object in the workflows called
-    from murfey.server.ispyb import TransportManager
-
     mocker.patch(
         "murfey.server._transport_object", new=TransportManager("PikaTransport")
     )
 
     # Mock the metadata returned from the image file
-    import murfey.workflows.fib.register_atlas
-
     extracted = {
         "voltage": 2000,
         "shift_x": 0,
@@ -109,17 +107,18 @@ def test_run_with_db(
         "rotation": -1.309,
         "tilt_alpha": 0.8,
         "tilt_beta": 0,
-        "pixels_x": 3072,
-        "pixels_y": 2048,
+        "pixels_x": 1500,
+        "pixels_y": 1000,
         "pixel_size_x": 1e-6,
         "pixel_size_y": 1e-6,
     }
-    extracted["slot_number"] = get_slot_number(
+    slot_number = get_slot_number(
         x=extracted["pos_x"],
         y=extracted["pos_y"],
         rotation=extracted["rotation"],
         rotation_offset=rotation_offset,
     )
+    extracted["slot_number"] = slot_number
     mock_metadata = [
         FIBImageMetadata(
             visit_name=visit_name,
@@ -132,15 +131,36 @@ def test_run_with_db(
         "murfey.workflows.fib.register_atlas.parse_image_metadata",
         return_value=extracted,
     )
-    spy_register = mocker.spy(
-        murfey.workflows.fib.register_atlas,
-        "_register_fib_imaging_site",
-    )
 
     # Mock 'PIL.Image.open' and create a test image
     mock_open = mocker.patch("murfey.workflows.fib.register_atlas.PIL.Image.open")
     mock_open.__enter__.return_value = PIL.Image.fromarray(
-        np.ones((2048, 1152), dtype=np.uint8)
+        np.ones((1500, 1000), dtype=np.uint8)
+    )
+    # Build the name of the expected test images
+    thumbnails: list[Path] = []
+    for file in test_files:
+        image_number = number_from_name(file.stem)
+        thumbnail = (
+            visit_dir
+            / "processed"
+            / visit_name
+            / f"grid_{slot_number}"
+            / "atlas"
+            / f"atlas_{str(image_number).zfill(2)}.png"
+        )
+        thumbnail.parent.mkdir(parents=True, exist_ok=True)
+        thumbnail.touch(exist_ok=True)
+        thumbnails.append(thumbnail)
+
+    # Set up spies for the different steps in 'run()'
+    spy_thumbnail = mocker.spy(
+        murfey.workflows.fib.register_atlas,
+        "_make_thumbnail",
+    )
+    spy_register = mocker.spy(
+        murfey.workflows.fib.register_atlas,
+        "_register_fib_imaging_site",
     )
 
     # Run the function and check that it's run through to completion
@@ -154,6 +174,9 @@ def test_run_with_db(
             murfey_db=murfey_db_session,
         )
     assert mock_parse.call_count == len(test_files)
+    assert spy_thumbnail.call_count == len(test_files)
+    for thumbnail in thumbnails:
+        assert thumbnail.is_file()
     assert spy_register.call_count == len(test_files)
 
     # Murfey's ImagingSite should have an entry