Add Trackastra and Ultrack cell-tracking workers#150
Conversation
Both workers link existing tagged polygon segmentations across time into lineage tracks and upload the result as parent-child connections (like connect_timelapse), handling cell divisions as one parent to two daughters. - trackastra: transformer-based tracking (PyTorch, GPU with CPU fallback); rasterizes tagged polygons + loads channel intensity per frame, runs Trackastra.track, maps the track graph back to annotations by (time, label). - ultrack: global-optimization tracking (CPU, CBC solver); rasterizes tagged polygons, runs Ultrack on the label stack, maps tracks back by nearest mask centroid, and derives division links from parent_track_id. Each worker has an entrypoint, Dockerfile, environment.yml, docs, and a pytest suite (20 tests total) covering rasterization, graph/tracks-to-connections mapping, divisions, and end-to-end compute with the tracker mocked. Registered in build_machine_learning_workers.sh and REGISTRY.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0185y3NmemLJQdz5pY1yYQPQ
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: db848ef451
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| }, | ||
| 'Tracking mode': { | ||
| 'type': 'select', | ||
| 'items': ['greedy', 'greedy_nodiv', 'ilp'], |
There was a problem hiding this comment.
If a user selects ilp, the shipped Trackastra image only installs base trackastra (workers/annotations/trackastra/Dockerfile uses pip install trackastra) and does not install the optional ILP/solver dependencies required by Trackastra's ILP path, so model.track(..., mode='ilp') will fail at runtime instead of creating connections. Either install the ILP extras and solver in the image or omit this choice from the worker interface.
Useful? React with 👍 / 👎.
| if parent_track_id is None or parent_track_id < 0: | ||
| continue |
There was a problem hiding this comment.
When Ultrack emits parent_track_id == 0 for root tracks (the comment above says some versions use this), this condition treats zero as a real parent. If a track with ID 0 exists, founder tracks can be connected as division children of track 0, corrupting the lineage; normalize 0 to -1 before storing it or skip <= 0 here if zero means no parent.
Useful? React with 👍 / 👎.
- trackastra: drop the "ilp" tracking mode from the interface; the ILP
solver stack (ilpy + SCIP/Gurobi) is not installed in the image, so
selecting it would crash at runtime. Only greedy modes are exposed.
- Both workers now always stamp a descriptive connection tag ("Trackastra" /
"Ultrack") in addition to any configured output tags, so tracking
connections are identifiable and filterable in the UI (matching the
convention of connect_timelapse / sam2_video).
- Remove the misleading annotationShape="polygon" Docker label from both
workers; they create parent-child connections, not polygon annotations.
- ultrack: remove the broad try/except around config assignments so a real
misconfiguration surfaces instead of silently ignoring user settings; the
field names are correct against the current Ultrack API.
- ultrack: collapse the double tracks_df.iterrows() pass into one and drop the
unused has_z variable.
- REGISTRY.md: revert an unrelated Cellpose-SAM row change that the registry
regeneration swept in, keeping this PR focused.
- Document tracking limitations (non-overlapping instance masks, out-of-bounds
polygons, nearest-centroid mapping, soft division penalty) in both docs.
- Add tests asserting the default descriptive connection tag.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0185y3NmemLJQdz5pY1yYQPQ
Codex flagged that the division-link guard skipped only parent_track_id < 0, so a founder track whose parent_track_id is 0 (a "no parent" sentinel some Ultrack versions emit) would be wired as a child of track 0 if such a track existed, corrupting the lineage. Ultrack track ids are positive, so anything <= 0 means no parent. Fix the guard and add a regression test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0185y3NmemLJQdz5pY1yYQPQ
2 participants
Summary
Adds two new NimbusImage annotation workers that link existing tagged polygon segmentations across time into lineage tracks and upload the result as parent–child connections — the same integration model as
connect_timelapse. Both capture cell divisions as one parent linked to two daughters. You segment first (Cellpose, Stardist, etc.), tag the objects, then run one of these to build the lineage.Shared pipeline: load tagged polygons → group by XY/Z → rasterize each time frame into an integer label mask → run the tracker → map tracks back to the original annotation IDs → create connections.
workers/annotations/trackastra/) — transformer-based tracking on PyTorch (GPU with CPU fallback). Also loads the chosen intensity channel per frame, runsTrackastra.track, and maps the returned track graph back to annotations by each node's(time, label). Modes:greedy(default, allows divisions) andgreedy_nodiv.workers/annotations/ultrack/) — global-optimization tracking on CPU (open-source COIN-OR CBC solver, no Gurobi required). Runs Ultrack on the label stack and maps tracks back to annotations by nearest mask centroid, deriving division links fromparent_track_id. Exposes max linking distance and an allow-divisions toggle.Both connection sets always carry a descriptive tag (
Trackastra/Ultrack) so they can be filtered and managed as a group in the UI.What's included per worker
entrypoint.py,Dockerfile,environment.yml, aWORKERNAME.mddoc, and a pytest suite. Trackastra also bakes its pretrainedgeneral_2dmodel into the image. Both are registered inbuild_machine_learning_workers.shandREGISTRY.md(annotation workers 26 → 28).Testing
compute.Trackastra.trackreturns(graph, masks); node attributestime/label; Ultracktrack(labels=…, overwrite=True)andto_tracks_layer(include_parents=True)withtrack_id/t/y/x/parent_track_id; config attribute paths).Not yet exercised (require a GPU/Docker build and live data, unavailable in the authoring environment): the Docker image builds themselves and an end-to-end tracking run against real NimbusImage data.
Code review incorporated
ilpmode — the ILP solver stack (ilpy+ SCIP/Gurobi) is not installed, so selecting it would crash at runtime.annotationShape="polygon"Docker label (these produce connections, not polygons).try/exceptaround Ultrack config assignment so real misconfiguration surfaces.tracks_df.iterrows()pass into one and dropped an unused variable.REGISTRY.mdregeneration swept in.Limitations
Each frame is rasterized to a single label image, so objects are expected to be non-overlapping instance segmentations within the image bounds; overlapping or out-of-bounds objects may be dropped. Ultrack's map-back is nearest-centroid with no distance cutoff, so it assumes distinct per-frame centroids. "Forbid divisions" is a strong solver penalty, not a hard constraint.
🤖 Generated with Claude Code
https://claude.ai/code/session_0185y3NmemLJQdz5pY1yYQPQ