Relevant work

Cell tracking methods can be categorized into the following two groups: (1) detect and link, and (2) registration-based. Detect and link algorithms have two distinct steps [6, 11, 12, 14, 21, 22]: (a) Detection, where identity-blind candidate locations for objects are proposed by a segmentation or keypoint detection algorithm at each time frame independently. (b) Linking, where temporal associations between detected objects are determined to establish a single continuous worldline across all frames for each individual object. A major drawback of this two-step approach is the propagation of errors from the detection step. Errors that occur in the detection step are difficult to recover from, and they can have detrimental effects on linking and overall tracking quality. Several linking methods have been proposed that are robust to detection outliers, but they either require training with large amounts of manually or synthetically produced ground-truth data, or are not scalable to lengthy videos [11, 21].

An alternative approach is registration, which directly operates in the image space and optimize some transformation parameters that align a frame to some other frame [4, 15, 23–27]. This is done by mapping the underlying image grid from the source to the reference space using the transformation parameters and interpolated pixel values. The transformation parameters must be optimized for each new image over a number of iterations.

Fortunately, recent advances in spatial transformers and differentiable grid sampling have dramatically decreased computational burden and increased performance via GPU acceleration for registration-based approaches [24, 27–30]. Similarly, modern optimization packages such as PyTorch allow the construction of dynamic computational graphs that support more complex nonlinear transformation families and novel cost functions with various regularizers.

Here, we build upon these recent advances to develop a new registration-based approach: ZephIR, a semi-supervised multiple object tracking algorithm with a novel cost function that can incorporate a diverse set of spatio-temporal constraints that can change dynamically during optimization. Our proposed method is capable of efficiently and accurately tracking keypoints in a wide range of 2D or 3D videos. It allows the user to tune a number of parameters controlling the relative strengths of the registration loss and other constraints, and hence generalizes well to a wide range of biological assumptions. To showcase the efficacy and versatility of our method, we demonstrate its performance on a number of biological applications, including cell tracking and posture tracking.

Frame sorting

ZephIR tracks a fixed set of keypoints within a volume over time by matching keypoints between an annotated reference frame and an unlabeled child frame. To fully analyze a movie, we need to register every frame to a reference frame.

For many datasets, it is best to register every child frame directly to a coarsely similar reference frame, and let annotations for that reference frame provide initial guesses for keypoints in the child frame (Fig 1A). For this, we must identify a set of representative reference frames that capture the range of deformation patterns present in the movie, and we must assign each remaining frame to one of those reference frames. A pairwise distance between all pairs of frames is determined by a similarity metric (e.g. correlation coefficient) applied to low-resolution thumbnails. A k-medoids clustering algorithm is applied in the chosen “similarity” space to identify a small number of median frames to best serve as reference frames for all other frames in the corresponding cluster (Fig 1A and S2 Fig) [4, 25].

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 1. Overview of frame sorting strategies.

Orange indicates fully annotated reference frames, blue indicates parent frames with at least one child frame, and green indicates child frames. A: In the simplest strategy, all frames are initialized by the closest reference frame. B: Frames are sorted into ordered queues based on similarity. Each of these branches start with a reference frame, and new child frames are added such that the parent-child similarity distance is minimized, naturally clustering similar frames around each reference frame. C: Frames are sorted chronologically, branching both forward and backwards from each reference frame.

https://doi.org/10.1371/journal.pcbi.1012075.g001

In other datasets, the tracking results from one frame may provide useful insight into the solution for another frame. For example, a frame that is close (in similarity space) to the reference may be easy to track. The tracked results from that frame, in turn, may provide a better initial guess for keypoints in a frame that is further away from the reference. This can reduce the distance between the initial guess and correct positions as well as the difficulty of the optimization problem (S2 Fig). Thus, every child frame being registered is associated not only with a reference frame (a registration target), but also a previously registered parent frame, which provides the initialization prior to optimization (Fig 1B).

Additionally, the learning rate offers a way to improve the optimization trajectory by controlling the rate at which the optimized parameters, , are updated after each registration iteration. The learning rate for the child frame is partly determined by the distance between the parent and child frames. We expect that when a parent-child pair are close in the similarity space, the keypoints do not undergo significant displacements. Hence, given a set number of registration iterations, we may apply a lower learning rate for a similar parent-child pair to limit displacements between the two frames, and scale up to a higher learning rate in the case of a dissimilar pair to allow tracking of features much further away from the initial coordinates. The combination of these effects produces a flexible limit on the range of possible optimization results for the child frame based on coarse similarity to its parent frame [31–33].

To take full advantage of this parent-child interaction, we sort all frames into distinct sequences of parent-child frames based on similarity. Each of the resulting branches begin from a previously selected median reference frame. The subsequent child frames are selected to minimize the distance from a parent frame until every frame is assigned to a branch. Doing so produces unique sets of frames that stem from each reference frame, naturally forming clusters that separate similar frames from dissimilar ones. This is particularly useful for datasets that repeatedly sample from a limited set of postures or global spatial structures (e.g. locomotion).

However, not all datasets have temporal patterns that can reliably make use of the similarity-based initialization method. For such datasets, a chronologically sorted queue may be more reasonable and provide better accuracy overall, where a branch simply stems from each reference frame both forwards and backwards in time until it encounters the first frame, the last frame, or another branch (Fig 1C). Note that the parent-child interactions during tracking are still the same regardless of the sorting method. For a chronologically sorted queue, the controlled variation of learning rates effectively allows us to adapt to different capture frame rates. A high frame rate video often captures smooth motion that benefits from low learning rates but a low frame rate video does not.

ZephIR tracks all keypoints in an image or volume (Fig 1A) simultaneously before moving onto a subsequent frame in the sorted queue (Fig 1B). This matching is done by minimizing a loss function with four contributions:

We measure overlap of local image features around the keypoint via . We measure relative elastic motion between keypoints via . We measure the distance of each keypoint to the nearest candidate location from a precomputed set via . We measure smoothness of keypoint-determined dynamical features (e.g. fluorescence or motion) via . Each is described in more detail below.

The relative weights of each term, λ, can be freely adjusted by the user to better fit a particular dataset. The user can also set the relative weights to change while tracking a single frame to allow the algorithm to shift focus to different loss components over a number of optimization iterations.

Image registration,

The first term of our algorithm measures overlap of local image descriptors.

For each keypoint i in a child frame, I^(c), we generate a sampling grid centered around that keypoint’s coordinates in 3D space, (origin at frame center, with fixed z = 0 for 2D images). A sampling grid defines a map for a pixel (or an interpolated subpixel) from the original image to each pixel of an output image of the same spatial shape as the sampling grid, therefore producing arbitrary transformations of the image, such as simple crops or complex nonlinear deformations. We define a set of transformation parameters, , that is closely related to but may include additional transformation models, such as rotation, to characterize and construct a sampling grid centered at the keypoint coordinate. We use the grid to sample from the child frame accordingly and produce a representation of the local image information, i.e. an image descriptor, around each keypoint (Fig 2C, top left, blue): .

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2. Overview of ZephIR algorithm.

A: Examples of input datasets, created with BioRender.com. ZephIR can track keypoints in various biological systems, including fluorescent cellular nuclei in a tissue and body parts that summarize a posture. Input dimensions can range from 3D (time, XY) to 5D (time, channel, XYZ). Colored dots indicate example keypoints to be tracked. B: Frame sorting schemes. A branch defines an ordered queue of frames to be tracked. Each branch begins at a manually annotated reference frame (orange), but subsequent parent (blue) and child (green) frames in a single branch can be sorted either by chronology (top) or by minimizing the similarity distance between each parent-child pair (bottom). C: Overview of tracking loss. Tracking loss is comprised of four terms: 1) (top left), overlap of local image features around each keypoint, sampled from the current frame and its nearest reference frame, 2) (top right), elastic connections between neighboring keypoints with varying stiffnesses based on covariance of the connected keypoints, 3) (bottom left), proximity to features detected by a shallow model selector network that takes in a number of existing feature detection software as input channels, 4) (bottom right), smoothness of temporal dynamics at each keypoint position. D: Overview of steps for manual verification and additional supervision. Users can verify tracking results as correct or identify incorrect results. After fixing a few key incorrect results, ZephIR can use those new annotations as well as the verified correct tracking results to improve tracking results for all other keypoints in that frame (and all its child frames).

https://doi.org/10.1371/journal.pcbi.1012075.g002

The sampled descriptors are foveated to prioritize more local information relative to the neighboring features. In lieu of image pyramids [34], we apply a Gaussian blur to reduce the effective resolution of the descriptors at the start of optimization. The resolution is dynamically increased every few registration iterations as the magnitude of the blur is decreased. We empirically observed that doing so avoids gradient values that are too small or large, both of which can occur in regions with sharp, well-defined edges surrounded by a uniform background. On the other hand, restoring the original resolution of the image still provides the best available information for fine-tuning tracking results towards the end of the optimization loop.

Similarly, a set of reference descriptors that serve as registration targets are sampled from a reference frame, I^(r). The user-defined annotations for that reference frame are used to calculate fixed set of transformation parameters, , and thus a fixed set of sampling grids that generate a reference descriptor centered around each annotation (Fig 2C, top left, orange): .

Using the two sets of image descriptors, our registration loop optimizes the transformation parameters for the child frame, , to minimize the following loss term:

The optimized parameters are then used to calculate the desired results, the keypoint coordinates for the child frame, . Note that these coordinates are also used for different loss components below, but as is calculated from , gradients are always accumulated at .

Spatial regularization,

Cellular motion within a tissue tends to be highly correlated, but these correlations can be hidden in sparse fluorescent movies that only highlight a small number of cells (or subcellular features) [35]. Even in less sparse movies, correlations between nearby keypoints may not be well-captured by descriptors, especially when deformations, noise, or lighting conditions prevent descriptor alignment. In order to maintain a similar spatial structure during motion without relying on highly specialized skeletal models, such as those used in animal posture trackers [36–38], we add an elastic spring network between neighboring keypoints [35, 39, 40]. The resulting penalty to relative displacement of neighboring keypoints prevents unreasonable deformations, providing a simple and flexible spatial heuristic of the global structure and motion present in the data (Fig 2C, top right).

Despite ZephIR tracking a fixed number of keypoints across the video, the spring network makes it robust to fluctuations in the number of keypoints visible in a frame. In frames where a keypoint may not be visible or present, it fails to produce useful image descriptors for registration, but the spring connections to its neighboring keypoints allow us to keep track of its approximate location.

Each of the i keypoints being tracked is connected to j nearest neighbors to define the following loss term: where describes the distance between keypoints i and j in the frame t.

The stiffness of each spring connection, k_ij, is initially set to a constant 1.0 for all connections. However, when multiple reference frames are available, k_ij is further adjusted to better model the spatial patterns in the data by accounting for the covariance of connected keypoints across all available n reference frames, r_1:n:

This ensures that connections between highly covariant keypoints are made stronger while connections between keypoints with more weakly correlated motion are weakened or cut accordingly.

Feature detection,

For this component of the algorithm, we solve an easier problem of identity-blind feature detection, as such detection algorithms have been shown to be fruitful in the context of tracking [6]. Namely, we identify key features (such as the center of a cell) present in a volume without matching them to a specific feature in some other volume. Adding this information allows us to dynamically combine both registration and detect-and-link strategies into a single loss function.

This object or feature detection problem has been well-studied, and a wide variety solutions have been proposed. Solutions can range from more parameter-free algorithms (e.g. Richardson-Lucy deconvolution [41, 42]), to algorithms requiring more fine-tuning (e.g. Viola-Jones [43]). More recently, deep convolutional neural networks have shown to be powerful, effective solutions as well (e.g. StarDist [10]). Importantly, each of these approaches may work better or worse on different classes of images. Generalization to new datasets can be hard to predict, especially for neural networks that are trained on data generated from a single source.

Our approach is to automatically evaluate simple combinations of these established algorithms by using a simple model-selecting network. After identifying a set of candidate models, we provide the outputs of these models as input channels to a small convolutional neural network (CNN), capable of producing optimized combinations of some or all of the selected models. If a particular model is best suited for a dataset, network weights for the corresponding input channel are increased during training while suppressing other channels, thus it is not critical for the user to select an “optimal” set of models since even one well-suited model can sufficiently increase performance. The model selector network can be trained using any fully annotated frame, and the low number of learnable parameters in the network also allows fast training for each new type of data or imaging condition, which in turn allows rapid experimentation with new selections of models to test as inputs.

The ultimate output of this selector network, C(I^(c)), is formulated as a probability map, where each pixel of the original image is assigned some probability of being a desired feature (Fig 2C, bottom left). By indexing the map to obtain the assigned probability at the keypoint coordinates, , we use this information to push tracking results towards detected features:

Temporal smoothing,

Given a sufficiently fast imaging rate, we expect pixel intensity values to be smooth across a small local patch of frames, even for cellular datasets where pixel intensities represent smoothly-varying dynamical signals [44, 45]. Thus, we attempt to maintain smoothly-varying local pixel intensities as a form of temporal regularization. For datasets where expected dynamics are appreciably slower than the imaging rate, the strongest version of this regularization is to penalize any deviation from a local zeroth-order fit. We apply this across a small (ϵ ≈ 2 frames) patch of frames from c − ϵ to c + ϵ that are registered simultaneously, and add this to the loss for the center frame, c (Fig 2C, bottom right):

Note that since the loss term is applied for the center frame only, it does not affect the results for the other frames despite registering all frames in the patch together. Additionally, this component of the algorithm requires registration (or approximate registration) of nearby frames, making it more appropriate in low-motion conditions or after initial coarse registration is complete.

User intervention

Our pipeline allows a user to dramatically improve tracking quality in various ways by providing further supervision. Providing additional fully annotated frames will improve registration targets to better match descriptors from similar frames. Strategically selecting a new reference frame can have dramatic impacts on frame sorting as well, creating opportunities to form tighter clusters of parent-child branches.

Furthermore, when multiple reference frames are present, covariance of keypoints in those frames helps better define an implicit global spatial structure by modulating stiffnesses of the spring connections between neighboring keypoints, k_ij. Any additional reference frames can provide more accurate covariances, and thus a spatial model that is more accurately tailored for that particular dataset.

Partially annotated frames are not used to seed sorted frame branches nor used to sample reference descriptors. Still, all user annotations present in the frame are utilized to improve the tracking quality of the remaining keypoints in that frame (Fig 2D).

Firstly, prior to gradient descent, displacements between any annotations in the child frame and their corresponding coordinates in the parent frame are used to interpolate a flow field, Φ, which serves as a rough model of the global motion between the two frames [15, 26, 40]. The flow field provides an estimated displacement between the two frames for the remaining unannotated keypoints which can be applied to their coordinates in the parent frame to initialize those keypoints closer to their new positions in the child frame. This is particularly helpful for pairs of parent-child frames with large motion between them. While randomly annotating a small number of keypoints distributed across the frame is effective, the flow field can always be improved in both precision and accuracy by adding more annotations for the child frame (S3 Fig).

Secondly, the spatial regularization during the optimization process, , also makes good use of any partial annotations. The annotations are fixed in place, but the spring connections to their neighbors remain a crucial component of the backwards gradient calculations and helps to “pull” the connected keypoints into place.

To streamline the process of providing user supervision, we offer a browser-based graphical user interface that provides an intuitive, simple environment to produce and save further annotations. Since our approach lacks a slow “training” phase, any new annotations can be applied to tracking a frame directly from the GUI. A macro available in the GUI executes a temporary state of the algorithm quickly and efficiently, allowing users to see the precipitated improvements immediately.

Additionally, the GUI provides an opportunity for users to provide supervision without creating new annotations. The user may upgrade individual results into annotations or entire frames into new reference frames by marking them as correctly tracked. These user-confirmed frames will be treated as a regular reference frame next time the algorithm is executed, benefiting from all the improvements to tracking quality discussed previously. These improvements to the rest of the results can be observed immediately by executing the algorithm from the GUI. A high-level summary of this procedure is provided as Algorithm 1 below.

Algorithm 1 ZephIR optimization loop

for c ∈ sorted_frame_list do

 if then

                 ⊳ Partial annotations

     ⊳ Interpolate flow field

    ⊳ Initialize at prediction

 else

          ⊳ Initialize at parent results

 end if

 for n ← 1, n_epoch do

  Backwards          ⊳ Backpropagate gradients

  Update          ⊳ Gradient descent

 end for

        ⊳ Get keypoint coordinates

 Write           ⊳ Save coordinates

end for

Neurons in crawling worms (C. elegans)

Optical methods based on fluorescence activity of calcium binding indicators has become a standard tool for observing neuronal activity in C. elegans. To do so, it is necessary to track fluorescent signals from individual neurons across every frame in a recording. This poses a significant challenge, particularly when the animal is allowed to freely crawl. The worm’s brain undergoes fast, dramatic, nonaffine deformations, exhibiting a large variety (forward and backward motion, omega turns, coils, pharyngeal pumping, etc.) and magnitude (up to ten microns relative to an internal reference frame) of movements as the animal behaves [20, 47–49].

Many solutions have been proposed to track fluorescent neurons in C. elegans. Two step (detect and link) approaches often suffer from the lack of reliable detection algorithms and require relatively low frame-to-frame motion in order to accurately link the detected neurons [6, 14, 23]. Similarly, deep learning approaches are limited by insufficient training data, often failing to generalize across different animals, even those within the same strain [11, 25]. While these approaches have provided important insight and progress, there remains substantial need for improvement in accuracy and efficiency when tracking many neurons in freely behaving worms.

Fig 3 describes ZephIR’s general analysis workflow and performance on tracking a set of 178 fluorescent neurons in the head of a freely behaving worm across a 3D recording of approximately 4.4 minutes (1060 frames @ 4Hz). We collected this data for the purpose of testing this algorithm using a microscope and technique similar to that described in [47]. The video has been centered and rigidly rotated to maintain a consistent orientation of the worm, but no further straightening has been done. With only a few manually annotated reference frames (78 mins/frame on average with the provided GUI), ZephIR already achieves state-of-the-art accuracy as reported in recently published works on neuron tracking [11, 14, 23, 46] (Fig 3B, Table 1), where accuracy is measured as the average percentage of neurons correctly tracked, and a neuron is considered correctly tracked if the tracked coordinate is within the volume of the neuron as identified via a manual annotator. First three median frames determined via k-medoids clustering and recommended for annotation as reference frames clearly represent the principal postures that are repeatedly sampled during locomotion (Fig 3A). Fully annotating those three frames (534 total annotations) is sufficient to produce tracking accuracy of 89.27%, comparable to previously published works on neuron tracking (Fig 3B, Table 1). Annotating additional median frames can increase this accuracy up to 91.49% for a total of 10 reference frames.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3. ZephIR analysis workflow and results for tracking GCaMP fluorescence from neuronal nuclei in 3D volumes of freely behaving C. elegans.

A: Plot of mean distance (in similarity space) to the nearest reference frame vs the number of reference frames (left), and the first three median frames (maximum intensity projections of shape 200 x 512) recommended by ZephIR’s k-medoids clustering algorithm (right). The first three median frames clearly represent the three main postures that the worm cycles through as it crawls. B: Accuracy (higher is better) and precision (lower is better) vs the number of reference frames. Accuracy is measured as the average percentage of neurons correctly tracked, where a neuron is considered correctly tracked if the tracked coordinate is within the volume of the neuron as identified via a manual annotator. Precision is presented as the average RMS error between the predicted position and the manually annotated position of each neuron in pixels. Note that once the majority of the postures present in the data is well-represented by the first three reference frames, subsequent additions returns diminished improvements. Last data point shows ZephIR’s accuracy using 10 reference frames with 10 partial annotations across all frames (panel C). C: 10 neurons were randomly selected to be verified or corrected to serve as partial annotations. Traces of 5 of these neurons extracted using the initial ZephIR results with 10 reference frames (left), and those using verified true positions (right) are shown, along with 5 other randomly selected neurons. Traces are calculated as fold change of the ratio between GCaMP and RFP fluorescence of each neuronal nuclei over the baseline, where the baseline is defined as the ratio in the first frame. Tracking quality for these 10 neurons can also be seen in individual crops around the neurons averaged across all frames (sharper image of the cell at the center reflects better accuracy and precision in tracking). Note how the five unannotated neurons show improvements in tracking quality after the addition of partial annotations, exemplifying the effects of partial annotations on the unannotated neurons in the same frame. D: Neuronal activity traces from 178 neurons, extracted using results from ZephIR with 10 reference frames and 10 partial annotations in all frames. Traces are calculated as fold change of the ratio between GCaMP and RFP fluorescence of each neuronal nuclei over the baseline, where the baseline is defined as the ratio in the first frame. Behavior is shown in the ethogram below the heatmap. Trajectory of the worm (t = 0 at bottom right) is also colored with the behavior state at the time. Trajectory of the worm matches changes in behavior over time as expected, and many of the neuronal activity traces show strong correlation with behavior. E: Accuracy vs the number of reference frames for tracking 79 neurons in a publicly available dataset of freely moving C. elegans [14]. Since the spatiotemporal patterns in the data are similar to the previously tracked data, we can reuse the same parameters and follow the same procedure to track the 79 neurons in the head of the worm. Each volume has been centered and rotated but no further straightening has been done. Since this particular dataset has also been used to benchmark a number of recently developed algorithms, we may also compare ZephIR’s accuracy with Neuron Registration Vector Encoding (NeRVE) [14], fast Deep Neural Correspondence (fDNC) [11], 3DeeCellTracker [23], and CeNDeR [46].

https://doi.org/10.1371/journal.pcbi.1012075.g003

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Table 1. Summary of results and performance of ZephIR and other recently developed algorithms for tracking neurons in freely moving C. elegans.

When benchmarking on a common dataset (provided by [14]), ZephIR achieves top 2 accuracy with significantly less annotations and faster performance. Note that tracking inference speed does not include time spent for any potential training or detection.

https://doi.org/10.1371/journal.pcbi.1012075.t001

We further improve on the accuracy of the initial results by providing additional supervision. We randomly selected ten neurons uniformly distributed throughout the brain to verify and use as partial annotations across all frames (Fig 3C). Because the initial results already achieved high accuracy, they only required correction for a subset of frames (≈ 15%), taking an average of 51 mins/neuron of human time to complete across the entire movie. After this correction and validation, annotations for these ten neurons were re-classified as manual annotations in all frames. The partial annotations produce a dramatic improvement in accuracy (red data point in Fig 3B) without the need to verify entire frames.

Through this workflow, we are able to achieve a sufficiently high accuracy to extract neuronal activity traces across the entire recording (Fig 3D) [44, 45]. Neuronal activity is calculated as the average ratio of GCaMP and RFP fluorescence intensities of the 9 brightest pixels in a 3x7x7 volume (roughly the size of the cell nucleus in the image) centered around the tracked neuron coordinates after masking out pixels that overlap with the volumes around 5 nearest neighbors. Traces are adjusted for photobleaching (double exponential fit) and are shown relative to the ratio of intensities seen in the first frame. Many neurons show clear correlation with observed behaviors, and the activity patterns are comparable to previously published works [14, 49–51].

In order to verify and benchmark the accuracy of our method, we apply the same workflow to a similar dataset of a freely moving worm provided by Nguyen, et al. [14], a publicly available dataset with ground-truth results provided that has already been used as a benchmark for other recently published algorithms. [11, 14, 23, 46] Since this dataset is similar in behavior, motion, and imaging conditions, we may reuse the same parameters and follow the same workflow as the previously discussed C. elegans dataset. Doing so, we are again able to achieve state-of-the-art accuracy (84.0%) with only a few reference frames, and adding additional partial annotations increases our top accuracy further (94.48%) (Fig 3E). The top accuracy, the required annotations or training frames, and tracking inference speed (not including any training or detection) from other algorithms recently developed for tracking neurons in the same C. elegans dataset are summarized in Table 1.

Comparison of loss terms and their combinations

By tracking the same set of keypoints in a single dataset with varying combinations for the tracking loss, we may examine how each component affects the accuracy and precision. We analyze the same dataset as discussed in the preceding section and Fig 3A–3D, and limit tracking to three reference frames with fixed weights and parameters (λ_R, λ_N, λ_D). Fig 4 summarizes the resulting tracking accuracy and precision as the loss components are changed. We can confirm the positive contributions of both and as adding either or both to image registration, , increases accuracy and precision (Fig 4B). We also note that have similar local information as the image descriptors when represented as a probability map with peaks at the neuron centers (Fig 4C). When combined with the spring network, , it reaches similar accuracy as image registration, and the sum of the three terms, , result in only a small increase in accuracy compared to either combinations, or (Fig 4D).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 4. Results for tracking GCamP fluorescent neuron nuclei in 3D volumes of freely behaving C. elegans (same dataset as shown in Fig 3A–3D) with varying combinations for the tracking loss.

Three reference frames and fixed weights (λ_R, λ_N, λ_D) were used for all results shown. A. One of the three reference frames. B. Network of connections between neighboring neurons in the frame shown in panel A. The edge weights represent their relative stiffness for calculating the spring loss, . The connections and their stiffness are modulated by the distance between the neurons and the covariance of the connected pair in reference frames. C. Results from feature detection, , on the frame shown in panel A. Feature detection achieves an average precision and recall of 0.948 ± 0.024 and 0.931 ± 0.016, respectively. D. Tracking accuracy (top) and precision (bottom) as the loss components are changed. We can confirm the positive contributions of both and as adding either or both to image registration, , increases accuracy and precision. We also note that have similar local information as the image descriptors when represented as a probability map with peaks at the neuron centers. When combined with the spring network, , it reaches similar accuracy as image registration, and the sum of the three terms, , result in only a small increase in accuracy compared to either combinations, or .

https://doi.org/10.1371/journal.pcbi.1012075.g004

Posture of a behaving mouse

Here, we demonstrate how ZephIR can be used for behavioral tracking in natural movies by analyzing the pose of a head-fixed mouse performing a motor task. The richness of local image features present in natural images lend themselves to registration. In addition, by connecting key points along the mouse’s body, our spring network loss () can implicitly capture the scaffold underlying the mouse’s posture.

There exist many solutions for similar problems in posture tracking. In particular, convolutional neural networks have been successfully implemented for posture analysis in both laboratory and natural settings [3–5]. Notably, DeepLabCut adapts a ResNet CNN architecture to track postures of various animals without any physical markers. DeepLabCut utilizes transfer learning, where a “base” model is trained on a publicly available dataset of various “natural” images prior to specializing the weights to a particular dataset. Their method reduces the amount of training data required to achieve state-of-the-art results by orders of magnitude (from hundreds of thousands to just a couple hundred labeled images), and thus reduces the amount of manual labor required by the experimentalist.

Fig 5 compares the performance of our algorithm and that of DeepLabCut on the same dataset. This dataset of a mouse performing a motor task was produced by the Churchland Lab and published in [52]. We track ten points that summarize the mouse’s posture as it performs the task in a 2D movie with a total of 5400 frames. We show that for low numbers of reference frames, i.e. low numbers of training data, ZephIR can produce much better quality tracking than DeepLabCut, achieving good results with less than 20 reference frames (Fig 5B and 5C). ZephIR is also able to produce this result with significantly less total computation time as it does not require a slow training phase (table of Fig 5, middle row).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 5. Results for tracking posture of a behaving mouse in 2D.

We compare performances of ZephIR and DeepLabCut on tracking 10 body parts that characterize the mouse’s posture over time. A: Accuracy (average percentage of keypoints correctly tracked in unannotated frames) and precision (average distance between predicted position and ground truth position of keypoints) vs the number of manually labeled or ground truth frames. These labeled frames are used as reference frames for ZephIR and as training data for DeepLabCut. The frames are selected based on automated recommendations from each algorithm, meaning the two sets of frames used may not be identical. The last data point (results with 200 training frames) for DeepLabCut are produced with training data generated by verifying and correcting ZephIR results with 10 reference frames. Note that ZephIR achieves better accuracy when only a few labeled frames are provided, but DeepLabCut ultimately reaches a higher accuracy when its training data was augmented with ZephIR. B, C: DeepLabCut and ZephIR results with 20 labeled frames (vertical line in panel A) for tracking mouse body parts as it raises its paws. Note that ZephIR is more stable during motion while DeepLabCut tends to jump between the different body parts. Table: Annotation and computation speed comparison. Annotation time is calculated for the same person, using the respective GUI’s provided with each software package. Training and inference times are tested on the same CPU and single GPU environment and with 20 reference frames (vertical line in panel A). While DeepLabCut is faster for inference, it requires a slow training phase, dramatically increasing the total computation time. This data was produced and provided by the Churchland Lab (UCLA). Raw data is available at: https://ibl.flatironinstitute.org/public/churchlandlab/Subjects/CSHL047/2020-01-20/001/raw_video_data/.

https://doi.org/10.1371/journal.pcbi.1012075.g005

It is important to note that DeepLabCut can ultimately produce more accurate results when provided with more training data as seen when using 200 manually labeled frames (right-most data points in Fig 5A). While we do not expect ZephIR to replace DeepLabCut in such applications where a higher accuracy than what ZephIR is able to provide on its own is required, ZephIR can easily fit into a DeepLabCut workflow to augment the amount of training data available. Instead of manually labeling the full list of 200 frames to produce the last data point in Fig 5A, we only annotated the first 10 of the recommended frames. We then run ZephIR using those frames as references, verify the tracking results for the remaining 190 frames, and correct any errors to produce the full set of 200 training images to use for DeepLabCut. Including this step cuts the total human time required for a DeepLabCut workflow, from an extrapolated 160 minutes to label 200 frames by hand to 53 minutes.

More recently developed variants of DeepLabCut, such as DeepGraphPose, can also reduce training data size by incorporating spatio-temporal priors and enabling semi-supervised training that uses both annotated and unannotated data [5]. However, these variants still require a significant amount of training data ( of DeepLabCut’s requirements, compared to ZephIR’s ≈ 5 − 10%) and often fail when analyzing sparse or volumetric datasets, making it difficult to employ for biological datasets.

Performance

We are able to compute the loss terms and optimize the tracking parameters efficiently by utilizing modern deep learning tools with, in particular, differentiable grid sampling and GPU acceleration offered by PyTorch [53]. Since our approach does not require a training phase, which is often the most significant bottleneck in both time and resources, it is fast without being computationally costly. We also sacrifice a small amount of performance to reduce the amount of memory required for both CPU and GPU to levels that are reasonable for commercial laptops. This balance can be manually adjusted by the user depending on their computing environment.

We run the following tests on a PC with a 16-core AMD Ryzen Threadripper 1950X processor @ 3.40GHz, 64GB RAM, and an Nvidia GTX 1080Ti GPU, 11GB VRAM. The tests were carried out on the freely behaving worm dataset (Fig 3).

In the default configuration, ZephIR registers 178x1x5x25x25 (NxCxDxHxW) descriptors () with spatial regularization () over 40 optimization epochs for an average of 1.24s total computation time spent per volume. In comparison, similar algorithms such as NeRVE takes an approximate 50 s/vol on over 200 computing cores [14] and 3DeeCellTracker approximately 1 min/vol on a desktop PC with an NVIDIA GeForce GTX 1080 GPU (inference only) [23].

During the test, the process utilized a maximum of 1.84GB RAM and 0.89GB VRAM. The number of descriptors does not significantly affect performance as descriptors are registered in parallel, but the number of epochs will impact speed linearly. The size of descriptors may slightly affect performance as well as memory consumption. The number of volumes in the dataset also affect the speed of frame sorting (Fig 1). When sorting frames by similarity (i.e. minimizing distance between parent and child frames), performance can be affected as O(N²), adding approximately 2.85 minutes to sort 1060 frames of shape 23 x 512 x 512.