python-microscopy · Yujin-Bao · Mar 25, 2024 · Mar 25, 2024 · Mar 26, 2024 · Mar 26, 2024
diff --git a/PYME/DSView/modules/drift_correct.py b/PYME/DSView/modules/drift_correct.py
@@ -0,0 +1,49 @@
+import numpy as np
+from ._base import Plugin
+
+class DriftCorrector(Plugin):
+    def __init__(self, dsviewer):
+        Plugin.__init__(self, dsviewer)
+
+        dsviewer.AddMenuItem('Processing', 'Sub-pixel Drift Correction', self.OnSubPixelDriftCorrect)
+
+
+    def OnSubPixelDriftCorrect(self, event):
+        """sub-pixel lateral drift correction based on the drift tracking events
+        in the transmitted-light channel
+        """
+        #from scipy import ndimage
+        #from PYME.IO.DataSources import ArrayDataSource
+        #from PYME.IO.DataSources import BaseDataSource
+        from PYME.IO.image import ImageStack
+        from PYME.DSView import ViewIm3D
+        from PYME.LMVis import pipeline
+        from PYME.IO.DataSources import DriftCorrectDataSource
+
+        # read lateral drift values from events
+        ev_mappings, _ = pipeline._processEvents(self.image.data_xyztc, self.image.events, self.image.mdh)
+        driftx = ev_mappings['driftx']
+        drifty = ev_mappings['drifty']
+        #dx = driftx(np.arange(1, self.image.data_xyztc.shape[3]+1))   # in pixel unit
+        #dy = drifty(np.arange(1, self.image.data_xyztc.shape[3]+1))   # in pixel unit
+
+        """
+        # correct lateral drift
+        drift_corrected = np.ones_like(self.image.data_xyztc[:,:,:,:,:].squeeze())
+        for t in range(self.image.data_xyztc.shape[3]):
+            corrected = ndimage.shift(self.image.data_xyztc.getSlice(t), shift=[-dx[t], -dy[t]], order=3, mode='nearest')
+            drift_corrected[:,:,t] *= corrected
+
+        im = BaseDataSource.XYZTCWrapper(ArrayDataSource.ArrayDataSource(drift_corrected), 'XYZTC', self.image.data_xyztc.shape[2], drift_corrected.shape[2], 1)
+        im = ImageStack(im)
+        im.mdh.copyEntriesFrom(self.image.mdh)
+        """
+
+        ds = DriftCorrectDataSource.XYZTCDriftCorrectSource(self.image.data_xyztc, driftx, drifty, x_scale=1.0, y_scale=1.0)
+        im = ImageStack(ds[:,:,:,:,:], mdh=ds.mdh)
+        ViewIm3D(im, mode=self.dsviewer.mode, glCanvas=self.dsviewer.glCanvas)
+
+
+def Plug(dsviewer):
+    return DriftCorrector(dsviewer)
+
diff --git a/PYME/IO/DataSources/DriftCorrectDataSource.py b/PYME/IO/DataSources/DriftCorrectDataSource.py
@@ -0,0 +1,64 @@
+from .BaseDataSource import XYZTCDataSource, XYZTCWrapper
+import numpy as np
+from scipy import ndimage
+
+class XYZTCDriftCorrectSource(XYZTCDataSource):
+    moduleName = 'DriftCorrectDataSource'
+    '''
+    def __init__(self, datasource, x_mapping, y_mapping, x_scale=1.0, y_scale=1.0):
+        if (not isinstance(datasource, XYZTCDataSource)) and (not datasource.ndim == 5) :
+            datasource = XYZTCWrapper.auto_promote(datasource)
+
+        self._datasource = datasource
+        size_z, size_t, size_c = datasource.shape[2:]
+
+        self._x_map = x_mapping # a piecewise mapping object
+        self._x_scale = x_scale #allows conversion between different camera pixel units. TODO - make it an affine transformation matrix to allow for rotation as well.
+        self._y_map = y_mapping
+        self._y_scale = y_scale
+
+        XYZTCDataSource.__init__(self, input_order=datasource._input_order, size_z=size_z, size_t=size_t, size_c=size_c)
+    '''
+    def __init__(self, datasource, x_mapping, y_mapping, px0, py0, px1, py1, relative_rotation_angle):
+        if (not isinstance(datasource, XYZTCDataSource)) and (not datasource.ndim == 5) :
+            datasource = XYZTCWrapper.auto_promote(datasource)
+
+        self._datasource = datasource
+        size_z, size_t, size_c = datasource.shape[2:]
+
+        self._x_map = x_mapping # a piecewise mapping object
+        self._y_map = y_mapping
+        self._xx_scale = px0/px1
+        self._xy_scale = px0/py1
+        self._yx_scale = py0/px1
+        self._yy_scale = py0/py1
+        self._theta = relative_rotation_angle     # prepare for an affine transformation to allow for scaling and rotation
+
+        XYZTCDataSource.__init__(self, input_order=datasource._input_order, size_z=size_z, size_t=size_t, size_c=size_c)
+
+    def getSlice(self, ind):
+        sl = self._datasource.getSlice(ind)
+        x0 = self._x_map(np.array([ind+1]))
+        y0 = self._y_map(np.array([ind+1]))
+        x1 = x0*self._xx_scale*np.cos(self._theta) + y0*self._yx_scale*np.sin(self._theta)
+        y1 = y0*self._yy_scale*np.cos(self._theta) - x0*self._xy_scale*np.sin(self._theta)
+        #return ndimage.shift(sl, [-self._x_scale*self._x_map(np.array([ind+1])), -self._y_scale*self._y_map(np.array([ind+1]))], order=3, mode='nearest') 
+        return ndimage.shift(sl, [-x1, -y1], order=3, mode='nearest') 
+
+   # proxy original data source attributes 
+    def __getattr__(self, item):
+        return getattr(self._datasource, item)
+
+    def getSliceShape(self):
+        return self._datasource.getSliceShape()
+
+    def getNumSlices(self):
+        return self._datasource.getNumSlices()
+
+    def getEvents(self):
+        return self._datasource.getEvents()
+
+    @property
+    def is_complete(self):
+        return self._datasource.is_complete()
+
diff --git a/PYME/LMVis/pipeline.py b/PYME/LMVis/pipeline.py
@@ -273,6 +273,23 @@ def _processEvents(ds, events, mdh):
 
             # self.eventCharts = eventCharts
             # self.ev_mappings = ev_mappings
+
+        if b'PYME2ShiftMeasure' in evKeyNames:
-        if b'PYME2ShiftMeasure' in evKeyNames:
+        elif b'PYME2ShiftMeasure' in evKeyNames:
+            # Standard ShiftMeasure events should have priority (i.e. this should only execute if we haven't found any standard ShiftMeasure events
+            # TODO - find a way to remove the need for this special case / event type.
-        if b'PYME2ShiftMeasure' in evKeyNames:
+        elif b'PYME2ShiftMeasure' in evKeyNames:
+            # Standard ShiftMeasure events should have priority (i.e. this should only execute if we haven't found any standard ShiftMeasure events
+            # TODO - find a way to remove the need for this special case / event type.
+            driftx = piecewiseMapping.GeneratePMFromEventList(events, mdh, mdh.getEntry('StartTime'), 0, b'PYME2ShiftMeasure',
+                                                              0)
+            drifty = piecewiseMapping.GeneratePMFromEventList(events, mdh, mdh.getEntry('StartTime'), 0, b'PYME2ShiftMeasure',
+                                                              1)
+            driftz = piecewiseMapping.GeneratePMFromEventList(events, mdh, mdh.getEntry('StartTime'), 0, b'PYME2ShiftMeasure',
+                                                              2)
+
+            ev_mappings['driftx'] = driftx
+            ev_mappings['drifty'] = drifty
+            ev_mappings['driftz'] = driftz
+
+            eventCharts.append(('X Drift [px]', driftx, 'PYME2ShiftMeasure'))
+            eventCharts.append(('Y Drift [px]', drifty, 'PYME2ShiftMeasure'))
+            eventCharts.append(('Z Drift [px]', driftz, 'PYME2ShiftMeasure'))
+
     elif all(k in mdh.keys() for k in ['StackSettings.FramesPerStep', 'StackSettings.StepSize',
                                        'StackSettings.NumSteps', 'StackSettings.NumCycles']):
         # TODO - Remove this code - anytime we get here it's generally the result of an error in the input data

diff --git a/PYME/recipes/processing.py b/PYME/recipes/processing.py
@@ -2905,4 +2905,59 @@ def applyFilter(self, data, chanNum, i, image0):
         if self.offset_selection == 'offset by a constant':
             return data - self.offset_constant
         elif self.offset_selection == 'offset by minimum':
-            return data - np.min(data)
+            return data - np.min(data)
+
+
+@register_module('SubPixelDriftCorrect')
+class SubPixelDriftCorrect(ModuleBase):
+    """
+    Sub-pixel lateral drift correction based on the drift tracking events
+    in the transmitted-light channel
+
+    Parameters
+    ----------
+    input : Input
+        PYME.IO.ImageStack
+    drift_tracking_channel_pixel_size_x : Float
+        Pixel value of the drift tracking channel in x, in nm unit
+    drift_tracking_channel_pixel_size_y : Float
+        Pixel value of the drift tracking channel in y, in nm unit
+    relative_rotation : Float
+        relative rotation degree from the drift tracking channel to the channel to be corrected, in degree
+
+    Returns
+    -------
+    output = Output
+        PYME.IO.ImageStack
+
+    """
+
+    input = Input('input')
+    drift_tracking_channel_pixel_size_x = Float(64.5)
+    drift_tracking_channel_pixel_size_y = Float(64.5)
+    relative_rotation = Float(0.0)
+    output = Output('sub_pixel_drift_corrected')
+
+    def run(self, input):
+        from PYME.IO.image import ImageStack
+        #from PYME.DSView import ViewIm3D
+        from PYME.LMVis import pipeline
+        from PYME.IO.DataSources import DriftCorrectDataSource
+
+        # read lateral drift values from events
+        ev_mappings, _ = pipeline._processEvents(input.data_xyztc, input.events, input.mdh)
+        driftx = ev_mappings['driftx']
+        drifty = ev_mappings['drifty']
+
+        px0 = self.drift_tracking_channel_pixel_size_x
+        py0 = self.drift_tracking_channel_pixel_size_y
+        px1 = input.mdh.voxelsize.x * 1e3
+        py1 = input.mdh.voxelsize.y * 1e3
+        theta = self.relative_rotation * np.pi/180
+
+        ds = DriftCorrectDataSource.XYZTCDriftCorrectSource(input.data_xyztc, driftx, drifty, px0, py0, px1, py1, theta)
+        im = ImageStack(ds[:,:,:,:,:], mdh=input.mdh)
+        #ViewIm3D(im, mode=self.dsviewer.mode, glCanvas=self.dsviewer.glCanvas)
+
+        return im
+