Trackastra - Tracking by Association with Transformers

Trackastra is a cell tracking approach that links already segmented cells in a microscopy timelapse by predicting assocations with a transformer model that was trained on a diverse set of microscopy videos.

If you are using this code in your research, please cite our preprint:

Benjamin Gallusser and Martin Weigert
Trackastra - Transformer-based cell tracking for live-cell microscopy
arXiv, 2024

Examples

Nuclei tracking	Bacteria tracking
deepcell_results.mp4	bacteria_results.mp4

Installation

This repository contains the Python implementation of Trackastra.

Please first set up a Python environment (with Python version 3.10 or higher), preferably via conda or mamba.

Trackastra can then be installed from PyPI using pip:

pip install trackastra

For tracking with an integer linear program (ILP, which is optional)

conda create --name trackastra python=3.10 --no-default-packages
conda activate trackastra
conda install -c conda-forge -c gurobi -c funkelab ilpy
pip install trackastra[ilp]

Notes:

• For the optional ILP linking, this will install motile and binaries for two discrete optimizers:

  1. The Gurobi Optimizer. This is a commercial solver, which requires a valid license. Academic licenses are provided for free, see here for how to obtain one.

  2. The SCIP Optimizer, a free and open source solver. If motile does not find a valid Gurobi license, it will fall back to using SCIP.

• On MacOS, installing packages into the conda environment before installing ilpy can cause problems.

• 2024-06-07: On Apple M3 chips, you might have to use the nightly build of torch and torchvision, or worst case build them yourself.

Usage

The input to Trackastra is a sequence of images and their corresponding cell (instance) segmentations.

Napari plugin

For a quick try of Trackastra on your data, please use our napari plugin, which already comes with pretrained models included.

Tracking with a pretrained model

The available pretrained models are described in detail here.

Consider the following python example script for tracking already segmented cells. All you need are the following two numpy arrays:

• imgs: a microscopy time lapse of shape time,(z),y,x.
• masks: corresponding instance segmentation of shape time,(z),y,x.

The predicted assocations can then be used for linked with several modes:

• greedy_nodiv (greedy linking with no division) - fast, no additional dependencies
• greedy (greedy linking with division) - fast, no additional dependencies
• ilp (ILP based linking) - slower but more accurate, needs motile

Otherwise, no hyperparameters to choose :)

import torch
from trackastra.model import Trackastra
from trackastra.tracking import graph_to_ctc, graph_to_napari_tracks
from trackastra.data import example_data_bacteria

device = "cuda" if torch.cuda.is_available() else "cpu"

# load some test data images and masks
imgs, masks = example_data_bacteria()

# Load a pretrained model
model = Trackastra.from_pretrained("general_2d", device=device)

# or from a local folder
# model = Trackastra.from_folder('path/my_model_folder/', device=device)

# Track the cells
track_graph = model.track(imgs, masks, mode="greedy")  # or mode="ilp", or "greedy_nodiv"


# Write to cell tracking challenge format
ctc_tracks, masks_tracked = graph_to_ctc(
      track_graph,
      masks,
      outdir="tracked",
)

You then can visualize the tracks with napari:

# Visualise in napari
napari_tracks, napari_tracks_graph, _ = graph_to_napari_tracks(track_graph)

import napari
v = napari.Viewer()
v.add_image(imgs)
v.add_labels(masks_tracked)
v.add_tracks(data=napari_tracks, graph=napari_tracks_graph)

Training a model on your own data

To run an example

• clone this repository and got into the scripts directory with cd trackastra/scripts.
• download the Fluo-N2DL-HeLa dataset from the Cell Tracking Challenge into data/ctc.

Now, run

python train.py --config example_config.yaml

Generally, training data needs to be provided in the Cell Tracking Challenge (CTC) format, i.e. annotations are located in a folder containing one or several subfolders named TRA, with masks and tracklet information.