Register lamella image

src/murfey/workflows/fib/register_lamella_evaluation_image.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import math
 import re
 from datetime import datetime
 from pathlib import Path
@@ -9,9 +10,10 @@
 from pydantic import BaseModel
 from sqlmodel import Session as SQLModelSession, select
 
+import murfey.server
 import murfey.util.db as MurfeyDB
 from murfey.util.config import get_machine_config
-from murfey.util.models import FIBImageMetadata
+from murfey.util.models import FIBImageMetadata, GridSquareParameters
 from murfey.workflows.fib.shared import (
     parse_image_metadata,
     populate_fib_imaging_site_entry,
@@ -177,6 +179,189 @@ def _register_dcg(
     murfey_db.add(imaging_site)
     murfey_db.commit()
 
+    return imaging_site
+
+
+def _register_grid_square(
+    session_id: int,
+    imaging_site: MurfeyDB.ImagingSite,
+    site_number: int,
+    murfey_db: SQLModelSession,
+):
+    """
+    Helper function to create a GridSquare entry in ISPyB if one doesn't already
+    exist, and to link it to the corresopnding ImagingSite entry.
+    """
+    # Early exits if values are missing/not configured
+    if murfey.server._transport_object is None:
+        raise RuntimeError("No TransportManager object was set up")
+    dcg_name = imaging_site.dcg_name
+    if dcg_name is None:
+        raise ValueError(
+            f"'dcg_name' field in ImagingSite entry for {imaging_site.image_path} is empty"
+        )
+
+    # Check if an atlas has been registered
+    atlas_search = murfey_db.exec(
+        select(MurfeyDB.ImagingSite)
+        .where(MurfeyDB.ImagingSite.session_id == session_id)
+        .where(MurfeyDB.ImagingSite.dcg_name == dcg_name)
+        .where(MurfeyDB.ImagingSite.data_type == "atlas")
+        .order_by(MurfeyDB.ImagingSite.id)  # Sort in ascending insertion order
+    ).all()
+    if not atlas_search:
+        logger.info(
+            f"No atlas has been registered for data collection group {dcg_name!r} yet"
+        )
+        return imaging_site
+    atlas = atlas_search[-1]
+
+    # Check if the atlas has the required values for the GridSquare registration
+    if not (
+        atlas.pos_x is not None
+        and atlas.pos_y is not None
+        and atlas.pos_z is not None
+        and atlas.rotation is not None
+        and atlas.tilt_alpha is not None
+        and atlas.len_x is not None
+        and atlas.len_y is not None
+        and atlas.thumbnail_pixels_x is not None
+        and atlas.thumbnail_pixels_y is not None
+    ):
+        logger.warning(f"Atlas {atlas.image_path} not populated with required values")
+        return imaging_site
+    atlas_x1 = atlas.pos_x + (atlas.len_x / 2)
+    atlas_y0 = atlas.pos_y - (atlas.len_y / 2)
+
+    # Check that imaging site has the required values for registration
+    if not (
+        imaging_site.pos_x is not None
+        and imaging_site.pos_y is not None
+        and imaging_site.pos_z is not None
+        and imaging_site.rotation is not None
+        and imaging_site.tilt_alpha is not None
+        and imaging_site.len_x is not None
+        and imaging_site.len_y is not None
+    ):
+        logger.warning(
+            f"ImagingSite for {imaging_site.image_path} not populated with required values"
+        )
+        return imaging_site
+
+    # Transform the imaging site coordinates into the atlas' frame of reference
+    # NOTE: This will require further investigation and tweaking, given the
+    # many axes and centres of rotation present in the FIB stage system.
+    # We start with a simple 2D rotation for now, and will adjust it as we observe
+    # the alignment accuracy
+    theta = math.radians(imaging_site.rotation - atlas.rotation)
+    sin = math.sin(theta)
+    cos = math.cos(theta)
+    x_transformed = (imaging_site.pos_x * cos) - (imaging_site.pos_y * sin)
+    y_transformed = (imaging_site.pos_x * sin) + (imaging_site.pos_y * cos)
+
+    # Find the pixel coordinates of the image on the atlas
+    # NOTE: On the atlas image, positive directions are LEFT (x) and DOWN (y)
+    x_mid_px = int(
+        round((atlas_x1 - x_transformed) / atlas.len_x * atlas.thumbnail_pixels_x) or 1
+    )
+    y_mid_px = int(
+        round((y_transformed - atlas_y0) / atlas.len_y * atlas.thumbnail_pixels_y) or 1
+    )
+
+    # Find the pixel width and height of the lamella image on the atlas
+    width_scaled = int(
+        round((imaging_site.len_x / atlas.len_x) * atlas.thumbnail_pixels_x) or 1
+    )
+    height_scaled = int(
+        round((imaging_site.len_y / atlas.len_y) * atlas.thumbnail_pixels_y) or 1
+    )
+
+    # Populate GridSquareParameters model
+    grid_square_params = GridSquareParameters(
+        tag=dcg_name,
+        x_location=x_transformed,
+        x_location_scaled=x_mid_px,
+        y_location=y_transformed,
+        y_location_scaled=y_mid_px,
+        readout_area_x=imaging_site.image_pixels_x,
+        readout_area_y=imaging_site.image_pixels_y,
+        thumbnail_size_x=imaging_site.thumbnail_pixels_x,
+        thumbnail_size_y=imaging_site.thumbnail_pixels_y,
+        width=imaging_site.image_pixels_x,
+        width_scaled=width_scaled,
+        height=imaging_site.image_pixels_y,
+        height_scaled=height_scaled,
+        x_stage_position=x_transformed,
+        y_stage_position=y_transformed,
+        pixel_size=imaging_site.image_pixel_size,
+        image=imaging_site.thumbnail_path,
+    )
+
+    # Register or update the grid square entry as required
+    if grid_square_entry := murfey_db.exec(
+        select(MurfeyDB.GridSquare)
+        .where(MurfeyDB.GridSquare.name == site_number)
+        .where(MurfeyDB.GridSquare.session_id == session_id)
+        .where(MurfeyDB.GridSquare.tag == grid_square_params.tag)
+    ).one_or_none():
+        # Update existing grid square entry on Murfey
+        grid_square_entry.x_location = grid_square_params.x_location
+        grid_square_entry.y_location = grid_square_params.y_location
+        grid_square_entry.x_stage_position = grid_square_params.x_stage_position
+        grid_square_entry.y_stage_position = grid_square_params.y_stage_position
+        grid_square_entry.readout_area_x = grid_square_params.readout_area_x
+        grid_square_entry.readout_area_y = grid_square_params.readout_area_y
+        grid_square_entry.thumbnail_size_x = grid_square_params.thumbnail_size_x
+        grid_square_entry.thumbnail_size_y = grid_square_params.thumbnail_size_y
+        grid_square_entry.pixel_size = grid_square_params.pixel_size
+        grid_square_entry.image = grid_square_params.image
+
+        # Update existing entry on ISPyB
+        murfey.server._transport_object.do_update_grid_square(
+            grid_square_id=grid_square_entry.id,
+            grid_square_parameters=grid_square_params,
+        )
+    else:
+        # Look up data collection group for current series
+        dcg_entry = murfey_db.exec(
+            select(MurfeyDB.DataCollectionGroup)
+            .where(MurfeyDB.DataCollectionGroup.session_id == session_id)
+            .where(MurfeyDB.DataCollectionGroup.tag == grid_square_params.tag)
+        ).one()
+        # Register to ISPyB
+        grid_square_ispyb_result = (
+            murfey.server._transport_object.do_insert_grid_square(
+                atlas_id=dcg_entry.atlas_id,
+                grid_square_id=site_number,
+                grid_square_parameters=grid_square_params,
+            )
+        )
+        # Create matching record in Murfey
+        grid_square_entry = MurfeyDB.GridSquare(
+            id=grid_square_ispyb_result.get("return_value", None),
+            name=site_number,
+            session_id=session_id,
+            tag=grid_square_params.tag,
+            x_location=grid_square_params.x_location,
+            y_location=grid_square_params.y_location,
+            x_stage_position=grid_square_params.x_stage_position,
+            y_stage_position=grid_square_params.y_stage_position,
+            readout_area_x=grid_square_params.readout_area_x,
+            readout_area_y=grid_square_params.readout_area_y,
+            thumbnail_size_x=grid_square_params.thumbnail_size_x,
+            thumbnail_size_y=grid_square_params.thumbnail_size_y,
+            pixel_size=grid_square_params.pixel_size,
+            image=grid_square_params.image,
+        )
+    murfey_db.add(grid_square_entry)
+
+    # Add grid square ID to existing CLEM image series entry
+    imaging_site.grid_square_id = grid_square_entry.id
+    murfey_db.add(imaging_site)
+    murfey_db.commit()
+
+    return imaging_site
+
 
 class FIBLamellaImageInfo(BaseModel):
     session_id: int
@@ -217,6 +402,10 @@ def run(
             rotation_offset=rotation_offset,
         ),
     )
+    if metadata.lamella_number is None:
+        raise ValueError(
+            f"No lamella number associated with lamella image {fib_info.lamella_image_file}"
+        )
     logger.info(
         "Extracted the following metadata from the image:\n"
         f"{json.dumps(metadata.model_dump(), indent=2, default=str)}"
@@ -231,18 +420,26 @@ def run(
 
     # 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(
+    fib_img_site = _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,
     )
 
+    # Register grid square in ISPyB
+    fib_img_site = _register_grid_square(
+        session_id=fib_info.session_id,
+        imaging_site=fib_img_site,
+        site_number=metadata.lamella_number,
+        murfey_db=murfey_db,
+    )
+
+    logger.info(
+        f"Registered lamella evaluation image {fib_info.lamella_image_file} "
+        f"for slot {metadata.slot_number} in Murfey database"
+    )
     return {"success": True}