Fix Dask mask relabeling for large datasets by bentaculum · Pull Request #73 · weigertlab/trackastra

bentaculum · 2026-09-02T00:18:37Z

Summary

replace per-object Dask mask mutation with lazy, framewise sparse relabeling
reuse the relabeling path for CTC conversion and make CTC label validation Dask-safe
fix the misleading incomplete greedy progress bar and log the mask-application phase
add focused NumPy/Dask regression tests covering laziness, chunking, CTC conversion, and missing labels

Addresses #72.

Full-scale simulation

A synthetic run matching the reported dimensions completed successfully:

masks: 1,441 x 2,304 x 4,096, int32 (50.66 GiB logical size)
500 objects/frame; 720,500 graph nodes and 720,000 edges
total: 501.447 s (8m 21.4s)
CTC relabeling and full validation: 498.191 s
peak RSS: 1.02 GiB
source-frame loads: 1,441, exactly once per frame

The old path would request 1,001 source loads per frame at 500 objects. The simulation isolates post-tracking relabeling and CTC validation; it excludes model inference, ILP solving, real TIFF/ND2 I/O, and output writing.

Validation

python -m pytest -q tests/test_tracking_utils.py (5 passed)
existing real-model 2D/3D greedy inference tests passed
existing empty-frame, empty-window, and empty-sequence greedy tests passed

tlnagy · 2026-09-02T01:33:31Z

This PR definitely fixes the primary problem with the greedy linking:

0.5.6.dev2+g933763790
/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 10 processes for feature extraction
Extracting features: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1441/1441 [02:27<00:00,  9.76it/s]
INFO:trackastra.model.model_api:Building windows
Building windows: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1438/1438 [00:00<00:00, 24320.10it/s]
INFO:trackastra.model.model_api:Predicting windows with batch size 4
Computing associations: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 360/360 [00:37<00:00,  9.65it/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: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 358952/358952 [00:02<00:00, 147138.40it/s]
INFO:trackastra.model.model_api:Applying solution graph to masks
Converting graph to CTC results: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18030/18030 [00:00<00:00, 81715.70it/s]
Saving masks:   1%|█                                                                                                                                                       | 10/1441 [02:36<6:16:24, 15.78s/it]

Looks like it successfully gets past the linking, but the mask saving slows down tremendously...approximately 15s to save a single mask (6+ hours to save all of them). The masks appear correct though.

I assume this is an unintended side effect of this PR? I would assume that saving the masks should be quite speedy.

tlnagy · 2026-09-02T17:20:26Z

I ran the mode=ilp version overnight and took almost 6.5 hours to save the masks

0.5.6.dev2+g933763790
/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 10 processes for feature extraction
Extracting features: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1441/1441 [02:20<00:00, 10.27it/s]
INFO:trackastra.model.model_api:Building windows
Building windows: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1438/1438 [00:00<00:00, 24551.46it/s]
INFO:trackastra.model.model_api:Predicting windows with batch size 4
Computing associations: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 360/360 [00:37<00:00,  9.65it/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
INFO:trackastra.model.model_api:Applying solution graph to masks
Converting graph to CTC results: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 24949/24949 [00:00<00:00, 99743.97it/s]
Saving masks: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1441/1441 [6:32:19<00:00, 16.34s/it]

The actual tracking/linking part was super fast.

bentaculum · 2026-09-02T23:31:36Z

@tlnagy I took another look and hopefully fixed the writer issues you reported.

On a full-scale writer benchmark on a simulated dataset of the size you report (1,441 × 2,304 × 4,096, ~500 masks/frame), using lazy Dask masks, I now get:

CTC (check=False, fastremap): 663.3 s / 11m03s (0.460 s/frame)
GEFF (chunks (1, 512, 512)): 590.5 s / 9m50s (0.410 s/frame)

Some practical takeaways:

use NumPy arrays for fastest tracking when the complete image and mask videos fit in RAM; use one delayed Dask chunk per frame for bounded-memory operation, accepting slower per-frame processing
For CTC, check=False skips the optional validation pass when write speed matters.

tlnagy · 2026-09-03T00:24:44Z

Hi @bentaculum, that seems to fix it!

0.5.6.dev8+g0cf524cbf
/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 10 processes for feature extraction
Extracting features: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1441/1441 [05:03<00:00,  4.76it/s]
INFO:trackastra.model.model_api:Building windows
Building windows: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1438/1438 [00:00<00:00, 11411.88it/s]
INFO:trackastra.model.model_api:Predicting windows with batch size 4
Computing associations: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 360/360 [00:40<00:00,  8.90it/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, 375312 edges                                                                                                                                          
INFO:trackastra.tracking.tracking:Running greedy tracker
Greedily matched edges: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 359755/359755 [00:03<00:00, 105675.03it/s]
INFO:trackastra.model.model_api:Applying solution graph to masks
Converting graph to CTC results: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20013/20013 [00:00<00:00, 88035.95it/s]
Saving masks: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1441/1441 [04:03<00:00,  5.92it/s]

Overall, the performance is really great now! The whole process takes just ~10 minutes to run 🚀

Fix Dask mask relabeling during tracking

9337637

bentaculum marked this pull request as ready for review September 2, 2026 00:24

bentaculum mentioned this pull request Sep 2, 2026

Trackastra stalls out during linking #72

Closed

Speed up saving Dask CTC masks

dd05b87

bentaculum added 5 commits September 2, 2026 10:59

Accelerate CTC relabeling with fastremap

55dc5fe

Stream Dask masks during GEFF export

9fbcc30

Use time-one chunks for GEFF masks

d19cb79

Format GEFF streaming changes

c71852a

Match pre-commit Ruff formatting

0cf524c

bentaculum added 2 commits September 3, 2026 17:36

Remove unrelated greedy progress test

8e343b1

Clarify CTC label validation comment

6e38aae

bentaculum merged commit aa57a95 into main Sep 4, 2026
10 checks passed

Navigation Menu

Repository navigation

Fix Dask mask relabeling for large datasets - #73

Fix Dask mask relabeling for large datasets#73
bentaculum merged 9 commits into
mainfrom
investigate-issue-72-memory

bentaculum commented Sep 2, 2026

Uh oh!

tlnagy commented Sep 2, 2026 •

edited

Loading

Uh oh!

tlnagy commented Sep 2, 2026

Uh oh!

bentaculum commented Sep 2, 2026 •

edited

Loading

Uh oh!

tlnagy commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

bentaculum commented Sep 2, 2026

Summary

Full-scale simulation

Validation

Uh oh!

tlnagy commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tlnagy commented Sep 2, 2026

Uh oh!

bentaculum commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tlnagy commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tlnagy commented Sep 2, 2026 •

edited

Loading

bentaculum commented Sep 2, 2026 •

edited

Loading