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 using pip directly from this repository:

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 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.

Usage

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

Tracking with a pretrained model

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
import numpy as np
from trackastra.utils import normalize
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()

# Normalize your images
imgs = np.stack([normalize(x) for x in imgs])

# 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)

Napari plugin

We additionally provide a napari plugin which allows one to quickly apply pretrained and custom models on custom timeseries.

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.