Make ISPyB records for atlas, grid squares and foil holes (#426)

This is complicated and almost certainly buggy. The problems arise because we need to allow for grid square and foil hole updates when metadata becomes available.

Author: d-j-hatton

src/murfey/client/contexts/spa.py

@@ -41,6 +41,7 @@ class FoilHole(NamedTuple):
     thumbnail_size_y: Optional[int] = None
     pixel_size: Optional[float] = None
     image: str = ""
+    diameter: Optional[float] = None
 
 
 class GridSquare(NamedTuple):
@@ -59,9 +60,18 @@ class GridSquare(NamedTuple):
     tag: str = ""
 
 
-def _get_grid_square_atlas_positions(
-    xml_path: Path, grid_square: str = ""
-) -> Dict[str, Tuple[Optional[int], Optional[int], Optional[float], Optional[float]]]:
+def _get_grid_square_atlas_positions(xml_path: Path, grid_square: str = "") -> Dict[
+    str,
+    Tuple[
+        Optional[int],
+        Optional[int],
+        Optional[float],
+        Optional[float],
+        Optional[int],
+        Optional[int],
+        Optional[float],
+    ],
+]:
     with open(
         xml_path,
         "r",
@@ -71,7 +81,16 @@ def _get_grid_square_atlas_positions(
         "TileXml"
     ]
     gs_pix_positions: Dict[
-        str, Tuple[Optional[int], Optional[int], Optional[float], Optional[float]]
+        str,
+        Tuple[
+            Optional[int],
+            Optional[int],
+            Optional[float],
+            Optional[float],
+            Optional[int],
+            Optional[int],
+            Optional[float],
+        ],
     ] = {}
     for ti in tile_info:
         try:
@@ -96,6 +115,13 @@ def _get_grid_square_atlas_positions(
                     * 1e9,
                     float(gs["value"]["b:PositionOnTheAtlas"]["c:Physical"]["d:y"])
                     * 1e9,
+                    int(
+                        float(gs["value"]["b:PositionOnTheAtlas"]["c:Size"]["d:width"])
+                    ),
+                    int(
+                        float(gs["value"]["b:PositionOnTheAtlas"]["c:Size"]["d:height"])
+                    ),
+                    float(gs["value"]["b:PositionOnTheAtlas"]["c:Rotation"]),
                 )
                 if grid_square:
                     break
@@ -221,6 +247,7 @@ def _foil_hole_data(
         for fh_block in serialization_array[required_key]:
             pix = fh_block["b:value"]["PixelCenter"]
             stage = fh_block["b:value"]["StagePosition"]
+            diameter = fh_block["b:value"]["PixelWidthHeight"]["c:width"]
             if int(fh_block["b:key"]) == foil_hole:
                 return FoilHole(
                     id=foil_hole,
@@ -236,6 +263,7 @@ def _foil_hole_data(
                     thumbnail_size_y=None,
                     pixel_size=float(pixel_size) if image_path else None,
                     image=str(image_path),
+                    diameter=diameter,
                 )
     logger.warning(
         f"Foil hole positions could not be determined from metadata file {xml_path} for foil hole {foil_hole}"
@@ -557,7 +585,10 @@ def _position_analysis(
                 Optional[int],
                 Optional[float],
                 Optional[float],
-            ] = (None, None, None, None)
+                Optional[int],
+                Optional[int],
+                Optional[float],
+            ] = (None, None, None, None, None, None, None)
             data_collection_group = (
                 requests.get(
                     f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/data_collection_groups"
@@ -611,6 +642,9 @@ def _position_analysis(
                     "y_location": gs_pix_position[1],
                     "x_stage_position": gs_pix_position[2],
                     "y_stage_position": gs_pix_position[3],
+                    "width": gs_pix_position[4],
+                    "height": gs_pix_position[5],
+                    "angle": gs_pix_position[6],
                 },
             )
         foil_hole = _foil_hole_from_file(transferred_file)
@@ -651,6 +685,7 @@ def _position_analysis(
                         "thumbnail_size_x": fh.thumbnail_size_x,
                         "thumbnail_size_y": fh.thumbnail_size_y,
                         "pixel_size": fh.pixel_size,
+                        "diameter": fh.diameter,
                         "tag": str(source),
                         "image": str(image_path),
                     },

src/murfey/client/contexts/spa_metadata.py

@@ -4,6 +4,7 @@
 
 import requests
 import xmltodict
+from PIL import Image
 
 from murfey.client.context import Context
 from murfey.client.contexts.spa import _get_grid_square_atlas_positions, _get_source
@@ -64,6 +65,38 @@ def post_transfer(
             visitless_path = Path(
                 str(transferred_file).replace(f"/{environment.visit}", "")
             )
+            visit_index_of_transferred_file = transferred_file.parts.index(
+                environment.visit
+            )
+            atlas_xml_path = list(
+                (
+                    Path(
+                        "/".join(
+                            transferred_file.parts[
+                                : visit_index_of_transferred_file + 1
+                            ]
+                        )
+                    )
+                    / partial_path
+                ).parent.glob("Atlas_*.xml")
+            )[0]
+            with open(atlas_xml_path, "rb") as atlas_xml:
+                atlas_xml_data = xmltodict.parse(atlas_xml)
+                atlas_original_pixel_size = atlas_xml_data["MicroscopeImage"][
+                    "SpatialScale"
+                ]["pixelSize"]["x"]["numericValue"]
+                readout_width = float(
+                    atlas_xml_data["MicroscopeImage"]["SpatialScale"]["pixelSize"]["x"][
+                        "numericValue"
+                    ]
+                )
+
+            # need to calculate the pixel size of the downscaled image
+            atlas_im = Image.open(atlas_xml_path.with_suffix(".jpg"))
+            atlas_pixel_size = atlas_original_pixel_size * (
+                readout_width / atlas_im.width
+            )
+
             source = _get_source(
                 visitless_path.parent / "Images-Disc1" / visitless_path.name,
                 environment,
@@ -90,6 +123,7 @@ def post_transfer(
                         / environment.samples[source].atlas
                     ),
                     "sample": environment.samples[source].sample,
+                    "atlas_pixel_size": atlas_pixel_size,
                 }
                 capture_post(url, json=dcg_data)
                 registered_grid_squares = (
@@ -121,5 +155,8 @@ def post_transfer(
                                     "y_location": pos_data[1],
                                     "x_stage_position": pos_data[2],
                                     "y_stage_position": pos_data[3],
+                                    "width": pos_data[4],
+                                    "height": pos_data[5],
+                                    "angle": pos_data[6],
                                 },
                             )

src/murfey/server/__init__.py

@@ -23,6 +23,7 @@
 from fastapi.templating import Jinja2Templates
 from importlib_metadata import EntryPoint  # For type hinting only
 from ispyb.sqlalchemy._auto_db_schema import (
+    Atlas,
     AutoProcProgram,
     Base,
     DataCollection,
@@ -43,6 +44,7 @@
 from werkzeug.utils import secure_filename
 
 import murfey
+import murfey.server.ispyb
 import murfey.server.prometheus as prom
 import murfey.server.websocket
 import murfey.util.db as db
@@ -2629,8 +2631,17 @@ def feedback_callback(header: dict, message: dict) -> None:
                         experimentTypeId=message["experiment_type_id"],
                     )
                     dcgid = _register(record, header)
+                    atlas_record = Atlas(
+                        dataCollectionGroupId=dcgid,
+                        atlasImage=message.get("atlas", ""),
+                        pixelSize=message.get("atlas_pixel_size", 0),
+                        cassetteSlot=message.get("sample"),
+                    )
+                    if _transport_object:
+                        atlas_id = _transport_object.do_insert_atlas(atlas_record)
                     murfey_dcg = db.DataCollectionGroup(
                         id=dcgid,
+                        atlas_id=atlas_id,
                         session_id=message["session_id"],
                         tag=message.get("tag"),
                     )
@@ -2655,6 +2666,14 @@ def feedback_callback(header: dict, message: dict) -> None:
                     }
                 _transport_object.transport.ack(header)
             return None
+        elif message["register"] == "atlas_update":
+            if _transport_object:
+                _transport_object.do_update_atlas(
+                    message["atlas_id"],
+                    message["atlas"],
+                    message["atlas_pixel_size"],
+                    message["sample"],
+                )
         elif message["register"] == "data_collection":
             murfey_session_id = message["session_id"]
             ispyb_session_id = murfey.server.ispyb.get_session_id(

src/murfey/server/api/__init__.py

@@ -502,7 +502,21 @@ def register_grid_square(
         grid_square.y_location = grid_square_params.y_location
         grid_square.x_stage_position = grid_square_params.x_stage_position
         grid_square.y_stage_position = grid_square_params.y_stage_position
+        if _transport_object:
+            _transport_object.do_update_grid_square(grid_square.id, grid_square_params)
     except Exception:
+        if _transport_object:
+            dcg = db.exec(
+                select(DataCollectionGroup)
+                .where(DataCollectionGroup.session_id == session_id)
+                .where(DataCollectionGroup.tag == grid_square_params.tag)
+            ).one()
+            gs_ispyb_response = _transport_object.do_insert_grid_square(
+                dcg.atlas_id, gsid, grid_square_params
+            )
+        else:
+            # mock up response so that below still works
+            gs_ispyb_response = {"success": False, "return_value": None}
         secured_grid_square_image_path = secure_filename(grid_square_params.image)
         if (
             secured_grid_square_image_path
@@ -512,6 +526,11 @@ def register_grid_square(
         else:
             jpeg_size = (0, 0)
         grid_square = GridSquare(
+            id=(
+                gs_ispyb_response["return_value"]
+                if gs_ispyb_response["success"]
+                else None
+            ),
             name=gsid,
             session_id=session_id,
             tag=grid_square_params.tag,
@@ -552,16 +571,13 @@ def register_foil_hole(
     db=murfey_db,
 ):
     try:
-        gsid = (
-            db.exec(
-                select(GridSquare)
-                .where(GridSquare.tag == foil_hole_params.tag)
-                .where(GridSquare.session_id == session_id)
-                .where(GridSquare.name == gs_name)
-            )
-            .one()
-            .id
-        )
+        gs = db.exec(
+            select(GridSquare)
+            .where(GridSquare.tag == foil_hole_params.tag)
+            .where(GridSquare.session_id == session_id)
+            .where(GridSquare.name == gs_name)
+        ).one()
+        gsid = gs.id
     except NoResultFound:
         log.debug(
             f"Foil hole {sanitise(str(foil_hole_params.name))} could not be registered as grid square {sanitise(str(gs_name))} was not found"
@@ -572,21 +588,53 @@ def register_foil_hole(
         jpeg_size = Image.open(secured_foil_hole_image_path).size
     else:
         jpeg_size = (0, 0)
-    foil_hole = FoilHole(
-        name=foil_hole_params.name,
-        session_id=session_id,
-        grid_square_id=gsid,
-        x_location=foil_hole_params.x_location,
-        y_location=foil_hole_params.y_location,
-        x_stage_position=foil_hole_params.x_stage_position,
-        y_stage_position=foil_hole_params.y_stage_position,
-        readout_area_x=foil_hole_params.readout_area_x,
-        readout_area_y=foil_hole_params.readout_area_y,
-        thumbnail_size_x=foil_hole_params.thumbnail_size_x or jpeg_size[0],
-        thumbnail_size_y=foil_hole_params.thumbnail_size_y or jpeg_size[1],
-        pixel_size=foil_hole_params.pixel_size,
-        image=secured_foil_hole_image_path,
-    )
+    try:
+        foil_hole = db.exec(
+            select(FoilHole)
+            .where(FoilHole.name == foil_hole_params.name)
+            .where(FoilHole.grid_square_id == gsid)
+            .where(FoilHole.session_id == session_id)
+        ).one()
+        foil_hole.x_location = foil_hole_params.x_location
+        foil_hole.y_location = foil_hole_params.y_location
+        foil_hole.x_stage_position = foil_hole_params.x_stage_position
+        foil_hole.y_stage_position = foil_hole_params.y_stage_position
+        foil_hole.readout_area_x = foil_hole_params.readout_area_x
+        foil_hole.readout_area_y = foil_hole_params.readout_area_y
+        foil_hole.thumbnail_size_x = foil_hole_params.thumbnail_size_x or jpeg_size[0]
+        foil_hole.thumbnail_size_y = foil_hole_params.thumbnail_size_y or jpeg_size[1]
+        foil_hole.pixel_size = foil_hole_params.pixel_size
+        if _transport_object:
+            _transport_object.do_update_foil_hole(
+                foil_hole.id, gs.thumbnail_size_x / gs.readout_area_x, foil_hole_params
+            )
+    except Exception:
+        if _transport_object:
+            fh_ispyb_response = _transport_object.do_insert_foil_hole(
+                gsid.id, gs.thumbnail_size_x / gs.readout_area_x, foil_hole_params
+            )
+        else:
+            fh_ispyb_response = {"success": False, "return_value": None}
+        foil_hole = FoilHole(
+            id=(
+                fh_ispyb_response["return_value"]
+                if fh_ispyb_response["success"]
+                else None
+            ),
+            name=foil_hole_params.name,
+            session_id=session_id,
+            grid_square_id=gsid,
+            x_location=foil_hole_params.x_location,
+            y_location=foil_hole_params.y_location,
+            x_stage_position=foil_hole_params.x_stage_position,
+            y_stage_position=foil_hole_params.y_stage_position,
+            readout_area_x=foil_hole_params.readout_area_x,
+            readout_area_y=foil_hole_params.readout_area_y,
+            thumbnail_size_x=foil_hole_params.thumbnail_size_x or jpeg_size[0],
+            thumbnail_size_y=foil_hole_params.thumbnail_size_y or jpeg_size[1],
+            pixel_size=foil_hole_params.pixel_size,
+            image=secured_foil_hole_image_path,
+        )
     db.add(foil_hole)
     db.commit()
     db.close()
@@ -1128,6 +1176,7 @@ async def request_spa_preprocessing(
                 "image_number": proc_file.image_number,
                 "microscope": get_microscope(),
                 "mc_uuid": murfey_ids[0],
+                "foil_hole_id": foil_hole_id,
                 "ft_bin": proc_params["motion_corr_binning"],
                 "fm_dose": proc_params["dose_per_frame"],
                 "gain_ref": proc_params["gain_ref"],
@@ -1377,15 +1426,30 @@ def register_dc_group(
     ).all():
         dcg_murfey[0].atlas = dcg_params.atlas
         dcg_murfey[0].sample = dcg_params.sample
+        dcg_murfey[0].atlas_pixel_size = dcg_params.atlas_pixel_size
         db.add(dcg_murfey[0])
         db.commit()
+        if _transport_object:
+            _transport_object.send(
+                _transport_object.feedback_queue,
+                {
+                    "register": "atlas_update",
+                    "atlas_id": dcg_murfey.atlas_id,
+                    "atlas": dcg_params.atlas,
+                    "sample": dcg_params.sample,
+                    "atlas_pixel_size": dcg_params.atlas_pixel_size,
+                },
+            )
     else:
         dcg_parameters = {
             "start_time": str(datetime.datetime.now()),
             "experiment_type": dcg_params.experiment_type,
             "experiment_type_id": dcg_params.experiment_type_id,
             "tag": dcg_params.tag,
             "session_id": session_id,
+            "atlas": dcg_params.atlas,
+            "sample": dcg_params.sample,
+            "atlas_pixel_size": dcg_params.atlas_pixel_size,
         }
 
         if _transport_object: