8472170680

Committed 28 Mar 2024 06:42PM UTC coverage: 77.01% (+0.2%) from 76.797%

Build # 8472170680

Build Type

Pull #26

github

Committed by

WardLT

Commit Message

Add a drift module

Pull Request Pull Request #26: Add more algorithms for drift correction

Run Details

23 of 26 new or added lines in 1 file covered. (88.46%)

479 of 622 relevant lines covered (77.01%)

0.77 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

88.46

/rtdefects/drift.py

"""Algorithms for correcting drift in microscopy images"""
import numpy as np
import pandas as pd
from scipy.signal import fftconvolve


def compute_drift_from_tracks(tracks: pd.DataFrame, minimum_tracks: int = 1) -> np.ndarray:
    """Estimate the drift for each frame from the positions of voids that were mapped between multiple frames

    We determine the "drift" based on the median displacement of all voids, which is based
    on the assumption that there is no net motion of all the voids.

    We compute the drift in each frame and assume the drift remains unchanged if there are no voids matched
    between a frame and the previous.

    In contrast, trackpy uses the mean and only computes drift when there are matches between frames.

    Args:
        tracks: Track information generated by trackpy.
        minimum_tracks: The minimum number of tracks a void must appear to be used in drift correction
    Returns:
        Drift correction for each frame
    """

    # We'll assume that the first frame has a void
    drifts = [(0, 0)]

    # We're going to go frame-by-frame and guess the drift from the previous frame
    last_frame = tracks.query('frame==0')
    for fid in range(1, tracks['frame'].max() + 1):
        # Join the two frames
        my_frame = tracks.query(f'frame=={fid}')
        aligned = last_frame.merge(my_frame, on='particle')

        # The current frame will be the previous for the next iteration
        last_frame = my_frame

        # If there are no voids in both frames, assign a drift change of 0
        if len(aligned) < minimum_tracks:
            drifts.append(drifts[-1])
            continue

        # Get the median displacements displacements
        last_pos = aligned[['x_x', 'y_x']].values
        cur_pos = aligned[['x_y', 'y_y']].values
        median_disp = np.mean(cur_pos - last_pos, axis=0)

        # Add the drift to that of the previous image
        drift = np.add(drifts[-1], median_disp)
        drifts.append(drift)

    return np.array(drifts)


def compute_drift_from_images(image_1: np.ndarray, image_2: np.ndarray, return_conv: bool = False) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
    """Compute the drift between two different frames

    Args:
        image_1: Starting image
        image_2: Next image
        return_conv: Whether to return the convolution between the two images
    Returns:
        - The optimal translation between the two images
        - Convolution used to make the judgement, if ``return_conv`` is True
    """

    # Compute the correlations between the two images using FFT
    #  You must reverse the second signal/image for this trick
    conv = fftconvolve(image_1, image_2[::-1, ::-1], mode='same')

    # Find the location of the maximum
    peak_loc = np.unravel_index(np.argmax(conv), conv.shape)

    # Find its deplacement from the image center, that's the location
    drift = [peak_loc[1] - conv.shape[0] // 2, peak_loc[0] - conv.shape[1] // 2]
    if return_conv:
        return -np.array(drift), conv
    return -np.array(drift)

1	"""Algorithms for correcting drift in microscopy images"""
2	import numpy as np	1✔
3	import pandas as pd	1✔
4	from scipy.signal import fftconvolve	1✔
5
6
7	def compute_drift_from_tracks(tracks: pd.DataFrame, minimum_tracks: int = 1) -> np.ndarray:	1✔
8	"""Estimate the drift for each frame from the positions of voids that were mapped between multiple frames
9
10	We determine the "drift" based on the median displacement of all voids, which is based
11	on the assumption that there is no net motion of all the voids.
12
13	We compute the drift in each frame and assume the drift remains unchanged if there are no voids matched
14	between a frame and the previous.
15
16	In contrast, trackpy uses the mean and only computes drift when there are matches between frames.
17
18	Args:
19	tracks: Track information generated by trackpy.
20	minimum_tracks: The minimum number of tracks a void must appear to be used in drift correction
21	Returns:
22	Drift correction for each frame
23	"""
24
25	# We'll assume that the first frame has a void
26	drifts = [(0, 0)]	1✔
27
28	# We're going to go frame-by-frame and guess the drift from the previous frame
29	last_frame = tracks.query('frame==0')	1✔
30	for fid in range(1, tracks['frame'].max() + 1):	1✔
31	# Join the two frames
32	my_frame = tracks.query(f'frame=={fid}')	1✔
33	aligned = last_frame.merge(my_frame, on='particle')	1✔
34
35	# The current frame will be the previous for the next iteration
36	last_frame = my_frame	1✔
37
38	# If there are no voids in both frames, assign a drift change of 0
39	if len(aligned) < minimum_tracks:	1✔
NEW 40	drifts.append(drifts[-1])	×
NEW 41	continue	×
42
43	# Get the median displacements displacements
44	last_pos = aligned[['x_x', 'y_x']].values	1✔
45	cur_pos = aligned[['x_y', 'y_y']].values	1✔
46	median_disp = np.mean(cur_pos - last_pos, axis=0)	1✔
47
48	# Add the drift to that of the previous image
49	drift = np.add(drifts[-1], median_disp)	1✔
50	drifts.append(drift)	1✔
51
52	return np.array(drifts)	1✔
53
54
55	def compute_drift_from_images(image_1: np.ndarray, image_2: np.ndarray, return_conv: bool = False) -> np.ndarray \| tuple[np.ndarray, np.ndarray]:	1✔
56	"""Compute the drift between two different frames
57
58	Args:
59	image_1: Starting image
60	image_2: Next image
61	return_conv: Whether to return the convolution between the two images
62	Returns:
63	- The optimal translation between the two images
64	- Convolution used to make the judgement, if ``return_conv`` is True
65	"""
66
67	# Compute the correlations between the two images using FFT
68	# You must reverse the second signal/image for this trick
69	conv = fftconvolve(image_1, image_2[::-1, ::-1], mode='same')	1✔
70
71	# Find the location of the maximum
72	peak_loc = np.unravel_index(np.argmax(conv), conv.shape)	1✔
73
74	# Find its deplacement from the image center, that's the location
75	drift = [peak_loc[1] - conv.shape[0] // 2, peak_loc[0] - conv.shape[1] // 2]	1✔
76	if return_conv:	1✔
NEW 77	return -np.array(drift), conv	×
78	return -np.array(drift)	1✔

ivem-argonne / real-time-defect-analysis / 8472170680

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous