[论文解读] Trackastra: Transformer-based cell tracking for live-cell microscopy

Benjamin Gallusser, Martin Weigert|arXiv (Cornell University)|May 24, 2024

Cell Image Analysis Techniques被引用 4

一句话总结

Trackastra 是一种基于 Transformer 的细胞追踪方法，通过仅使用物体位置和基本形状特征，在短时间内序窗口内直接学习细胞检测之间的成对关联。它通过采用父系 Softmax 机制提升分裂细胞的追踪性能，在仅使用简单贪心链接的情况下也能实现高精度结果，从而减少对复杂优化流水线的依赖，在多种生物数据集（包括细菌、细胞培养物和荧光颗粒）上实现了最先进性能。

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.

研究动机与目标

• 解决活细胞显微成像中准确细胞追踪的挑战，其中细胞在视觉上相似、频繁分裂且表现出复杂的运动模式。
• 通过神经网络直接学习关联代价，减少在链接步骤中对计算成本高昂的离散优化方法（如 ILP）的依赖。
• 开发一种无需针对特定数据集调优的通用追踪框架，适用于多样化生物数据集。
• 通过新颖的父系 Softmax 归一化显式建模母细胞-子细胞关系，提升分裂细胞的追踪性能。
• 通过直接从数据中学习生物上合理的关联，实现仅使用简单贪心链接的稳健追踪，从而减少对复杂优化流程的依赖。

提出的方法

• Trackastra 将来自短时间窗口（例如 3–6 帧）的物体检测作为 Transformer 编码器-解码器架构中的标记进行处理。
• 每个检测通过位置和基本形状特征（如质心、面积、偏心率）进行嵌入，形成输入标记。
• 模型使用自注意力和交叉注意力机制，计算跨时间的所有成对检测关联。
• 对关联矩阵应用一种新颖的父系 Softmax 归一化，优先考虑生物上合理的母细胞-子细胞链接，尤其是在细胞分裂期间。
• 预测的关联矩阵用于构建候选图，通过贪心链接或整数线性规划（ILP）进行剪枝，生成最终轨迹。
• 模型通过联合使用交叉熵损失计算关联以及专门设计的父系关系损失组件进行端到端训练。

实验结果

研究问题

• RQ1 纯 Transformer 架构是否能在不依赖复杂离散优化的情况下，学习活细胞显微成像中生物上合理的细胞关联？
• RQ2 所提出的父系 Softmax 归一化是否能提升追踪精度，特别是对分裂细胞？
• RQ3 使用多大时间上下文（窗口大小）足以实现基于 Transformer 方法的有效追踪？
• RQ4 单一 Trackastra 模型在不同成像模式和细胞类型的多样化生物数据集上，其泛化能力有多强？
• RQ5 Trackastra 是否能通过简单贪心链接实现具有竞争力的性能，从而减少对计算密集型 ILP 求解器的依赖？

主要发现

• Trackastra 在多个数据集（包括细菌、细胞培养物和荧光颗粒）上实现了最先进性能，优于或匹配高度调优的基线方法。
• 在 Bacteria 数据集上，父系 Softmax 将 AOGM（真实匹配平均重叠）降低了约 20%，无论使用贪心链接还是 ILP。
• 使用 3–6 帧的窗口大小可获得优异性能，而将窗口缩小至 s=2 会导致精度显著下降，表明短时间上下文已足够。
• 移除位置编码或形状描述符会显著增加错误，表明其在空间和形态推理中的重要性。
• 该模型在数据集间具有良好的泛化能力：单一训练好的 Trackastra 模型无需微调即可在 Bacteria 和 HeLa 细胞数据集上表现良好。
• 即使使用贪心链接，Trackastra 在 Bacteria 上仍实现了 23.0 的 AOGM（对比 ILP 的 18.8），表明高精度可不依赖复杂优化实现。