Skip to content

Trackastra stalls out during linking #72

Description

@tlnagy

Hi @bentaculum and @maweigert, thanks for the fantastic tool. I've had great success with using Trackastra, but I sometimes struggle when using it with certain datasets where it seems to stall out for many hours when linking (both in greedy and ilp modes).

These datasets are fairly large: 1441x2304x4096 with ~ 500 cells per frame.

Here's the code that I use to run Trackastra:

Analysis code
import torch
import nd2
import tifffile
import dask.array as da
from dask import delayed
from trackastra.model import Trackastra
from trackastra.tracking import graph_to_ctc, graph_to_napari_tracks
from pathlib import Path

device = "cuda"

model = Trackastra.from_pretrained("general_2d", device=device)

folder = sorted(Path("../segs").iterdir())[0]

sample = tifffile.imread(sorted(folder.glob('*.tif*'))[0])

lazy_imread = delayed(tifffile.imread)

filenames = sorted(folder.glob('*.tif*'))

lazy_arrays = [lazy_imread(fn) for fn in filenames]

dask_arrays = [
    da.from_delayed(lazy_reader, shape=sample.shape, dtype=sample.dtype)
    for lazy_reader in lazy_arrays
]

segs = da.stack(dask_arrays, axis=0)


p = 0
with nd2.ND2File('../timelapse.nd2') as img:
    ximg = img.to_xarray(delayed = True)
    track_graph, masks_tracked = model.track(ximg.sel(C = 'Cy5').isel(P = p).data, segs[:, :, :], mode="ilp")  # or mode="ilp", or "greedy_nodiv"

# Relabel the masks and write to cell tracking challenge format (CTC), 
ctc_tracks, ctc_masks = graph_to_ctc(
    track_graph,
    masks_tracked,
    outdir=Path("tracks_{}".format(ximg.coords['P'][p].item())),
)

And then it stalls with the following error when using mode=greedy:

Greedy mode stall output

Everything moves quickly until it hits 97% and then it stalls at this point till I kill the process (it's pegging all threads on my Threadripper at 1% or so).

/home/tnagy/.local/share/trackastra/models/general_2d already downloaded, skipping.
INFO:trackastra.model.model:Loading model state from /home/tnagy/.local/share/trackastra/models/general_2d/model.pt
INFO:trackastra.model.model_api:Using device cuda
INFO:trackastra.model.model_api:Default batch size = 4 for model on cuda.
INFO:trackastra.model.model_api:Predicting weights for candidate graph
INFO:trackastra.data.wrfeat:Extracting features from 1441 frames.
INFO:trackastra.data.wrfeat:Using single process for feature extraction
Extracting features: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1441/1441 [04:38<00:00,  5.18it/s]
INFO:trackastra.model.model_api:Building windows
Building windows: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1438/1438 [00:00<00:00, 7360.70it/s]
INFO:trackastra.model.model_api:Predicting windows with batch size 4
Computing associations: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 360/360 [00:38<00:00,  9.41it/s]
INFO:trackastra.model.model_api:Running greedy tracker
INFO:trackastra.tracking.tracking:Build candidate graph with delta_t=1
INFO:trackastra.tracking.tracking:Added 377122 vertices, 368811 edges                                                                                                                                          
INFO:trackastra.tracking.tracking:Running greedy tracker
Greedily matched edges:  97%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍   | 358952/368811 [00:03<00:00, 115164.27it/s]

And then it stalls for a long time (>24 hours) when using mode=ilp and then crashes during graph_to_ctc:

ILP error
/home/tnagy/.local/share/trackastra/models/general_2d already downloaded, skipping.
INFO:trackastra.model.model:Loading model state from /home/tnagy/.local/share/trackastra/models/general_2d/model.pt
INFO:trackastra.model.model_api:Using device cuda
INFO:trackastra.model.model_api:Default batch size = 4 for model on cuda.
INFO:trackastra.model.model_api:Predicting weights for candidate graph
INFO:trackastra.data.wrfeat:Extracting features from 1441 frames.
INFO:trackastra.data.wrfeat:Using single process for feature extraction
Extracting features: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 1441/1441 [04:51<00:00,  4.95it/s]
INFO:trackastra.model.model_api:Building windows
Building windows: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 1438/1438 [00:00<00:00, 20071.38it/s]
INFO:trackastra.model.model_api:Predicting windows with batch size 4
Computing associations: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 360/360 [00:38<00:00,  9.46it/s]
INFO:trackastra.model.model_api:Running greedy tracker
INFO:trackastra.tracking.tracking:Build candidate graph with delta_t=1
INFO:trackastra.tracking.tracking:Added 377122 vertices, 368811 edges                                                                                                 
INFO:trackastra.tracking.ilp:Using `gt` ILP config.
INFO:motile.solver:Adding NodeSelectedCost cost...
INFO:motile.solver:Adding NodeSelected variables...
INFO:motile.solver:Adding EdgeSelectedCost cost...
INFO:motile.solver:Adding EdgeSelected variables...
INFO:motile.solver:Adding NodeAppearCost cost...
INFO:motile.solver:Adding NodeAppear variables...
INFO:motile.solver:Adding NodeDisappearCost cost...
INFO:motile.solver:Adding NodeDisappear variables...
INFO:motile.solver:Adding NodeSplitCost cost...
INFO:motile.solver:Adding NodeSplit variables...
INFO:motile.solver:Adding MaxParents constraint...
INFO:motile.solver:Adding MaxChildren constraint...
INFO:motile.solver:Computing costs...
INFO:motile.solver:ILP solver returned with: OPTIMAL

Candidate graph         377122 nodes    368811 edges
Solution graph          377122 nodes    361969 edges
Converting graph to CTC results:   0%|                                                                                                      | 0/24949 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/tiogapool/data/2026/20260721_MyoG_ntdT_pTN73_diff_dll1_dapt_ctrl/tracks/nd2-trackastra.py", line 41, in <module>
    ctc_tracks, ctc_masks = graph_to_ctc(
                            ^^^^^^^^^^^^^
  File "/tiogapool/data/2026/20260721_MyoG_ntdT_pTN73_diff_dll1_dapt_ctrl/tracks/.venv/lib/python3.12/site-packages/trackastra/tracking/utils.py", line 336, in graph_to_ctc
    ss = regions[t][lab]
         ~~~~~~~~~~^^^^^
KeyError: np.int32(1)

Any suggestions for what to try would be very helpful!

The files are pretty chunky, so I'm happy to help diagnosis, otherwise I'm can try to slim the data down into a more manageable size to share with you all via email.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions