feat(flomni): add functionality to save reference and alignment images as HDF5 files

csaxs_bec/bec_ipython_client/plugins/flomni/flomni.py

@@ -5,6 +5,7 @@ import subprocess
 import time
 from pathlib import Path
 
+import h5py
 import numpy as np
 from bec_lib import bec_logger
 from bec_lib.alarm_handler import AlarmBase
@@ -804,11 +805,21 @@ class FlomniSampleTransferMixin:
 
         if self.OMNYTools.yesno("All OK? Continue?", "y"):
             print("OK. continue.")
+            reference_image = dev.cam_flomni_gripper.image.get()
+            self.save_reference_image(reference_image)
             dev.ftransy.controller.socket_put_confirmed("confirm=1")
         else:
             print("Stopping.")
             raise FlomniError("User abort sample transfer.")
 
+    def save_reference_image(self, image):
+        # Save the reference image to a file
+        timestamp = time.strftime("%Y%m%d_%H%M%S")
+        file = os.path.expanduser(f"~/data/raw/logs/sample_transfer_reference_image_{timestamp}.h5")
+
+        with h5py.File(file, "w") as f:
+            f.create_dataset("reference_image", data=image)
+
     def ftransfer_gripper_is_open(self) -> bool:
         status = bool(float(dev.ftransy.controller.socket_put_and_receive("MG @OUT[9]").strip()))
         return status
@@ -2599,13 +2610,10 @@ class Flomni(
         if self.tomo_type == 1:
             print("\x1b[1mTomo type 1:\x1b[0m 8 equally spaced sub-tomograms")
             print(f"Angular range                    =  {self.tomo_angle_range} degrees")
-            # N, step, total_projections all come from the same helper
-            # sub_tomo_scan() effectively uses internally - see
-            # _tomo_type1_actual_grid() for why this can't just read
-            # self.tomo_angle_stepsize directly.
-            _, achievable_step, total_projections = self._tomo_type1_actual_grid()
-            print(f"Total number of projections: {total_projections}")
-            print(f"Angular step within sub-tomogram:   {achievable_step:.3f} degrees")
+            print(
+                f"Total number of projections: {(self.tomo_angle_range/self.tomo_angle_stepsize)*8}"
+            )
+            print(f"Angular step within sub-tomogram:   {self.tomo_angle_stepsize} degrees")
             print(
                 "Angular step of the final (combined) tomogram:"
                 f"   {self.tomo_angle_range / total_projections:.3f} degrees"

csaxs_bec/bec_ipython_client/plugins/flomni/x_ray_eye_align.py

@@ -5,9 +5,10 @@ import os
 import time
 from typing import TYPE_CHECKING
 
+import h5py
 import numpy as np
-
 from bec_lib import bec_logger
+from bec_lib.endpoints import MessageEndpoints
 
 # from csaxs_bec.bec_ipython_client.plugins.cSAXS import epics_get, epics_put, fshopen, fshclose
 
@@ -42,6 +43,7 @@ class XrayEyeAlign:
         self.device_manager = client.device_manager
         self.scans = client.scans
         self.alignment_values = {}
+        self.alignment_images = []
         # Deliberately NOT calling self.flomni.reset_correction() /
         # reset_tomo_alignment_fit() here: XrayEyeAlign is constructed every
         # time Flomni() is instantiated (i.e. every new client session, not
@@ -69,6 +71,19 @@ class XrayEyeAlign:
         # fit point instead of triggering another height correction.
         self._height_centered = False
 
+    # def _fetch_alignment_image(self):
+    #     msg = self.client.connector.get_last(MessageEndpoints.device_preview("cam_xeye", "image"))
+    #     if msg is None:
+    #         return
+    #     self.alignment_images.append(msg["data"].data)
+
+    def _save_alignment_data(self, file_path: str):
+        with h5py.File(os.path.expanduser(file_path), "w") as f:
+            f.create_dataset(
+                "alignment_values", data=np.array(list(self.alignment_values.values()))
+            )
+            f.create_dataset("alignment_images", data=np.array(self.alignment_images))
+
     def update_frame(self, keep_shutter_open=False):
         if self.flomni._flomnigui_check_attribute_not_exists("xeyegui"):
             self.flomni.flomnigui_show_xeyealign()
@@ -135,6 +150,7 @@ class XrayEyeAlign:
                 "This routine can be called with paramter keep_shutter_open=True to keep the shutter always open"
             )
         self.send_message("Getting things ready. Please wait...")
+        self.alignment_images = []
 
         self.gui.enable_submit_button(False)
 
@@ -204,6 +220,7 @@ class XrayEyeAlign:
                 print(f"Current rtx position {rtx_position}")
                 self.alignment_values[k] -= rtx_position
                 print(f"Corrected position {k}: x {self.alignment_values[k]}")
+                self.alignment_images.append(dev.cam_xeye.image.get())
                 # reset submit channel
                 dev.omny_xray_gui.submit.set(0)
 
@@ -355,10 +372,14 @@ class XrayEyeAlign:
         umv(dev.rtx, 0)
         print("You are ready to remove the xray eye and start ptychography scans.")
         print("Fine alignment: flomni.tomo_parameters() , then flomni.tomo_alignment_scan()")
-        print("After that, run the fit in Matlab and load the new fit flomni.read_alignment_offset()")
+        print(
+            "After that, run the fit in Matlab and load the new fit flomni.read_alignment_offset()"
+        )
 
     def write_output(self):
         file = os.path.expanduser("~/data/raw/logs/xrayeye_alignmentvalues")
+        timestamp = time.strftime("%Y%m%d_%H%M%S")
+        self._save_alignment_data(file + f"_image_data_{timestamp}.h5")
         if not os.path.exists(file):
             os.makedirs(os.path.dirname(file), exist_ok=True)