[feat] Use CoM for FM fiducial refinement

Author: tmoerkerken

src/odemis/gui/cont/acquisition/cryo_z_localization.py

@@ -243,7 +243,7 @@ def _on_target_focus_pos(self, target_coordinates: List[float]) -> None:
         """
         # Set the target Z ctrl with the focus position
         self._panel.ctrl_target_z.SetValue(target_coordinates[2])
-        save_project(self._tab_data_model.main)
+        save_project(self._tab_data.main)
 
     def _on_ctrl_target_z_change(self) -> List[float]:
         """

src/odemis/gui/cont/multi_point_correlation.py

@@ -35,7 +35,7 @@
 # This is not related to any particular wxPython version and is most likely permanent.
 
 from odemis import model, util
-from odemis.acq.align.tdct import get_optimized_z_gauss, _convert_das_to_numpy_stack, run_tdct_correlation
+from odemis.acq.align.tdct import _convert_das_to_numpy_stack, run_tdct_correlation
 from odemis.acq.feature import FIBFMCorrelationData, Target, TargetType
 from odemis.acq.stream import StaticFluoStream, StaticSEMStream, StaticFIBStream, FluoStream
 from odemis.gui import conf
@@ -44,7 +44,7 @@
 from odemis.gui.util import call_in_wx_main
 from odemis.model import ListVA
 from odemis.util.dataio import data_to_static_streams
-from odemis.util.interpolation import interpolate_z_stack
+from odemis.util.img import get_brightest_channel, compute_center_of_mass
 from odemis.util.units import readable_str
 
 # create an enum with column labels and position
@@ -62,12 +62,14 @@ class GridColumns(Enum):
 FIDUCIAL_PATTERN = r"^[^-]+-"
 RIM_COR_DEFAULT = 0.495  # See MD_RIM_COR. This value works fine for 50x objectives, which are common
 
-# Both functions getPixel3DCoordinates(args*, kwargs*) and getPhysical3DCoordinates(args*, kwargs*) need special
+# Both functions get_pixel_3d_coordinates(args*, kwargs*) and get_physical_3d_coordinates(args*, kwargs*) need special
 # conditions to convert between physical and pixel coordinate systems in order for multipoint correlation to operate.
 # For coordinate conversions, we assume the pixels in 3D are isosymmetric
 # i.e. size in pixel[0]=pixel[1]=pixel[2].
+COM_ROI_PADDING = 12  # Padding (pixels) for center of mass ROI extraction
+# TODO: Could adapt padding based on pixel spacing for more flexibility if needed
 
-def getPixel3DCoordinates(stream: FluoStream, p_pos: Tuple[float, float, float], check_bbox: bool = False) \
+def get_pixel_3d_coordinates(stream: FluoStream, p_pos: Tuple[float, float, float], check_bbox: bool = False) \
         -> Optional[Tuple[float, float, float]]:
     """
     Translate 3D physical coordinates into 3D pixel coordinates. The z coordinate is computed assuming iso-voxel
@@ -84,20 +86,19 @@ def getPixel3DCoordinates(stream: FluoStream, p_pos: Tuple[float, float, float],
 
     raw = stream.raw[0]
     md = stream._find_metadata(raw.metadata)
-    pxs = md.get(model.MD_PIXEL_SIZE, (1e-6, 1e-6))
+    pxs = md.get(model.MD_PIXEL_SIZE, (1e-6, 1e-6, 1e-6))
     # For multipoint correlation, we assume that the pixel size in x is the same as in y
     if not util.almost_equal(pxs[0], pxs[1], atol=1e-9):
         logging.warning("Pixel size in x and y are not equal while computing pixel coordinates")
 
-    # Z position is found by taking into account MD_POS and subtracting it from the physical coordinates.
-    # Pixel value used for Z enforces the iso-voxel condition between x, y and z. It is not the real pixel value in z.
     tpos = md.get(model.MD_POS, (0, 0, 0))
     tpos_z = tpos[2] if len(tpos) >= 3 else 0.0
-    z = (p_pos[2] - tpos_z) / pxs[1]
+    z = (p_pos[2] - tpos_z) / pxs[2]
     pixel_pos = (pixel_pos[0], pixel_pos[1], z)
+
     return pixel_pos
 
-def getPhysical3DCoordinates(stream: FluoStream, pixel_pos: Tuple[float, float, float])\
+def get_physical_3d_coordinates(stream: FluoStream, pixel_pos: Tuple[float, float, float])\
                              -> Optional[Tuple[float, float, float]]:
     """
     Translate 3D pixel coordinates into 3D physical coordinates. The z coordinate is computed assuming iso-voxel
@@ -109,10 +110,11 @@ def getPhysical3DCoordinates(stream: FluoStream, pixel_pos: Tuple[float, float,
     p_pos = stream.getPhysicalCoordinates(pixel_pos[:2])
     raw = stream.raw[0]
     md = stream._find_metadata(raw.metadata)
-    pxs = md.get(model.MD_PIXEL_SIZE, (1e-6, 1e-6))[0:2]
+    pxs = md.get(model.MD_PIXEL_SIZE, (1e-6, 1e-6, 1e-6))
     tpos = md.get(model.MD_POS, (0, 0, 0))
     tpos_z = tpos[2] if len(tpos) >= 3 else 0.0
-    p_pos_z = pixel_pos[2] * pxs[1] + tpos_z
+    # Account for slice thickness, aka, z distance between slices
+    p_pos_z = pixel_pos[2] * pxs[2] + tpos_z
     return (p_pos[0], p_pos[1], p_pos_z)
 
 def update_feature_correlation_target(correlation_target: FIBFMCorrelationData,
@@ -146,8 +148,7 @@ def update_feature_correlation_target(correlation_target: FIBFMCorrelationData,
         fm_fiducials.sort(key=lambda x: x.index.value)
     correlation_target.fm_fiducials = fm_fiducials
 
-    acq_conf = conf.get_acqui_conf()
-    save_project(acq_conf.pj_last_path, tab_data.main.features.value, tab_data.main.overviews.value)
+    save_project(tab_data.main)
 
     return correlation_target
 
@@ -176,20 +177,20 @@ def __init__(self, frame):
         # Access the correlation points table (wxListCtrl)
         self.grid = self._panel.table_grid
 
-        # Access the Refine Z text (to check if refine_z is working or not)
-        self.txt_refinez_active = self._panel.txt_refinez_active
-        self.txt_refinez_active.Show(True)
+        # Access the Refine XYZ status text (to check if XYZ targeting is working or not)
+        self.txt_refine_xyz_active = self._panel.txt_refine_xyz_active
+        self.txt_refine_xyz_active.Show(True)
 
-        # Access the Z-targeting button
-        self.z_targeting_btn = self._panel.btn_z_targeting
-        self.z_targeting_btn.Bind(wx.EVT_BUTTON, self._on_z_targeting)
-        self.z_targeting_btn.Enable(False)
-        # Disable Z-targeting button if super z stream is available as Z-targeting is not required in that case
-        self.refinez_active = True
+        # Access the XYZ-targeting button
+        self.xyz_targeting_btn = self._panel.btn_xyz_targeting
+        self.xyz_targeting_btn.Bind(wx.EVT_BUTTON, self._on_xyz_targeting)
+        self.xyz_targeting_btn.Enable(False)
+        # Disable XYZ-targeting button if super z stream is available as XYZ-targeting is not required in that case
+        self.refine_xyz_active = True
         if self._tab_data_model.main.currentFeature.value.superz_stream_name:
-            self.z_targeting_btn.SetToolTip("Super Z information available, Refine Z disabled")
-            self.txt_refinez_active.SetLabel("Super Z information in use")
-            self.refinez_active = False
+            self.xyz_targeting_btn.SetToolTip("Super Z information available, Refine XYZ disabled")
+            self.txt_refine_xyz_active.SetLabel("Super Z information in use")
+            self.refine_xyz_active = False
 
         self._panel.btn_delete_row.Bind(wx.EVT_BUTTON, self._on_delete_row)
 
@@ -478,11 +479,11 @@ def _do_correlation(self):
             fib_coords.append(fib_coord)
         fib_coords = numpy.array(fib_coords, dtype=numpy.float32)
         for fm_coord in self.correlation_target.fm_fiducials:
-            fm_coord_px = getPixel3DCoordinates(self.correlation_target.fm_streams[0], fm_coord.coordinates.value)
+            fm_coord_px = get_pixel_3d_coordinates(self.correlation_target.fm_streams[0], fm_coord.coordinates.value)
             fm_coords.append(fm_coord_px)
         fm_coords = numpy.array(fm_coords, dtype=numpy.float32)
         poi_coord = self.correlation_target.fm_pois[0]
-        poi_coord_px = getPixel3DCoordinates(self.correlation_target.fm_streams[0], poi_coord.coordinates.value)
+        poi_coord_px = get_pixel_3d_coordinates(self.correlation_target.fm_streams[0], poi_coord.coordinates.value)
         poi_coords.append(poi_coord_px)
         poi_coords = numpy.array(poi_coords, dtype=numpy.float32)
         # Run the correlation
@@ -663,7 +664,7 @@ def _on_cell_changing(self, event) -> None:
                 elif col_name == GridColumns.Z.name and (
                         self._tab_data_model.main.currentTarget.value.type.value != TargetType.FibFiducial):
                     self._tab_data_model.main.currentTarget.value.coordinates.value[2] = \
-                    getPhysical3DCoordinates(self.correlation_target.fm_streams[0], (x, y, float(new_value)))[2]
+                    get_physical_3d_coordinates(self.correlation_target.fm_streams[0], (x, y, float(new_value)))[2]
             except ValueError:
                 wx.MessageBox("X, Y, Z values must be a float!", "Invalid Input", wx.OK | wx.ICON_ERROR)
                 event.Veto()  # Prevent the change
@@ -695,19 +696,19 @@ def _on_current_target_changes(self, target: Target) -> None:
         # For new targets, automatically perform Z targeting if MIP is checked for at least one FM stream
         if not target:
             self.grid.ClearSelection()
-            self.z_targeting_btn.Enable(False)
+            self.xyz_targeting_btn.Enable(False)
             return
 
         mip_enabled = any([stream.max_projection.value for stream in self.correlation_target.fm_streams])
 
-        # Refine z should be disabled if the the Z information was obtained using SuperZ
-        if self.refinez_active and (target.type.value in self.grid_targets):
+        # Refine xyz should be disabled if the the Z information was obtained using SuperZ
+        if self.refine_xyz_active and (target.type.value in self.grid_targets):
             if TargetType.FibFiducial == target.type.value:
-                self.z_targeting_btn.Enable(False)
+                self.xyz_targeting_btn.Enable(False)
             else:
-                self.z_targeting_btn.Enable(True)
+                self.xyz_targeting_btn.Enable(True)
                 if mip_enabled:
-                    self._on_z_targeting(None)
+                    self._on_xyz_targeting(None)
 
         for row in range(self.grid.GetNumberRows()):
             if self._selected_target_in_grid(target, row):
@@ -739,7 +740,7 @@ def _on_current_coordinates_changes(self, coordinates: ListVA) -> None:
                     pixel_coords = self.correlation_target.fib_stream.getPixelCoordinates(
                         (target.coordinates.value[0], target.coordinates.value[1]), check_bbox=False)
                 else:
-                    pixel_coords = getPixel3DCoordinates(self.correlation_target.fm_streams[0], target.coordinates.value)
+                    pixel_coords = get_pixel_3d_coordinates(self.correlation_target.fm_streams[0], target.coordinates.value)
                     if (self.grid.GetCellValue(row,
                                                GridColumns.Z.value)) != f"{pixel_coords[2]:.{GRID_PRECISION}f}":
                         temp_check = True
@@ -783,7 +784,7 @@ def _on_target_changes(self, targets: List[Target]) -> None:
                         (target.coordinates.value[0], target.coordinates.value[1]), check_bbox=False)
                     self.grid.SetCellValue(current_row_count, GridColumns.Z.value, "")
                 else:
-                    pixel_coords = getPixel3DCoordinates(self.correlation_target.fm_streams[0], target.coordinates.value)
+                    pixel_coords = get_pixel_3d_coordinates(self.correlation_target.fm_streams[0], target.coordinates.value)
                     self.grid.SetCellValue(current_row_count, GridColumns.Z.value,
                                            f"{pixel_coords[2]:.{GRID_PRECISION}f}")
                 # Set x and y position in the grid
@@ -807,34 +808,58 @@ def _on_target_changes(self, targets: List[Target]) -> None:
         if self.check_correlation_conditions():
             self._need_reprocessing()
 
-    def _on_z_targeting(self, evt) -> None:
+    def _on_xyz_targeting(self, evt) -> None:
         """
-        Handle Z-targeting when the Z-targeting button is clicked.
+        Handle targeting when the targeting button is clicked.
+        Refactored to perform full 3D Center of Mass targeting (X, Y, Z).
         """
         if self._tab_data_model.main.currentTarget.value:
 
-            # Select the streams which are visible in the view for Z-targeting
+            # Select the streams which are visible in the view for targeting
             streams_projections = self._tab_data_model.views.value[0].stream_tree.flat.value
             if not streams_projections:
-                wx.MessageBox("FM streams are not available for refining Z", "Error", wx.OK | wx.ICON_ERROR)
+                wx.MessageBox("FM streams are not available for refining targets", "Error", wx.OK | wx.ICON_ERROR)
                 return
 
-            self.txt_refinez_active.SetLabel("active ...")
-            wx.CallLater(1000, self.txt_refinez_active.SetLabel, "")
+            self.txt_refine_xyz_active.SetLabel("active ...")
+            wx.CallLater(1000, self.txt_refine_xyz_active.SetLabel, "")
 
             coords = self._tab_data_model.main.currentTarget.value.coordinates.value
-            pixel_coords = getPixel3DCoordinates(self.correlation_target.fm_streams[0], coords)
-            das = [interpolate_z_stack(da=stream_projection.stream.raw[0]
-                                       [:,
-                                       int(pixel_coords[1]):int(pixel_coords[1])+1,
-                                       int(pixel_coords[0]):int(pixel_coords[0])+1],
-                                       method="linear")
-                   for stream_projection in streams_projections]
-
-            z = float(get_optimized_z_gauss(das, int(0), int(0), int(pixel_coords[2])))
-            z_p = getPhysical3DCoordinates(self.correlation_target.fm_streams[0],
-                                 (pixel_coords[0],pixel_coords[1], z))[2]
-            self._tab_data_model.main.currentTarget.value.coordinates.value[2] = z_p
+            pixel_coords = get_pixel_3d_coordinates(self.correlation_target.fm_streams[0], coords)
+
+            # We are going to refine around the clicked position
+            target_x, target_y = int(pixel_coords[0]), int(pixel_coords[1])
+            # Ensure multi-channel compatibility
+            raw_multi = numpy.asarray([d.stream.raw[0] for d in streams_projections])
+            shape_y, shape_x = raw_multi.shape[-2], raw_multi.shape[-1]
+            # Get boundary-safe slice & crop
+            y_start = max(0, target_y - COM_ROI_PADDING)
+            y_end = min(shape_y, target_y + COM_ROI_PADDING + 1)
+            x_start = max(0, target_x - COM_ROI_PADDING)
+            x_end = min(shape_x, target_x + COM_ROI_PADDING + 1)
+            roi = numpy.s_[:, y_start:y_end, x_start:x_end]
+            multi_crop = raw_multi[(slice(None),) + roi]  # We search along all stack slices (first axis)
+            # Find best channel and compute COM
+            best_c = get_brightest_channel(multi_crop)
+            com = compute_center_of_mass(multi_crop[best_c], baseline_percentile=95.0)
+            com_z = com[0]
+            com_y_crop = com[1] + roi[1].start
+            com_x_crop = com[2] + roi[2].start
+
+            # Map back to physical coordinates using optimized X, Y, and Z
+            physical_coords = get_physical_3d_coordinates(
+                self.correlation_target.fm_streams[0],
+                (com_x_crop, com_y_crop, com_z)
+            )
+
+            # Update the model with the refined 3D coordinates
+            target_coords = self._tab_data_model.main.currentTarget.value.coordinates.value
+            target_coords[0] = physical_coords[0]
+            target_coords[1] = physical_coords[1]
+            target_coords[2] = physical_coords[2]
+
+            for vp in self._viewports:
+                vp.canvas.update_drawing()
 
     def _reorder_grid(self) -> None:
         """

src/odemis/gui/main_xrc.py

@@ -118,8 +118,8 @@ def __init__(self, parent):
         self.fp_correlation_panel = xrc.XRCCTRL(self, "fp_correlation_panel")
         self.pnl_correlation = xrc.XRCCTRL(self, "pnl_correlation")
         self.btn_delete_row = xrc.XRCCTRL(self, "btn_delete_row")
-        self.btn_z_targeting = xrc.XRCCTRL(self, "btn_z_targeting")
-        self.txt_refinez_active = xrc.XRCCTRL(self, "txt_refinez_active")
+        self.btn_xyz_targeting = xrc.XRCCTRL(self, "btn_xyz_targeting")
+        self.txt_refine_xyz_active = xrc.XRCCTRL(self, "txt_refine_xyz_active")
         self.table_grid = xrc.XRCCTRL(self, "table_grid")
         self.txt_correlation_rms = xrc.XRCCTRL(self, "txt_correlation_rms")
         self.fp_correlation_streams = xrc.XRCCTRL(self, "fp_correlation_streams")
@@ -2010,14 +2010,14 @@ def __init_resources():
                                   <flag>wxALL|wxEXPAND</flag>
                                   <border>10</border>
                                 </object>
-                                <!-- Z-targeting button -->
+                                <!-- XYZ-targeting button -->
                                 <object class="sizeritem">
-                                  <object class="wxButton" name="btn_z_targeting">
-                                    <label>Refine Z</label>
+                                  <object class="wxButton" name="btn_xyz_targeting">
+                                    <label>Refine</label>
                                   </object>
                                 </object>
                                 <object class="sizeritem">
-                                  <object class="wxStaticText" name="txt_refinez_active">
+                                  <object class="wxStaticText" name="txt_refine_xyz_active">
                                   <label> </label>
                                   <fg>#E5E5E5</fg>
                                   <hidden>1</hidden>

src/odemis/gui/xmlh/resources/dialog_correlation_tdct.xrc

@@ -86,14 +86,14 @@
                                   <flag>wxALL|wxEXPAND</flag>
                                   <border>10</border>
                                 </object>
-                                <!-- Z-targeting button -->
+                                <!-- XYZ-targeting button -->
                                 <object class="sizeritem">
-                                  <object class="wxButton" name="btn_z_targeting">
-                                    <label>Refine Z</label>
+                                  <object class="wxButton" name="btn_xyz_targeting">
+                                    <label>Refine</label>
                                   </object>
                                 </object>
                                 <object class="sizeritem">
-                                  <object class="wxStaticText" name="txt_refinez_active">
+                                  <object class="wxStaticText" name="txt_refine_xyz_active">
                                   <label> </label>
                                   <fg>#E5E5E5</fg>
                                   <hidden>1</hidden>