Trackastra: Transformer-based cell tracking for live-cell microscopy (2405.15700v2)

Abstract: Cell tracking is a ubiquitous image analysis task in live-cell microscopy. Unlike multiple object tracking (MOT) for natural images, cell tracking typically involves hundreds of similar-looking objects that can divide in each frame, making it a particularly challenging problem. Current state-of-the-art approaches follow the tracking-by-detection paradigm, i.e. first all cells are detected per frame and successively linked in a second step to form biologically consistent cell tracks. Linking is commonly solved via discrete optimization methods, which require manual tuning of hyperparameters for each dataset and are therefore cumbersome to use in practice. Here we propose Trackastra, a general purpose cell tracking approach that uses a simple transformer architecture to directly learn pairwise associations of cells within a temporal window from annotated data. Importantly, unlike existing transformer-based MOT pipelines, our learning architecture also accounts for dividing objects such as cells and allows for accurate tracking even with simple greedy linking, thus making strides towards removing the requirement for a complex linking step. The proposed architecture operates on the full spatio-temporal context of detections within a time window by avoiding the computational burden of processing dense images. We show that our tracking approach performs on par with or better than highly tuned state-of-the-art cell tracking algorithms for various biological datasets, such as bacteria, cell cultures and fluorescent particles. We provide code at https://github.com/weigertlab/trackastra.

• The paper introduces a transformer architecture that directly predicts pairwise cell associations within temporal windows, simplifying the tracking process.
• It employs a blockwise parental softmax normalization to enforce biologically plausible, one-to-many associations, enhancing tracking accuracy across cell divisions.
• Evaluation on diverse datasets demonstrates that the model outperforms traditional and deep learning-based algorithms, particularly in high-density scenarios.

Transformer-based Cell Tracking for Live-cell Microscopy

The paper "Transformer-based cell tracking for live-cell microscopy" by Gallusser and Weigert presents a novel approach to cell tracking in microscopy videos, leveraging transformer architectures. This paper addresses a critical image analysis task similar to Multiple Object Tracking (MOT), but notably more challenging due to the presence of numerous similar-looking objects that may divide over time. The proposed solution deviates from the traditional tracking-by-detection paradigm which involves discrete optimization methods for linking detected cells across frames. Instead, the authors introduce a transformer-based method that directly learns the pairwise associations of cells within a temporal window from annotated data, significantly simplifying the linking process by enabling the use of a greedy algorithm.

Key Contributions

1. Transformer Architecture for Cell Tracking: The authors propose a plain transformer architecture, specifically designed to operate on the spatio-temporal context of detections within a temporal window. This architecture not only simplifies the computations by avoiding dense image processing but also directly accounts for cell divisions.
2. Parental Softmax Normalization: To enforce biologically plausible associations during training, the authors introduce a blockwise parental softmax normalization for the association matrix. This method ensures that each object's parent detection is unique while allowing for multiple child associations.
3. Evaluation Across Diverse Datasets: The method is evaluated on various biological datasets, including bacteria colonies, cell cultures, and fluorescent particles. The performance matches or surpasses state-of-the-art cell tracking algorithms, demonstrating its robustness and generalizability.

Methodology

Dataset Construction

The dataset construction involves the segmentation of raw image sequences into overlapping temporal windows. For each window, object features such as position and basic shape descriptors are extracted. These features, encoded as tokens, serve as input to the transformer model which predicts an association matrix representing the probabilities of pairwise associations between detections.

Transformer Model

The model consists of an encoder-decoder transformer with multi-head attention layers. This architecture allows for reasoning across all object detections within the temporal window. The input tokens are constructed by concatenating learned Fourier spatial positional encodings with object features, followed by linear projection. The transformer layers then process these tokens to predict the pairwise association matrix.

Training

The training process utilizes a binary cross-entropy loss function with the parental softmax normalization to guide the learning towards correct associations. Notably, the method emphasizes important biological constraints, such as allowing exactly one parent detection but multiple child detections for each object.

Inference and Linking

During inference, the predicted associations are averaged over all temporal windows to construct a global association matrix. The final tracking graph is generated using a greedy or ILP (Integer Linear Programming) linking algorithm, ensuring adherence to biological constraints such as non-fusion of objects.

Results and Implications

Quantitative Performance: The proposed transformer-based approach shows significant improvements in tracking performance compared to traditional and recent deep learning-based methods across multiple datasets. For example, in the Bacteria Colony dataset, the method achieved near-perfect tracking results, reducing errors significantly compared to Delta 2.0.

Versatility: The transformer-based model demonstrates its capability to generalize across different domains. It performs well even in high-density scenarios, such as vesicles in the ISBI particle tracking challenge, further proving its robustness.

Future Directions:

• End-to-end Training: Integrating detection and tracking into an end-to-end framework could enhance performance, especially in noisy scenarios or when working with low-quality input data.
• Higher Dimensionality: The extendibility of the transformer architecture to 3D datasets, such as volumetric imaging of biological samples, represents a promising future direction.
• Real-time Processing: Extending this model to support real-time processing applications in live-cell imaging could be highly beneficial for real-time diagnostics and research.

In summary, this paper presents a significant advancement in the field of cell tracking in live-cell microscopy by introducing a transformer-based method that simplifies the traditional tracking-by-detection paradigm. The impressive performance metrics across various datasets and the potential for further improvements underscore the importance of this research in advancing automated cell tracking methodologies.