diff --git a/REGISTRY.md b/REGISTRY.md index 57de4e9..09642e6 100644 --- a/REGISTRY.md +++ b/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 diff --git a/build_machine_learning_workers.sh b/build_machine_learning_workers.sh index 935ebdc..d6f5542 100755 --- a/build_machine_learning_workers.sh +++ b/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" diff --git a/workers/annotations/trackastra/Dockerfile b/workers/annotations/trackastra/Dockerfile new file mode 100644 index 0000000..1433f05 --- /dev/null +++ b/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"] diff --git a/workers/annotations/trackastra/TRACKASTRA.md b/workers/annotations/trackastra/TRACKASTRA.md new file mode 100644 index 0000000..80ea1d5 --- /dev/null +++ b/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). diff --git a/workers/annotations/trackastra/download_models.py b/workers/annotations/trackastra/download_models.py new file mode 100644 index 0000000..49f2b1a --- /dev/null +++ b/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") diff --git a/workers/annotations/trackastra/entrypoint.py b/workers/annotations/trackastra/entrypoint.py new file mode 100644 index 0000000..6cdc80f --- /dev/null +++ b/workers/annotations/trackastra/entrypoint.py @@ -0,0 +1,337 @@ +import argparse +import json +import sys +import timeit + +import annotation_client.annotations as annotations +import annotation_client.workers as workers +import annotation_client.tiles as tiles +from annotation_client.utils import sendProgress, sendError, sendWarning + +import annotation_utilities.annotation_tools as annotation_tools +import annotation_utilities.batch_argument_parser as batch_argument_parser + +import numpy as np + + +def interface(image, apiUrl, token): + client = workers.UPennContrastWorkerPreviewClient( + apiUrl=apiUrl, token=token) + + interface = { + 'Trackastra tracking': { + 'type': 'notes', + 'value': 'This tool links existing segmentation objects across time ' + 'into tracks using the Trackastra transformer tracking model. ' + 'It creates parent-child connections between objects, including ' + 'cell divisions (one parent linked to two daughters). ' + 'Segment your objects first (e.g. with Cellpose or Stardist), ' + 'tag them, then run this tool to build the lineage.', + 'displayOrder': 0, + }, + 'Tag of objects to track': { + 'type': 'tags', + 'tooltip': 'Track all objects that have this tag.', + 'displayOrder': 1, + }, + 'Channel': { + 'type': 'channel', + 'required': True, + 'tooltip': 'The intensity channel the tracking model looks at.\n' + 'This should be the channel the objects were segmented on.', + 'displayOrder': 2, + }, + 'Model': { + 'type': 'select', + 'items': ['general_2d'], + 'default': 'general_2d', + 'tooltip': 'Pretrained Trackastra model to use.', + 'displayOrder': 3, + }, + 'Tracking mode': { + 'type': 'select', + 'items': ['greedy', 'greedy_nodiv'], + 'default': 'greedy', + 'tooltip': 'Linking strategy. "greedy" allows divisions, ' + '"greedy_nodiv" forbids them.', + 'displayOrder': 4, + }, + 'Batch XY': { + 'type': 'text', + 'tooltip': 'Comma-separated list of XY positions to track, e.g. "1-3, 5".\n' + 'Leave blank to use the current XY position only.', + 'displayOrder': 5, + }, + 'Batch Z': { + 'type': 'text', + 'tooltip': 'Comma-separated list of Z slices to track, e.g. "1-3, 5".\n' + 'Leave blank to use the current Z slice only.', + 'displayOrder': 6, + }, + } + client.setWorkerImageInterface(image, interface) + + +def group_annotations_by_location(annotation_list): + """Group polygon annotations by (XY, Z) so each stack is tracked independently. + + Returns a dict mapping (XY, Z) -> list of annotations. + """ + groups = {} + for ann in annotation_list: + if ann.get('shape') != 'polygon': + continue + if len(ann.get('coordinates', [])) < 3: + continue + loc = ann['location'] + key = (loc['XY'], loc['Z']) + groups.setdefault(key, []).append(ann) + return groups + + +def build_label_stack(group_annotations, time_points, height, width): + """Rasterize polygon annotations into a (T, H, W) integer label stack. + + Each annotation is drawn with a unique per-frame label. Returns: + masks: (T, H, W) int32 label image stack + label_to_id: {(t_index, label): annotation_id} + label_centroid: {(t_index, label): (row, col)} centroid in image space + """ + # Lazy import: keeps scikit-image off the interface path; only needed during compute. + from skimage.draw import polygon as sk_polygon + + time_to_index = {t: i for i, t in enumerate(time_points)} + masks = np.zeros((len(time_points), height, width), dtype=np.int32) + label_to_id = {} + label_centroid = {} + + # Assign labels per time frame. + next_label = {t: 1 for t in time_points} + for ann in group_annotations: + t = ann['location']['Time'] + if t not in time_to_index: + continue + t_idx = time_to_index[t] + coords = ann['coordinates'] + if len(coords) < 3: + continue + # The 0.5 offset converts Girder (top-left origin) coordinates to + # scikit-image pixel-center coordinates. Rows are y, columns are x. + rr_poly = np.array([c['y'] - 0.5 for c in coords]) + cc_poly = np.array([c['x'] - 0.5 for c in coords]) + rr, cc = sk_polygon(rr_poly, cc_poly, shape=(height, width)) + if len(rr) == 0: + continue + label = next_label[t] + next_label[t] += 1 + masks[t_idx, rr, cc] = label + label_to_id[(t_idx, label)] = ann['_id'] + label_centroid[(t_idx, label)] = (float(rr.mean()), float(cc.mean())) + + return masks, label_to_id, label_centroid + + +def track_graph_to_connections(track_graph, label_to_id, datasetId, tags): + """Convert a Trackastra track graph into parent-child connection dicts. + + Trackastra returns a directed graph whose nodes carry a ``time`` (frame + index) and ``label`` (label id within that frame's mask). Each edge links a + detection to its successor; a node with two out-edges is a division. We map + each node back to the original annotation id via ``label_to_id`` and emit a + connection for every edge whose endpoints both map to real annotations. + """ + node_to_id = {} + node_time = {} + for node, data in track_graph.nodes(data=True): + t = data.get('time', data.get('t')) + label = data.get('label', data.get('index')) + node_time[node] = t + key = (t, label) + if key in label_to_id: + node_to_id[node] = label_to_id[key] + + connections = [] + seen = set() + for u, v in track_graph.edges(): + if u not in node_to_id or v not in node_to_id: + continue + tu = node_time.get(u) + tv = node_time.get(v) + # Orient parent -> child by time; skip same-frame or ambiguous edges. + if tu is None or tv is None or tu == tv: + continue + parent, child = (u, v) if tu < tv else (v, u) + parent_id = node_to_id[parent] + child_id = node_to_id[child] + if parent_id == child_id: + continue + pair = (parent_id, child_id) + if pair in seen: + continue + seen.add(pair) + connections.append({ + 'datasetId': datasetId, + 'parentId': parent_id, + 'childId': child_id, + 'tags': tags, + }) + + return connections + + +def run_trackastra(imgs, masks, model_name, mode): + """Load the pretrained Trackastra model and run tracking on a stack. + + Isolated so tests can patch it without importing torch/trackastra. + """ + # Lazy imports: keep torch/trackastra off the interface/startup path. + import torch + from trackastra.model import Trackastra + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Running Trackastra ({model_name}, mode={mode}) on device {device}") + model = Trackastra.from_pretrained(model_name, device=device) + result = model.track(imgs, masks, mode=mode) + # Newer Trackastra returns (track_graph, masks_tracked); older returns just + # the graph. Accept either. + if isinstance(result, tuple): + return result[0] + return result + + +def compute(datasetId, apiUrl, token, params): + """ + params (could change): + assignment, channel, connectTo, tags, tile, workerInterface + """ + start_time = timeit.default_timer() + + workerInterface = params['workerInterface'] + track_tags = list(set(workerInterface.get('Tag of objects to track', []) or [])) + channel = workerInterface.get('Channel', params.get('channel', 0)) + model_name = workerInterface.get('Model', 'general_2d') + mode = workerInterface.get('Tracking mode', 'greedy') + + if not track_tags: + sendError("No tag specified", + info="Please select at least one tag of objects to track.") + raise ValueError("No tag specified") + + tile = params['tile'] + # Always stamp a descriptive tag so tracking connections are identifiable + # (and filterable / bulk-selectable) in the UI even when the user set no + # output tags -- matching the convention of the other connection workers. + output_tags = list(dict.fromkeys((params.get('tags') or []) + ["Trackastra"])) + + batch_xy = batch_argument_parser.process_range_list( + workerInterface.get('Batch XY', None), convert_one_to_zero_index=True) + batch_z = batch_argument_parser.process_range_list( + workerInterface.get('Batch Z', None), convert_one_to_zero_index=True) + if batch_xy is None: + batch_xy = [tile['XY']] + if batch_z is None: + batch_z = [tile['Z']] + batch_xy = list(batch_xy) + batch_z = list(batch_z) + + annotationClient = annotations.UPennContrastAnnotationClient( + apiUrl=apiUrl, token=token) + tileClient = tiles.UPennContrastDataset( + apiUrl=apiUrl, token=token, datasetId=datasetId) + + sendProgress(0.05, "Loading objects", "Fetching annotations") + + blobAnnotationList = annotationClient.getAnnotationsByDatasetId( + datasetId, limit=1000000, shape='polygon') + objectList = annotation_tools.get_annotations_with_tags( + blobAnnotationList, track_tags, exclusive=False) + + if not objectList: + sendWarning("No annotations found", + info="No objects with the specified tag were found to track.") + return + + # Restrict to the requested XY/Z positions and group each stack. + objectList = [ann for ann in objectList + if ann['location']['XY'] in batch_xy + and ann['location']['Z'] in batch_z] + groups = group_annotations_by_location(objectList) + + if not groups: + sendWarning("No annotations found", + info="No objects with the specified tag were found in the " + "selected XY/Z positions.") + return + + # Image dimensions define the tracking canvas. + height = tileClient.tiles['sizeY'] + width = tileClient.tiles['sizeX'] + + all_connections = [] + total_groups = len(groups) + for group_idx, ((xy, z), group_annotations) in enumerate(groups.items()): + time_points = sorted(set(ann['location']['Time'] + for ann in group_annotations)) + if len(time_points) < 2: + # A single time point cannot form a track. + continue + + masks, label_to_id, _ = build_label_stack( + group_annotations, time_points, height, width) + + # Load the intensity images for this stack, one per time point. + imgs = np.zeros((len(time_points), height, width), dtype=np.float32) + for t_idx, t in enumerate(time_points): + frame = tileClient.coordinatesToFrameIndex(xy, z, t, channel) + image = tileClient.getRegion(datasetId, frame=frame).squeeze() + imgs[t_idx] = image.astype(np.float32) + + sendProgress( + (group_idx + 0.5) / total_groups, + "Tracking objects", + f"Tracking XY {xy + 1}, Z {z + 1} ({len(time_points)} time points)") + + track_graph = run_trackastra(imgs, masks, model_name, mode) + connections = track_graph_to_connections( + track_graph, label_to_id, datasetId, output_tags) + all_connections.extend(connections) + + sendProgress( + (group_idx + 1) / total_groups, + "Tracking objects", + f"Processed {group_idx + 1} of {total_groups} stacks") + + if not all_connections: + sendWarning("No connections created", + info="Trackastra did not find any links between objects.") + return + + sendProgress(0.95, "Uploading", f"Sending {len(all_connections)} connections") + annotationClient.createMultipleConnections(all_connections) + + elapsed = timeit.default_timer() - start_time + print(f"Created {len(all_connections)} connections in {elapsed:.1f} s") + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + description='Track objects across time with Trackastra') + + parser.add_argument('--datasetId', type=str, required=False, action='store') + parser.add_argument('--apiUrl', type=str, required=True, action='store') + parser.add_argument('--token', type=str, required=True, action='store') + parser.add_argument('--request', type=str, required=True, action='store') + parser.add_argument('--parameters', type=str, required=True, action='store') + + args = parser.parse_args(sys.argv[1:]) + + params = json.loads(args.parameters) + datasetId = args.datasetId + apiUrl = args.apiUrl + token = args.token + + match args.request: + case 'compute': + compute(datasetId, apiUrl, token, params) + case 'interface': + interface(params['image'], apiUrl, token) diff --git a/workers/annotations/trackastra/environment.yml b/workers/annotations/trackastra/environment.yml new file mode 100644 index 0000000..4f7c1f8 --- /dev/null +++ b/workers/annotations/trackastra/environment.yml @@ -0,0 +1,13 @@ +name: worker +channels: + - conda-forge + - defaults +dependencies: + - python=3.10 + - pip + - numpy + - scikit-image + - shapely + - rasterio + - pip: + - rtree diff --git a/workers/annotations/trackastra/tests/Dockerfile_Test b/workers/annotations/trackastra/tests/Dockerfile_Test new file mode 100644 index 0000000..2a5da8c --- /dev/null +++ b/workers/annotations/trackastra/tests/Dockerfile_Test @@ -0,0 +1,14 @@ +# Use the existing trackastra worker as the base +FROM annotations/trackastra:latest AS test + +# Install test dependencies +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] +RUN pip install pytest pytest-mock + +# Copy test files +RUN mkdir -p /tests +COPY ./workers/annotations/trackastra/tests/*.py /tests +WORKDIR /tests + +# Command to run tests +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python3", "-m", "pytest", "-v"] diff --git a/workers/annotations/trackastra/tests/__init__.py b/workers/annotations/trackastra/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/annotations/trackastra/tests/test_trackastra.py b/workers/annotations/trackastra/tests/test_trackastra.py new file mode 100644 index 0000000..5328942 --- /dev/null +++ b/workers/annotations/trackastra/tests/test_trackastra.py @@ -0,0 +1,277 @@ +import numpy as np +import networkx as nx +import pytest +from unittest.mock import patch, MagicMock + +from entrypoint import ( + interface, + compute, + group_annotations_by_location, + build_label_stack, + track_graph_to_connections, +) + + +@pytest.fixture +def mock_worker_preview_client(): + with patch('annotation_client.workers.UPennContrastWorkerPreviewClient') as mock_client: + yield mock_client.return_value + + +@pytest.fixture +def square_annotations(): + """One 4x4 square per time point, drifting diagonally over three frames.""" + def square(_id, t, x0, y0): + return { + '_id': _id, + 'shape': 'polygon', + 'coordinates': [ + {'x': x0, 'y': y0}, + {'x': x0 + 4, 'y': y0}, + {'x': x0 + 4, 'y': y0 + 4}, + {'x': x0, 'y': y0 + 4}, + {'x': x0, 'y': y0}, + ], + 'location': {'Time': t, 'XY': 0, 'Z': 0}, + 'tags': ['cell'], + } + return [ + square('t0', 0, 2, 2), + square('t1', 1, 4, 4), + square('t2', 2, 6, 6), + ] + + +def test_interface(mock_worker_preview_client): + interface('test_image', 'http://api', 'token') + mock_worker_preview_client.setWorkerImageInterface.assert_called_once() + interface_data = mock_worker_preview_client.setWorkerImageInterface.call_args[0][1] + assert 'Tag of objects to track' in interface_data + assert 'Channel' in interface_data + assert 'Model' in interface_data + assert 'Tracking mode' in interface_data + assert interface_data['Tag of objects to track']['type'] == 'tags' + assert interface_data['Channel']['type'] == 'channel' + + +def test_group_annotations_by_location(): + anns = [ + {'shape': 'polygon', 'coordinates': [1, 2, 3], 'location': {'XY': 0, 'Z': 0, 'Time': 0}}, + {'shape': 'polygon', 'coordinates': [1, 2, 3], 'location': {'XY': 0, 'Z': 0, 'Time': 1}}, + {'shape': 'polygon', 'coordinates': [1, 2, 3], 'location': {'XY': 1, 'Z': 0, 'Time': 0}}, + # Too few coordinates -> dropped + {'shape': 'polygon', 'coordinates': [1, 2], 'location': {'XY': 0, 'Z': 0, 'Time': 2}}, + # Wrong shape -> dropped + {'shape': 'point', 'coordinates': [1], 'location': {'XY': 0, 'Z': 0, 'Time': 0}}, + ] + groups = group_annotations_by_location(anns) + assert set(groups.keys()) == {(0, 0), (1, 0)} + assert len(groups[(0, 0)]) == 2 + assert len(groups[(1, 0)]) == 1 + + +def test_build_label_stack(square_annotations): + time_points = [0, 1, 2] + masks, label_to_id, label_centroid = build_label_stack( + square_annotations, time_points, height=20, width=20) + + assert masks.shape == (3, 20, 20) + # One object per frame, each gets label 1. + assert label_to_id[(0, 1)] == 't0' + assert label_to_id[(1, 1)] == 't1' + assert label_to_id[(2, 1)] == 't2' + # The label pixels exist in each frame. + for t_idx in range(3): + assert (masks[t_idx] == 1).sum() > 0 + # Centroid of the first square (x 2..6, y 2..6) is near (4, 4) in (row, col). + row, col = label_centroid[(0, 1)] + assert abs(row - 3.5) < 1.5 + assert abs(col - 3.5) < 1.5 + + +def test_build_label_stack_skips_short_polygons(): + anns = [ + {'_id': 'bad', 'shape': 'polygon', + 'coordinates': [{'x': 0, 'y': 0}, {'x': 1, 'y': 1}], + 'location': {'Time': 0, 'XY': 0, 'Z': 0}}, + ] + masks, label_to_id, _ = build_label_stack(anns, [0], height=10, width=10) + assert masks.sum() == 0 + assert label_to_id == {} + + +def test_track_graph_to_connections_linear(): + """A -> B -> C linear track becomes two connections.""" + g = nx.DiGraph() + g.add_node(0, time=0, label=1) + g.add_node(1, time=1, label=1) + g.add_node(2, time=2, label=1) + g.add_edge(0, 1) + g.add_edge(1, 2) + label_to_id = {(0, 1): 't0', (1, 1): 't1', (2, 1): 't2'} + + conns = track_graph_to_connections(g, label_to_id, 'ds', ['Trackastra']) + pairs = {(c['parentId'], c['childId']) for c in conns} + assert pairs == {('t0', 't1'), ('t1', 't2')} + assert all(c['datasetId'] == 'ds' for c in conns) + assert all(c['tags'] == ['Trackastra'] for c in conns) + + +def test_track_graph_to_connections_division(): + """A parent dividing into two daughters yields two connections.""" + g = nx.DiGraph() + g.add_node(0, time=0, label=1) + g.add_node(1, time=1, label=1) + g.add_node(2, time=1, label=2) + g.add_edge(0, 1) + g.add_edge(0, 2) + label_to_id = {(0, 1): 'parent', (1, 1): 'childA', (1, 2): 'childB'} + + conns = track_graph_to_connections(g, label_to_id, 'ds', ['Trackastra']) + pairs = {(c['parentId'], c['childId']) for c in conns} + assert pairs == {('parent', 'childA'), ('parent', 'childB')} + + +def test_track_graph_orients_by_time_not_edge_direction(): + """Even if an edge is stored later->earlier, parent is the earlier node.""" + g = nx.DiGraph() + g.add_node(0, time=5, label=1) + g.add_node(1, time=2, label=1) + g.add_edge(0, 1) # stored high-time -> low-time + label_to_id = {(5, 1): 'late', (2, 1): 'early'} + + conns = track_graph_to_connections(g, label_to_id, 'ds', []) + assert len(conns) == 1 + assert conns[0]['parentId'] == 'early' + assert conns[0]['childId'] == 'late' + + +def test_track_graph_skips_unmapped_nodes(): + g = nx.DiGraph() + g.add_node(0, time=0, label=1) + g.add_node(1, time=1, label=99) # not in label_to_id + g.add_edge(0, 1) + label_to_id = {(0, 1): 't0'} + conns = track_graph_to_connections(g, label_to_id, 'ds', []) + assert conns == [] + + +def test_compute_no_tag_raises(): + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': [], + 'channel': 0, + 'workerInterface': { + 'Tag of objects to track': [], + 'Channel': 0, + }, + } + with pytest.raises(ValueError, match="No tag specified"): + compute('ds', 'http://api', 'token', params) + + +@patch('entrypoint.run_trackastra') +@patch('annotation_client.tiles.UPennContrastDataset') +@patch('annotation_client.annotations.UPennContrastAnnotationClient') +def test_compute_end_to_end(mock_ann_client_cls, mock_tile_cls, mock_run, + square_annotations): + ann_client = mock_ann_client_cls.return_value + ann_client.getAnnotationsByDatasetId.return_value = square_annotations + ann_client.createMultipleConnections.return_value = {} + + tile_client = mock_tile_cls.return_value + tile_client.tiles = {'sizeX': 20, 'sizeY': 20} + tile_client.coordinatesToFrameIndex.return_value = 0 + tile_client.getRegion.return_value = np.zeros((20, 20), dtype=np.float32) + + # Fake track graph: t0 -> t1 -> t2 (single object per frame, all label 1). + g = nx.DiGraph() + g.add_node(0, time=0, label=1) + g.add_node(1, time=1, label=1) + g.add_node(2, time=2, label=1) + g.add_edge(0, 1) + g.add_edge(1, 2) + mock_run.return_value = g + + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': ['Trackastra'], + 'channel': 0, + 'workerInterface': { + 'Tag of objects to track': ['cell'], + 'Channel': 0, + 'Model': 'general_2d', + 'Tracking mode': 'greedy', + 'Batch XY': '', + 'Batch Z': '', + }, + } + + compute('ds', 'http://api', 'token', params) + + ann_client.createMultipleConnections.assert_called_once() + conns = ann_client.createMultipleConnections.call_args[0][0] + pairs = {(c['parentId'], c['childId']) for c in conns} + assert pairs == {('t0', 't1'), ('t1', 't2')} + # run_trackastra was called once (single XY/Z group). + mock_run.assert_called_once() + + +@patch('entrypoint.run_trackastra') +@patch('annotation_client.tiles.UPennContrastDataset') +@patch('annotation_client.annotations.UPennContrastAnnotationClient') +def test_compute_stamps_descriptive_tag_when_no_output_tags( + mock_ann_client_cls, mock_tile_cls, mock_run, square_annotations): + """Connections carry a 'Trackastra' tag even when no output tags are set.""" + ann_client = mock_ann_client_cls.return_value + ann_client.getAnnotationsByDatasetId.return_value = square_annotations + tile_client = mock_tile_cls.return_value + tile_client.tiles = {'sizeX': 20, 'sizeY': 20} + tile_client.coordinatesToFrameIndex.return_value = 0 + tile_client.getRegion.return_value = np.zeros((20, 20), dtype=np.float32) + + g = nx.DiGraph() + g.add_node(0, time=0, label=1) + g.add_node(1, time=1, label=1) + g.add_edge(0, 1) + mock_run.return_value = g + + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': [], # no output tags configured + 'channel': 0, + 'workerInterface': { + 'Tag of objects to track': ['cell'], + 'Channel': 0, + 'Batch XY': '', + 'Batch Z': '', + }, + } + compute('ds', 'http://api', 'token', params) + + conns = ann_client.createMultipleConnections.call_args[0][0] + assert conns, "expected at least one connection" + assert all('Trackastra' in c['tags'] for c in conns) + + +@patch('entrypoint.run_trackastra') +@patch('annotation_client.tiles.UPennContrastDataset') +@patch('annotation_client.annotations.UPennContrastAnnotationClient') +def test_compute_no_annotations(mock_ann_client_cls, mock_tile_cls, mock_run): + ann_client = mock_ann_client_cls.return_value + ann_client.getAnnotationsByDatasetId.return_value = [] + + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': ['Trackastra'], + 'channel': 0, + 'workerInterface': { + 'Tag of objects to track': ['cell'], + 'Channel': 0, + 'Batch XY': '', + 'Batch Z': '', + }, + } + compute('ds', 'http://api', 'token', params) + ann_client.createMultipleConnections.assert_not_called() + mock_run.assert_not_called() diff --git a/workers/annotations/ultrack/Dockerfile b/workers/annotations/ultrack/Dockerfile new file mode 100644 index 0000000..375e345 --- /dev/null +++ b/workers/annotations/ultrack/Dockerfile @@ -0,0 +1,58 @@ +FROM ubuntu:22.04 as base +LABEL isUPennContrastWorker=True + +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 (installs Ultrack via pip) +COPY ./workers/annotations/ultrack/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 + +# Copy entrypoint +COPY ./workers/annotations/ultrack/entrypoint.py / + +WORKDIR / + +LABEL isUPennContrastWorker="" \ + isAnnotationWorker="" \ + interfaceName="Ultrack tracking" \ + interfaceCategory="Ultrack" \ + description="Links objects across time into tracks using Ultrack global optimization" \ + defaultToolName="Ultrack 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"] diff --git a/workers/annotations/ultrack/ULTRACK.md b/workers/annotations/ultrack/ULTRACK.md new file mode 100644 index 0000000..12d529d --- /dev/null +++ b/workers/annotations/ultrack/ULTRACK.md @@ -0,0 +1,50 @@ +# Ultrack Worker + +This worker links **existing** segmentation objects across time into tracks using [Ultrack](https://github.com/royerlab/ultrack), which solves a global optimization over the segmentations rather than linking frame-by-frame. 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. The centroid of each label and its `(label, time) → annotation id` mapping are retained. +4. **Track**: Ultrack ingests the label stack, derives detection/contour hypotheses, and solves a global integer linear program to produce a tracks table with `track_id`, `t`, `y`, `x`, and `parent_track_id`. +5. **Map back and connect**: Each track node is matched to the nearest original annotation (by mask centroid at that time point). Consecutive nodes within a track become motion connections, and each `parent_track_id` produces a division connection from the parent track's last node to the child track's first node. + +## Interface Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| **Tag of objects to track** | tags | -- | Tag identifying which polygon annotations to track | +| **Max distance** | number (pixels) | 50 | Maximum distance an object may move between consecutive time points to be linked | +| **Allow divisions** | checkbox | True | Allow one object to split into two (cell division). Uncheck to forbid 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 + +Ultrack relabels objects internally, so track nodes are matched back to the original annotations by nearest mask centroid at the same time point (both are expressed in `(row, col)` = `(y, x)` image space). This keeps the worker independent of Ultrack's internal id scheme. The match is nearest-neighbour with no distance cutoff, so it assumes objects within a frame have distinct centroids; two annotations sharing a centroid can be mapped ambiguously. + +### Divisions + +Ultrack encodes lineage via `parent_track_id`. When a track has a parent, a connection is created from the parent track's final node to that child track's first node — so a division yields two such connections. Unchecking **Allow divisions** applies a strong penalty to division events in the solver (a heavy discouragement in the objective, not a hard prohibition). + +### Output Tags + +Created connections always carry an `Ultrack` 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. + +### Working Directory + +Ultrack persists intermediate results to a SQLite database; each XY/Z stack is tracked inside its own temporary directory that is cleaned up afterwards. The tracks table is fully read into memory before the directory is removed. + +## Notes + +- **CPU**: Ultrack's tracking is CPU-based and uses the open-source COIN-OR CBC solver by default (Gurobi is optional and not required). +- 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; polygons entirely outside the image bounds are skipped. +- Only polygon annotations are tracked. Objects at a single time point (no second frame in their XY/Z stack) cannot form a track and are skipped. +- Related workers: **Trackastra** (transformer-based tracking of the same segmentations), **Connect Time Lapse** and **Connect Sequential** (simpler nearest-neighbor linking), **SAM2 video** (segment-and-track in one step). diff --git a/workers/annotations/ultrack/entrypoint.py b/workers/annotations/ultrack/entrypoint.py new file mode 100644 index 0000000..01d1e40 --- /dev/null +++ b/workers/annotations/ultrack/entrypoint.py @@ -0,0 +1,381 @@ +import argparse +import json +import sys +import tempfile +import timeit + +import annotation_client.annotations as annotations +import annotation_client.workers as workers +import annotation_client.tiles as tiles +from annotation_client.utils import sendProgress, sendError, sendWarning + +import annotation_utilities.annotation_tools as annotation_tools +import annotation_utilities.batch_argument_parser as batch_argument_parser + +import numpy as np + + +def interface(image, apiUrl, token): + client = workers.UPennContrastWorkerPreviewClient( + apiUrl=apiUrl, token=token) + + interface = { + 'Ultrack tracking': { + 'type': 'notes', + 'value': 'This tool links existing segmentation objects across time ' + 'into tracks using Ultrack, which solves a global optimization ' + 'over the segmentations. It creates parent-child connections ' + 'between objects, including cell divisions (one parent linked to ' + 'two daughters). Segment your objects first (e.g. with Cellpose ' + 'or Stardist), tag them, then run this tool to build the lineage.', + 'displayOrder': 0, + }, + 'Tag of objects to track': { + 'type': 'tags', + 'tooltip': 'Track all objects that have this tag.', + 'displayOrder': 1, + }, + 'Max distance': { + 'type': 'number', + 'min': 0, + 'max': 1000, + 'default': 50, + 'unit': 'pixels', + 'tooltip': 'The maximum distance (in pixels) an object may move\n' + 'between consecutive time points to be linked.', + 'displayOrder': 2, + }, + 'Allow divisions': { + 'type': 'checkbox', + 'default': True, + 'tooltip': 'Allow one object to split into two (cell division).\n' + 'Uncheck to forbid divisions.', + 'displayOrder': 3, + }, + 'Batch XY': { + 'type': 'text', + 'tooltip': 'Comma-separated list of XY positions to track, e.g. "1-3, 5".\n' + 'Leave blank to use the current XY position only.', + 'displayOrder': 4, + }, + 'Batch Z': { + 'type': 'text', + 'tooltip': 'Comma-separated list of Z slices to track, e.g. "1-3, 5".\n' + 'Leave blank to use the current Z slice only.', + 'displayOrder': 5, + }, + } + client.setWorkerImageInterface(image, interface) + + +def group_annotations_by_location(annotation_list): + """Group polygon annotations by (XY, Z) so each stack is tracked independently. + + Returns a dict mapping (XY, Z) -> list of annotations. + """ + groups = {} + for ann in annotation_list: + if ann.get('shape') != 'polygon': + continue + if len(ann.get('coordinates', [])) < 3: + continue + loc = ann['location'] + key = (loc['XY'], loc['Z']) + groups.setdefault(key, []).append(ann) + return groups + + +def build_label_stack(group_annotations, time_points, height, width): + """Rasterize polygon annotations into a (T, H, W) integer label stack. + + Each annotation is drawn with a unique per-frame label. Returns: + masks: (T, H, W) int32 label image stack + label_to_id: {(t_index, label): annotation_id} + label_centroid: {(t_index, label): (row, col)} centroid in image space + """ + # Lazy import: keeps scikit-image off the interface path; only needed during compute. + from skimage.draw import polygon as sk_polygon + + time_to_index = {t: i for i, t in enumerate(time_points)} + masks = np.zeros((len(time_points), height, width), dtype=np.int32) + label_to_id = {} + label_centroid = {} + + next_label = {t: 1 for t in time_points} + for ann in group_annotations: + t = ann['location']['Time'] + if t not in time_to_index: + continue + t_idx = time_to_index[t] + coords = ann['coordinates'] + if len(coords) < 3: + continue + # The 0.5 offset converts Girder (top-left origin) coordinates to + # scikit-image pixel-center coordinates. Rows are y, columns are x. + rr_poly = np.array([c['y'] - 0.5 for c in coords]) + cc_poly = np.array([c['x'] - 0.5 for c in coords]) + rr, cc = sk_polygon(rr_poly, cc_poly, shape=(height, width)) + if len(rr) == 0: + continue + label = next_label[t] + next_label[t] += 1 + masks[t_idx, rr, cc] = label + label_to_id[(t_idx, label)] = ann['_id'] + label_centroid[(t_idx, label)] = (float(rr.mean()), float(cc.mean())) + + return masks, label_to_id, label_centroid + + +def _match_node_to_annotation(t_idx, y, x, centroids_by_time): + """Find the annotation id whose mask centroid is nearest to (y, x) at t_idx. + + centroids_by_time maps t_index -> list of (annotation_id, row, col). + Returns the nearest annotation id, or None if there are no candidates. + """ + candidates = centroids_by_time.get(t_idx) + if not candidates: + return None + best_id = None + best_dist = None + for ann_id, row, col in candidates: + dist = (row - y) ** 2 + (col - x) ** 2 + if best_dist is None or dist < best_dist: + best_dist = dist + best_id = ann_id + return best_id + + +def tracks_df_to_connections(tracks_df, label_centroid, label_to_id, datasetId, tags): + """Convert an Ultrack tracks DataFrame into parent-child connection dicts. + + Each row of ``tracks_df`` is a detection at time ``t`` with centroid + ``(y, x)`` belonging to ``track_id``; ``parent_track_id`` links a track to + the track it divided from (-1 for founder tracks). We map each detection to + the nearest original annotation (via mask centroids), then emit connections: + + * consecutive detections within a track (motion links), and + * the last detection of a parent track to the first detection of each of + its child tracks (division links). + """ + # Index the mask centroids by time for nearest-neighbour matching. + centroids_by_time = {} + for (t_idx, label), (row, col) in label_centroid.items(): + ann_id = label_to_id.get((t_idx, label)) + if ann_id is None: + continue + centroids_by_time.setdefault(t_idx, []).append((ann_id, row, col)) + + # Assign every track node to an annotation id, and record each track's + # parent in the same pass. node_ids[track_id] -> list of (t, annotation_id). + node_ids = {} + parent_of = {} + for _, r in tracks_df.iterrows(): + track_id = int(r['track_id']) + if track_id not in parent_of: + raw_parent = r.get('parent_track_id', -1) + # Founder tracks may encode "no parent" as -1, 0, or NaN depending + # on the Ultrack version. NaN/invalid becomes -1 here; the 0 vs -1 + # distinction is handled by the "<= 0" guard on the division loop. + try: + parent_of[track_id] = int(raw_parent) + except (ValueError, TypeError): + parent_of[track_id] = -1 + t_idx = int(r['t']) + y = float(r['y']) + x = float(r['x']) + ann_id = _match_node_to_annotation(t_idx, y, x, centroids_by_time) + if ann_id is None: + continue + node_ids.setdefault(track_id, []).append((t_idx, ann_id)) + + for track_id in node_ids: + node_ids[track_id].sort(key=lambda pair: pair[0]) + + connections = [] + seen = set() + + def add_connection(parent_id, child_id): + if parent_id is None or child_id is None or parent_id == child_id: + return + pair = (parent_id, child_id) + if pair in seen: + return + seen.add(pair) + connections.append({ + 'datasetId': datasetId, + 'parentId': parent_id, + 'childId': child_id, + 'tags': tags, + }) + + # Motion links within each track. + for track_id, nodes in node_ids.items(): + for i in range(len(nodes) - 1): + add_connection(nodes[i][1], nodes[i + 1][1]) + + # Division links: parent track's last node -> child track's first node. + # Ultrack track ids are positive, and founder tracks encode "no parent" as + # -1 or 0 depending on version, so anything <= 0 means no parent. + for track_id, parent_track_id in parent_of.items(): + if parent_track_id is None or parent_track_id <= 0: + continue + parent_nodes = node_ids.get(parent_track_id) + child_nodes = node_ids.get(track_id) + if not parent_nodes or not child_nodes: + continue + add_connection(parent_nodes[-1][1], child_nodes[0][1]) + + return connections + + +def run_ultrack(masks, max_distance, allow_division, working_dir): + """Run Ultrack on a (T, H, W) label stack and return the tracks DataFrame. + + Isolated so tests can patch it without importing ultrack. + """ + # Lazy import: keeps ultrack off the interface/startup path. + from ultrack import MainConfig, Tracker + + config = MainConfig() + # Ultrack persists intermediate results to a SQLite database in this dir. + config.data_config.working_dir = working_dir + config.linking_config.max_distance = float(max_distance) + if not allow_division: + # A strongly negative division weight makes the solver avoid splits. + # This is a strong penalty in the ILP objective, not a hard constraint. + config.tracking_config.division_weight = -1e6 + + tracker = Tracker(config) + tracker.track(labels=masks, overwrite=True) + tracks_df, _graph = tracker.to_tracks_layer(include_parents=True) + return tracks_df + + +def compute(datasetId, apiUrl, token, params): + """ + params (could change): + assignment, channel, connectTo, tags, tile, workerInterface + """ + start_time = timeit.default_timer() + + workerInterface = params['workerInterface'] + track_tags = list(set(workerInterface.get('Tag of objects to track', []) or [])) + max_distance = float(workerInterface.get('Max distance', 50)) + allow_division = bool(workerInterface.get('Allow divisions', True)) + + if not track_tags: + sendError("No tag specified", + info="Please select at least one tag of objects to track.") + raise ValueError("No tag specified") + + tile = params['tile'] + # Always stamp a descriptive tag so tracking connections are identifiable + # (and filterable / bulk-selectable) in the UI even when the user set no + # output tags -- matching the convention of the other connection workers. + output_tags = list(dict.fromkeys((params.get('tags') or []) + ["Ultrack"])) + + batch_xy = batch_argument_parser.process_range_list( + workerInterface.get('Batch XY', None), convert_one_to_zero_index=True) + batch_z = batch_argument_parser.process_range_list( + workerInterface.get('Batch Z', None), convert_one_to_zero_index=True) + if batch_xy is None: + batch_xy = [tile['XY']] + if batch_z is None: + batch_z = [tile['Z']] + batch_xy = list(batch_xy) + batch_z = list(batch_z) + + annotationClient = annotations.UPennContrastAnnotationClient( + apiUrl=apiUrl, token=token) + tileClient = tiles.UPennContrastDataset( + apiUrl=apiUrl, token=token, datasetId=datasetId) + + sendProgress(0.05, "Loading objects", "Fetching annotations") + + blobAnnotationList = annotationClient.getAnnotationsByDatasetId( + datasetId, limit=1000000, shape='polygon') + objectList = annotation_tools.get_annotations_with_tags( + blobAnnotationList, track_tags, exclusive=False) + + if not objectList: + sendWarning("No annotations found", + info="No objects with the specified tag were found to track.") + return + + objectList = [ann for ann in objectList + if ann['location']['XY'] in batch_xy + and ann['location']['Z'] in batch_z] + groups = group_annotations_by_location(objectList) + + if not groups: + sendWarning("No annotations found", + info="No objects with the specified tag were found in the " + "selected XY/Z positions.") + return + + height = tileClient.tiles['sizeY'] + width = tileClient.tiles['sizeX'] + + all_connections = [] + total_groups = len(groups) + for group_idx, ((xy, z), group_annotations) in enumerate(groups.items()): + time_points = sorted(set(ann['location']['Time'] + for ann in group_annotations)) + if len(time_points) < 2: + continue + + masks, label_to_id, label_centroid = build_label_stack( + group_annotations, time_points, height, width) + + sendProgress( + (group_idx + 0.5) / total_groups, + "Tracking objects", + f"Tracking XY {xy + 1}, Z {z + 1} ({len(time_points)} time points)") + + # Each stack gets an isolated working directory for Ultrack's database. + with tempfile.TemporaryDirectory() as working_dir: + tracks_df = run_ultrack(masks, max_distance, allow_division, working_dir) + + connections = tracks_df_to_connections( + tracks_df, label_centroid, label_to_id, datasetId, output_tags) + all_connections.extend(connections) + + sendProgress( + (group_idx + 1) / total_groups, + "Tracking objects", + f"Processed {group_idx + 1} of {total_groups} stacks") + + if not all_connections: + sendWarning("No connections created", + info="Ultrack did not find any links between objects.") + return + + sendProgress(0.95, "Uploading", f"Sending {len(all_connections)} connections") + annotationClient.createMultipleConnections(all_connections) + + elapsed = timeit.default_timer() - start_time + print(f"Created {len(all_connections)} connections in {elapsed:.1f} s") + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + description='Track objects across time with Ultrack') + + parser.add_argument('--datasetId', type=str, required=False, action='store') + parser.add_argument('--apiUrl', type=str, required=True, action='store') + parser.add_argument('--token', type=str, required=True, action='store') + parser.add_argument('--request', type=str, required=True, action='store') + parser.add_argument('--parameters', type=str, required=True, action='store') + + args = parser.parse_args(sys.argv[1:]) + + params = json.loads(args.parameters) + datasetId = args.datasetId + apiUrl = args.apiUrl + token = args.token + + match args.request: + case 'compute': + compute(datasetId, apiUrl, token, params) + case 'interface': + interface(params['image'], apiUrl, token) diff --git a/workers/annotations/ultrack/environment.yml b/workers/annotations/ultrack/environment.yml new file mode 100644 index 0000000..8167f26 --- /dev/null +++ b/workers/annotations/ultrack/environment.yml @@ -0,0 +1,14 @@ +name: worker +channels: + - conda-forge + - defaults +dependencies: + - python=3.10 + - pip + - numpy + - scikit-image + - shapely + - rasterio + - pip: + - rtree + - ultrack diff --git a/workers/annotations/ultrack/tests/Dockerfile_Test b/workers/annotations/ultrack/tests/Dockerfile_Test new file mode 100644 index 0000000..d14e666 --- /dev/null +++ b/workers/annotations/ultrack/tests/Dockerfile_Test @@ -0,0 +1,14 @@ +# Use the existing ultrack worker as the base +FROM annotations/ultrack:latest AS test + +# Install test dependencies +SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"] +RUN pip install pytest pytest-mock + +# Copy test files +RUN mkdir -p /tests +COPY ./workers/annotations/ultrack/tests/*.py /tests +WORKDIR /tests + +# Command to run tests +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "worker", "python3", "-m", "pytest", "-v"] diff --git a/workers/annotations/ultrack/tests/__init__.py b/workers/annotations/ultrack/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/annotations/ultrack/tests/test_ultrack.py b/workers/annotations/ultrack/tests/test_ultrack.py new file mode 100644 index 0000000..04db85e --- /dev/null +++ b/workers/annotations/ultrack/tests/test_ultrack.py @@ -0,0 +1,273 @@ +import numpy as np +import pandas as pd +import pytest +from unittest.mock import patch + +from entrypoint import ( + interface, + compute, + group_annotations_by_location, + build_label_stack, + tracks_df_to_connections, +) + + +@pytest.fixture +def mock_worker_preview_client(): + with patch('annotation_client.workers.UPennContrastWorkerPreviewClient') as mock_client: + yield mock_client.return_value + + +@pytest.fixture +def square_annotations(): + """One 4x4 square per time point, drifting diagonally over three frames.""" + def square(_id, t, x0, y0): + return { + '_id': _id, + 'shape': 'polygon', + 'coordinates': [ + {'x': x0, 'y': y0}, + {'x': x0 + 4, 'y': y0}, + {'x': x0 + 4, 'y': y0 + 4}, + {'x': x0, 'y': y0 + 4}, + {'x': x0, 'y': y0}, + ], + 'location': {'Time': t, 'XY': 0, 'Z': 0}, + 'tags': ['cell'], + } + return [ + square('t0', 0, 2, 2), + square('t1', 1, 4, 4), + square('t2', 2, 6, 6), + ] + + +def test_interface(mock_worker_preview_client): + interface('test_image', 'http://api', 'token') + mock_worker_preview_client.setWorkerImageInterface.assert_called_once() + interface_data = mock_worker_preview_client.setWorkerImageInterface.call_args[0][1] + assert 'Tag of objects to track' in interface_data + assert 'Max distance' in interface_data + assert 'Allow divisions' in interface_data + assert interface_data['Tag of objects to track']['type'] == 'tags' + assert interface_data['Allow divisions']['type'] == 'checkbox' + + +def test_group_annotations_by_location(): + anns = [ + {'shape': 'polygon', 'coordinates': [1, 2, 3], 'location': {'XY': 0, 'Z': 0, 'Time': 0}}, + {'shape': 'polygon', 'coordinates': [1, 2, 3], 'location': {'XY': 1, 'Z': 2, 'Time': 0}}, + ] + groups = group_annotations_by_location(anns) + assert set(groups.keys()) == {(0, 0), (1, 2)} + + +def test_build_label_stack(square_annotations): + masks, label_to_id, label_centroid = build_label_stack( + square_annotations, [0, 1, 2], height=20, width=20) + assert masks.shape == (3, 20, 20) + assert label_to_id[(0, 1)] == 't0' + assert label_to_id[(2, 1)] == 't2' + row, col = label_centroid[(1, 1)] + # Second square spans x 4..8, y 4..8 -> centroid near (5.5, 5.5). + assert abs(row - 5.5) < 1.5 + assert abs(col - 5.5) < 1.5 + + +def _centroid_maps(): + """Simple centroid maps: one object per frame at increasing positions.""" + label_centroid = { + (0, 1): (4.0, 4.0), + (1, 1): (6.0, 6.0), + (2, 1): (8.0, 8.0), + } + label_to_id = {(0, 1): 't0', (1, 1): 't1', (2, 1): 't2'} + return label_centroid, label_to_id + + +def test_tracks_df_to_connections_linear(): + label_centroid, label_to_id = _centroid_maps() + tracks_df = pd.DataFrame([ + {'track_id': 1, 't': 0, 'y': 4.0, 'x': 4.0, 'parent_track_id': -1}, + {'track_id': 1, 't': 1, 'y': 6.0, 'x': 6.0, 'parent_track_id': -1}, + {'track_id': 1, 't': 2, 'y': 8.0, 'x': 8.0, 'parent_track_id': -1}, + ]) + conns = tracks_df_to_connections(tracks_df, label_centroid, label_to_id, 'ds', ['Ultrack']) + pairs = {(c['parentId'], c['childId']) for c in conns} + assert pairs == {('t0', 't1'), ('t1', 't2')} + + +def test_tracks_df_to_connections_division(): + """Track 1 (t0,t1) divides into tracks 2 and 3 at t2.""" + label_centroid = { + (0, 1): (4.0, 4.0), + (1, 1): (6.0, 6.0), + (2, 1): (8.0, 8.0), # daughter A + (2, 2): (2.0, 2.0), # daughter B + } + label_to_id = {(0, 1): 'p0', (1, 1): 'p1', (2, 1): 'dA', (2, 2): 'dB'} + tracks_df = pd.DataFrame([ + {'track_id': 1, 't': 0, 'y': 4.0, 'x': 4.0, 'parent_track_id': -1}, + {'track_id': 1, 't': 1, 'y': 6.0, 'x': 6.0, 'parent_track_id': -1}, + {'track_id': 2, 't': 2, 'y': 8.0, 'x': 8.0, 'parent_track_id': 1}, + {'track_id': 3, 't': 2, 'y': 2.0, 'x': 2.0, 'parent_track_id': 1}, + ]) + conns = tracks_df_to_connections(tracks_df, label_centroid, label_to_id, 'ds', []) + pairs = {(c['parentId'], c['childId']) for c in conns} + # Motion link p0->p1, plus division links p1->dA and p1->dB. + assert pairs == {('p0', 'p1'), ('p1', 'dA'), ('p1', 'dB')} + + +def test_tracks_df_zero_parent_id_is_not_a_real_parent(): + """A founder track with parent_track_id == 0 must not be linked to track 0.""" + label_centroid = { + (0, 1): (4.0, 4.0), # track 0 node at t0 + (1, 1): (6.0, 6.0), # track 0 node at t1 + (0, 2): (50.0, 50.0), # founder track 5 node at t0 + } + label_to_id = {(0, 1): 'z0', (1, 1): 'z1', (0, 2): 'founder'} + # Track 0 is a real track; track 5 is a founder whose parent_track_id is 0. + tracks_df = pd.DataFrame([ + {'track_id': 0, 't': 0, 'y': 4.0, 'x': 4.0, 'parent_track_id': -1}, + {'track_id': 0, 't': 1, 'y': 6.0, 'x': 6.0, 'parent_track_id': -1}, + {'track_id': 5, 't': 0, 'y': 50.0, 'x': 50.0, 'parent_track_id': 0}, + ]) + conns = tracks_df_to_connections(tracks_df, label_centroid, label_to_id, 'ds', []) + pairs = {(c['parentId'], c['childId']) for c in conns} + # Only the motion link within track 0; the founder must NOT be wired to track 0. + assert pairs == {('z0', 'z1')} + assert ('z1', 'founder') not in pairs + + +def test_tracks_df_nearest_match_assigns_correct_annotation(): + """Each row is mapped to the annotation whose mask centroid is nearest.""" + # Two objects per frame: 'A' near the origin, 'B' near (100,100). + label_centroid = { + (0, 1): (0.0, 0.0), (0, 2): (100.0, 100.0), + (1, 1): (0.0, 0.0), (1, 2): (100.0, 100.0), + } + label_to_id = { + (0, 1): 'A0', (0, 2): 'B0', + (1, 1): 'A1', (1, 2): 'B1', + } + # Track 1 stays near the origin; track 2 stays near (100,100). + tracks_df = pd.DataFrame([ + {'track_id': 1, 't': 0, 'y': 1.0, 'x': 1.0, 'parent_track_id': -1}, + {'track_id': 1, 't': 1, 'y': 2.0, 'x': 2.0, 'parent_track_id': -1}, + {'track_id': 2, 't': 0, 'y': 98.0, 'x': 98.0, 'parent_track_id': -1}, + {'track_id': 2, 't': 1, 'y': 99.0, 'x': 99.0, 'parent_track_id': -1}, + ]) + conns = tracks_df_to_connections(tracks_df, label_centroid, label_to_id, 'ds', []) + pairs = {(c['parentId'], c['childId']) for c in conns} + assert pairs == {('A0', 'A1'), ('B0', 'B1')} + + +def test_compute_no_tag_raises(): + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': [], + 'workerInterface': {'Tag of objects to track': []}, + } + with pytest.raises(ValueError, match="No tag specified"): + compute('ds', 'http://api', 'token', params) + + +@patch('entrypoint.run_ultrack') +@patch('annotation_client.tiles.UPennContrastDataset') +@patch('annotation_client.annotations.UPennContrastAnnotationClient') +def test_compute_end_to_end(mock_ann_client_cls, mock_tile_cls, mock_run, + square_annotations): + ann_client = mock_ann_client_cls.return_value + ann_client.getAnnotationsByDatasetId.return_value = square_annotations + ann_client.createMultipleConnections.return_value = {} + + tile_client = mock_tile_cls.return_value + tile_client.tiles = {'sizeX': 20, 'sizeY': 20} + + # Ultrack returns one track linking the three squares. Centroids match the + # rasterized square centroids (row=y, col=x): (3.5,3.5),(5.5,5.5),(7.5,7.5). + def fake_run(masks, max_distance, allow_division, working_dir): + return pd.DataFrame([ + {'track_id': 1, 't': 0, 'y': 3.5, 'x': 3.5, 'parent_track_id': -1}, + {'track_id': 1, 't': 1, 'y': 5.5, 'x': 5.5, 'parent_track_id': -1}, + {'track_id': 1, 't': 2, 'y': 7.5, 'x': 7.5, 'parent_track_id': -1}, + ]) + mock_run.side_effect = fake_run + + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': ['Ultrack'], + 'workerInterface': { + 'Tag of objects to track': ['cell'], + 'Max distance': 50, + 'Allow divisions': True, + 'Batch XY': '', + 'Batch Z': '', + }, + } + + compute('ds', 'http://api', 'token', params) + + ann_client.createMultipleConnections.assert_called_once() + conns = ann_client.createMultipleConnections.call_args[0][0] + pairs = {(c['parentId'], c['childId']) for c in conns} + assert pairs == {('t0', 't1'), ('t1', 't2')} + mock_run.assert_called_once() + + +@patch('entrypoint.run_ultrack') +@patch('annotation_client.tiles.UPennContrastDataset') +@patch('annotation_client.annotations.UPennContrastAnnotationClient') +def test_compute_stamps_descriptive_tag_when_no_output_tags( + mock_ann_client_cls, mock_tile_cls, mock_run, square_annotations): + """Connections carry an 'Ultrack' tag even when no output tags are set.""" + ann_client = mock_ann_client_cls.return_value + ann_client.getAnnotationsByDatasetId.return_value = square_annotations + tile_client = mock_tile_cls.return_value + tile_client.tiles = {'sizeX': 20, 'sizeY': 20} + + def fake_run(masks, max_distance, allow_division, working_dir): + return pd.DataFrame([ + {'track_id': 1, 't': 0, 'y': 3.5, 'x': 3.5, 'parent_track_id': -1}, + {'track_id': 1, 't': 1, 'y': 5.5, 'x': 5.5, 'parent_track_id': -1}, + {'track_id': 1, 't': 2, 'y': 7.5, 'x': 7.5, 'parent_track_id': -1}, + ]) + mock_run.side_effect = fake_run + + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': [], # no output tags configured + 'workerInterface': { + 'Tag of objects to track': ['cell'], + 'Max distance': 50, + 'Allow divisions': True, + 'Batch XY': '', + 'Batch Z': '', + }, + } + compute('ds', 'http://api', 'token', params) + + conns = ann_client.createMultipleConnections.call_args[0][0] + assert conns, "expected at least one connection" + assert all('Ultrack' in c['tags'] for c in conns) + + +@patch('entrypoint.run_ultrack') +@patch('annotation_client.tiles.UPennContrastDataset') +@patch('annotation_client.annotations.UPennContrastAnnotationClient') +def test_compute_no_annotations(mock_ann_client_cls, mock_tile_cls, mock_run): + ann_client = mock_ann_client_cls.return_value + ann_client.getAnnotationsByDatasetId.return_value = [] + + params = { + 'tile': {'XY': 0, 'Z': 0, 'Time': 0}, + 'tags': ['Ultrack'], + 'workerInterface': { + 'Tag of objects to track': ['cell'], + 'Batch XY': '', + 'Batch Z': '', + }, + } + compute('ds', 'http://api', 'token', params) + ann_client.createMultipleConnections.assert_not_called() + mock_run.assert_not_called()