Add Trackastra and Ultrack cell-tracking workers

REGISTRY.md

@@ -7,13 +7,13 @@ Auto-generated by `generate_worker_docs.py` -- do not edit manually.
 
 | Category | Count |
 |----------|------:|
-| Annotation Workers | 26 |
+| Annotation Workers | 28 |
 | Property Workers -- Blobs | 10 |
 | Property Workers -- Points | 8 |
 | Property Workers -- Lines | 3 |
 | Property Workers -- Connections | 2 |
 | Test / Sample Workers | 7 |
-| **Total** | **56** |
+| **Total** | **58** |
 
 ## Annotation Workers
 
@@ -47,6 +47,8 @@ Create new annotations by segmenting images or connecting existing annotations.
 | SAM automatic mask generator | Uses SAM to find all masks in the image | Yes |  | [docs](workers/annotations/sam_automatic_mask_generator/SAM_AUTOMATIC_MASK_GENERATOR.md) |
 | SAM few-shot segmentation | Uses SAM1 ViT-H features for few-shot segmentation based on training annotations | Yes | Yes | [docs](workers/annotations/sam_fewshot_segmentation/SAM_FEWSHOT_SEGMENTATION.md) |
 | Stardist | Uses Stardist to find cells and nuclei | Yes | Yes | [docs](workers/annotations/stardist/STARDIST.md) |
+| Trackastra tracking | This tool links existing segmentation objects across time into tracks using the Trackas... | Yes | Yes | [docs](workers/annotations/trackastra/TRACKASTRA.md) |
+| Ultrack tracking | This tool links existing segmentation objects across time into tracks using Ultrack, wh... |  | Yes | [docs](workers/annotations/ultrack/ULTRACK.md) |
 
 ## Property Workers -- Blobs
 

build_machine_learning_workers.sh

@@ -61,6 +61,12 @@ docker build . -f ./workers/annotations/sam2_video/$DOCKERFILE -t annotations/sa
 echo "Building CondensateNet worker"
 docker build . -f ./workers/annotations/condensatenet/$DOCKERFILE -t annotations/condensatenet:latest $NO_CACHE
 
+echo "Building Trackastra worker"
+docker build . -f ./workers/annotations/trackastra/Dockerfile -t annotations/trackastra:latest $NO_CACHE
+
+echo "Building Ultrack worker"
+docker build . -f ./workers/annotations/ultrack/Dockerfile -t annotations/ultrack:latest $NO_CACHE
+
 # These are some legacy workers that are no longer used.
 #docker build ./workers/annotations/cellori_segmentation/ -t annotations/cellori_segmentation_worker:latest --label isUPennContrastWorker --label isAnnotationWorker --label "interfaceName=Cellori" --label "interfaceCategory=Cellori"
 #docker build ./workers/annotations/deepcell/ -t annotations/deepcell_worker:latest --label isUPennContrastWorker --label isAnnotationWorker --label "interfaceName=DeepCell" --label "interfaceCategory=Deepcell"

workers/annotations/trackastra/Dockerfile

@@ -0,0 +1,71 @@
+FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04 as base
+LABEL isUPennContrastWorker=True
+LABEL com.nvidia.volumes.needed="nvidia_driver"
+
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -qy tzdata && \
+  apt-get install -qy software-properties-common python3-software-properties && \
+  apt-get update && apt-get install -qy \
+  build-essential \
+  wget \
+  python3 \
+  libffi-dev \
+  libssl-dev \
+  libjpeg-dev \
+  zlib1g-dev \
+  git \
+  libpython3-dev && \
+  apt-get clean && rm -rf /var/lib/apt/lists/*
+
+ENV PATH="/root/miniforge3/bin:$PATH"
+ARG PATH="/root/miniforge3/bin:$PATH"
+
+RUN wget \
+    https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh \
+    && mkdir /root/.conda \
+    && bash Miniforge3-Linux-x86_64.sh -b \
+    && rm -f Miniforge3-Linux-x86_64.sh
+
+FROM base as build
+
+# Copy and create conda environment
+COPY ./workers/annotations/trackastra/environment.yml /
+RUN conda env create --file /environment.yml
+SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"]
+
+# Install NimbusImage annotation client
+RUN git clone https://github.com/arjunrajlaboratory/NimbusImage/
+RUN pip install -r /NimbusImage/devops/girder/annotation_client/requirements.txt
+RUN pip install -e /NimbusImage/devops/girder/annotation_client/
+
+# Install annotation_utilities
+COPY ./annotation_utilities /annotation_utilities
+RUN pip install /annotation_utilities
+
+# Install PyTorch (CUDA 12.1) and Trackastra. Installing torch from the CUDA
+# wheel index first ensures the GPU build is used rather than the CPU default.
+RUN pip install torch --index-url https://download.pytorch.org/whl/cu121
+RUN pip install trackastra
+
+# Pre-download the pretrained model (best-effort; falls back to runtime download)
+COPY ./workers/annotations/trackastra/download_models.py /
+RUN python /download_models.py || echo "Model pre-download skipped; will download at runtime"
+
+# Copy entrypoint
+COPY ./workers/annotations/trackastra/entrypoint.py /
+
+WORKDIR /
+
+LABEL isUPennContrastWorker="" \
+      isAnnotationWorker="" \
+      interfaceName="Trackastra tracking" \
+      interfaceCategory="Trackastra" \
+      description="Links objects across time into tracks using the Trackastra transformer model" \
+      defaultToolName="Trackastra tracking"
+
+# Fast env activation in place of `conda run` (see todo/worker-startup-latency.md)
+COPY ./workers/base_docker_images/run_worker.sh /usr/local/bin/run_worker.sh
+RUN chmod +x /usr/local/bin/run_worker.sh
+ENTRYPOINT ["/usr/local/bin/run_worker.sh", "/entrypoint.py"]

workers/annotations/trackastra/TRACKASTRA.md

@@ -0,0 +1,50 @@
+# Trackastra Worker
+
+This worker links **existing** segmentation objects across time into tracks using [Trackastra](https://github.com/weigertlab/trackastra), a transformer-based cell tracking model. It does not create new segmentations — you segment your objects first (e.g. with Cellpose, Cellpose-SAM, or Stardist), tag them, then run this tool to build the lineage as parent-child connections. Cell divisions are captured as one parent linked to two daughters.
+
+## How It Works
+
+1. **Load objects**: Fetches all polygon annotations carrying the selected tag.
+2. **Group by position**: Objects are grouped by XY position and Z slice; each stack is tracked independently across time.
+3. **Rasterize to label masks**: For each time point, the tagged polygons are drawn into an integer label image, giving a `(T, H, W)` mask stack. A `(label, time) → annotation id` map is retained.
+4. **Load intensity images**: The raw image for the selected channel is loaded for each time point, giving a matching `(T, H, W)` image stack.
+5. **Track**: The pretrained Trackastra model links detections across frames, producing a directed track graph (a node per detection, edges from a detection to its successor, and two out-edges where a cell divides).
+6. **Create connections**: Each graph edge is mapped back to the original annotations and uploaded as a parent-child connection.
+
+## Interface Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| **Tag of objects to track** | tags | -- | Tag identifying which polygon annotations to track |
+| **Channel** | channel | -- | Intensity channel the model looks at (should match the channel the objects were segmented on) |
+| **Model** | select | general_2d | Pretrained Trackastra model |
+| **Tracking mode** | select | greedy | `greedy` (allows divisions) or `greedy_nodiv` (no divisions) |
+| **Batch XY** | text | current XY | XY positions to track, e.g. `1-3, 5`. Blank = current position |
+| **Batch Z** | text | current Z | Z slices to track, e.g. `1-3, 5`. Blank = current slice |
+
+## Implementation Details
+
+### Coordinate Handling
+
+Polygon coordinates are converted from Girder's top-left-origin convention to scikit-image pixel centers (the 0.5 offset) when rasterizing. Rows correspond to `y`, columns to `x`. The tracking canvas uses the dataset's full image dimensions (`sizeX`, `sizeY`).
+
+### Mapping Tracks Back to Annotations
+
+Each mask label is unique within its frame and recorded against the annotation it came from. Trackastra graph nodes carry `time` (frame index) and `label`, so every node maps directly back to an annotation. Connections are oriented parent → child by time (not by stored edge direction), and duplicate parent/child pairs are de-duplicated.
+
+### Divisions
+
+In `greedy` mode a dividing cell produces a node with two outgoing edges, which becomes two connections (parent → daughter A, parent → daughter B). Use `greedy_nodiv` to forbid divisions.
+
+### Output Tags
+
+Created connections always carry a `Trackastra` tag (in addition to any output tags configured on the tool) so the tracking result can be filtered, selected, or bulk-deleted as a group in the UI.
+
+## Notes
+
+- **GPU**: Uses PyTorch and runs on GPU when available, falling back to CPU otherwise.
+- **Model weights** are pre-downloaded into the image (best-effort at build time; otherwise fetched on first run).
+- Trackastra also ships an `ilp` linking mode, but it requires an ILP solver (`ilpy` + SCIP/Gurobi) that is not installed in this image, so only the greedy modes are exposed.
+- Only polygon annotations are tracked, and each frame is rasterized into a single label image, so objects are expected to be **non-overlapping** instance segmentations. Where two tagged polygons overlap at the same XY/Z/Time, the later-drawn one wins and a fully occluded object is dropped from tracking. Polygons that fall entirely outside the image bounds rasterize to nothing and are likewise skipped.
+- Objects at a single time point (no second frame in their XY/Z stack) cannot form a track and are skipped.
+- Related workers: **Ultrack** (optimization-based tracking of the same segmentations), **Connect Time Lapse** and **Connect Sequential** (simpler nearest-neighbor linking), **SAM2 video** (segment-and-track in one step).

workers/annotations/trackastra/download_models.py

@@ -0,0 +1,14 @@
+"""Pre-download the pretrained Trackastra model into the image.
+
+Baking the weights into the Docker image avoids a network download on the
+first run (NimbusImage launches a fresh container per job). Best-effort: if the
+model host is unreachable at build time, the build continues and the model is
+downloaded on first use instead.
+"""
+
+from trackastra.model import Trackastra
+
+if __name__ == '__main__':
+    # device="cpu" is used at build time since no GPU is present during build.
+    Trackastra.from_pretrained("general_2d", device="cpu")
+    print("Downloaded Trackastra general_2d model")