Repository navigation
Add mask-tracked-crabs entry point - #293
Merged
Merged
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #293 +/- ##
==========================================
+ Coverage 50.89% 55.66% +4.77%
==========================================
Files 26 28 +2
Lines 1682 1870 +188
==========================================
+ Hits 856 1041 +185
- Misses 826 829 +3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
sfmig
force-pushed
the
smg/mask-tracked-crabs
branch
from
October 6, 2026 14:07
8fb338e to
b4458db
Compare
This was referenced Oct 8, 2026
`read_tracked_bboxes_from_zarr` turns one clip of one video group into the `tracked_bboxes_dict` the masking pass consumes: a map from clip frame index to boxes and individual labels, with no key for a frame with no boxes. It takes the already-open video dataset rather than a path, so a run over many clips opens the datatree once. The guard on the time axis counts finite `escape_state` rows rather than comparing `sizes["time"]`, because after the outer join along `clip_id` every clip of a video reports the video's longest axis. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`crabs/tracker/mask_video.py` holds the two functions the pass splits into, because SAM2 runs over one clip's video file at a time but a trajectories store holds a whole set of a video's clips: - `create_mask_store` writes the xarray template for one video group and returns its raw `labels` array. The store holds one label image per frame -- a uint16 array in which each crab's pixels carry that crab's own value -- mirroring the trajectories store's `clip_id`, `time` and `individual` coordinates, so the two align 1:1. The `individual` to pixel value mapping is explicit, in `label_of`. - `write_clip_masks_to_store` runs one clip, buffering a shard's worth of frames and writing them whole, so zarr never read-modify-writes a shard file. Neither constructs the predictor: it is the caller's, so one predictor is reused across every clip of every video in a run. Chunks are one frame and shards default to 32 frames, which decouples the access pattern from the file count. Prompts are chunked to bound peak memory, and sorted by box area first, so that painting chunk by chunk cannot invert the `smallest_wins` occlusion policy. The unit tests drive the pass through a fake predictor, so they need no SAM2 and no GPU. The `max_prompts_per_batch` parametrisation is what pins the output as independent of that knob. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
One command, and one `--boxes` argument whose suffix selects the reader. This PR accepts `.zarr` and rejects anything else by suffix: no stub and no dead branch, so `--help` never advertises something that does not work. The run opens the trajectories datatree once, resolves every clip video before SAM2 is loaded -- so a mis-pointed `--videos` costs a second rather than a model download -- and then loads one predictor for the whole run, reused across every clip of every video. `--zarr_mode_store` and `--zarr_mode_group` mirror `create-zarr-dataset`'s own, so a cluster array job with one job per video can append groups to a shared store. `zarr>=3` and `xarray` were already imported by `crabs/zarr` but only arrived transitively via movement; this adds a second direct importer of each, so declare them. SAM2 itself is an opt-in dependency group: it is not on PyPI, and groups are never written into distribution metadata, so the sdist stays publishable. The test runs the entry point end to end with SAM2 replaced by the fake predictor, and asserts that the individuals present in `labels` are exactly those the trajectories store has for that frame. That comparison is the contract of this feature. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The install, the run, the config knobs, the store layout, and the five things a reader cannot work out from the store itself: that occlusion is resolved at write time and the lost pixels are not recorded; that `individual` names mean something only within a clip, and only within a video, in value and in format; that `label_of` is the mapping in both directions; and that the trajectories store is the index for the mask store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The default store name carries a timestamp, computed per process. A SLURM array with one task per video would therefore have each task write its own store, so the `--zarr_mode_store a` append pattern `create-zarr-dataset` uses could not work here. `--output_store` names the store exactly and ignores `--output_dir`, mirroring `create-zarr-dataset`'s `--zarr_store`. An array job points every task at one store named after the array job ID, which keeps the property the timestamp was there for: two runs cannot collide. The test is the array pattern itself -- two runs, each masking one video group with `--match`, both appending to one store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`bash_scripts/run_mask_array.sh` is a SLURM array with one task per video, so all videos are masked in parallel. It follows `run_zarr_dataset.sh`: a uv environment per job under $TMPDIR, the list of tasks derived from the input, and every task appending its own group to one shared store. Two things it does that no other script here does. It installs SAM2 after the crabs package, with `--no-build-isolation` so it builds against the torch already in the environment rather than pulling a second one into an isolated build env. And it sets `HF_HOME` to /ceph/scratch, so 27 jobs share one download of the checkpoint instead of each fetching its own. The task list is the video groups of the input store, read straight off disk with `find | sort`, rather than re-derived from the metadata csv: the store is what the job actually consumes, and a stable sort keeps a task index meaning the same video across re-runs. `guides/MaskTrackedCrabsHPC.md` follows CreateZarrDatasetForTracks.md, and covers the two inputs and where they come from, the mask config knobs and what each one trades, the array sizing, the output store and how to read it, and re-running failed tasks into the same store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Turning off build isolation means SAM2's build requirements have to be satisfied in our own environment, and a `uv venv` starts with neither setuptools nor pip. Verified: the install fails with "ModuleNotFoundError: No module named 'setuptools'" from inside the build backend. It happens to work today only because torch declares setuptools as a runtime dependency, so installing the crabs package drags it in. That is not a thing to rely on, and the failure mode if a torch release drops it is opaque and 27 jobs deep. Install it explicitly. Also set SAM2_BUILD_CUDA=0. The sam2._C extension is only used to fill holes and remove sprinkles in predicted masks, which the image predictor never does: `SAM2ImagePredictor` defaults both `max_hole_area` and `max_sprinkle_area` to 0, and `postprocess_masks` only reaches `get_connected_components` when one is above 0. Left on, the build either spends minutes compiling an extension we do not call, or prints a traceback and carries on, since SAM2 allows build errors by default. Verified end to end in a clean venv: torch plus setuptools, then the no-isolation build, which installs sam-2 without touching the existing torch, builds no compiled extension, and imports. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ock to avoid jobs timing out waiting for the uv cache lock
SAM2ImagePredictor.predict copies every full-frame mask to the host as float32. At 4K with ~40 boxes per frame that copy and the astype(bool) after it took ~2.3 of ~2.7 s per frame on an L40S (measured with scripts/profile_mask_timing.py). Call the two private methods predict is made of instead, paint the label image on SAM2's device, and copy only the finished label frame back: ~0.23 s per frame. A randomly initialised SAM2 checks the private-API path gives the same masks as the public predict, so a SAM2 bump that changes it fails in the tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Since masks are flattened on the GPU, the batch size no longer sets host memory. Also note that one chunk per frame applies the occlusion policy exactly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Even at ~0.25 s per frame, the longest videos (~330k frames) take about a day, which a 24 h limit would cut off. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Blosc clevel=9 maps to zstd's maximum level. On 4K label images that took ~0.6 s per frame, more than SAM2's whole GPU pass, for files only ~25% smaller than at the default clevel=5 (~14 ms per frame). The label_of array is a few KB, so its level 9 bought nothing either. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On 4K label images zarr's default ZstdCodec wrote 24 kB per frame in ~13 ms, as small as Blosc zstd at clevel=9 (~600 ms) and smaller than Blosc zstd with bitshuffle at clevel=5 (31 kB). Label images are runs of equal values, which plain zstd handles without shuffling. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Rewrite the profiling script to run the current SAM2 pass, then write its label frames shard by shard with the old compressor (Blosc zstd, clevel=9) and with zarr's default, on the filesystem the jobs write to. It prints the time per frame of each, the speed-up, and the hours for the smallest, median and largest videos. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
sfmig
force-pushed
the
smg/mask-tracked-crabs
branch
from
October 13, 2026 16:38
d1a7004 to
8899bb7
Compare
sfmig
marked this pull request as ready for review
October 13, 2026 17:41
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
What is this PR?
Why is this PR needed?
The orientation work needs per-crab masks, not just boxes. We already have tracked boxes for the whole dataset, in the trajectories zarr store
create-zarr-datasetwrites. We can prompt SAM2 with them directly to generate the required masks.What does this PR does?
Adds a
mask-tracked-crabsentry point that prompts SAM2 with the boxes in a trajectories store, and writes a zarr store holding one label image per frame. A label image is auint16array in which each crab's pixels carry that crab's own label, with label0being background.The
clip_id,timeandindividualcoordinates match the trajectories store's ones, so masks and trajectories align 1:1 with no reindexing. Additionally, a label image array for one frame feedsregionpropsandnapari.add_labelswith no conversion step on read.Specifically, this PR:
mask-tracked-crabsentry point--output_storetest_sam2CI job.References
This is PR 1 of 3, per
claude-plans/proposal-mask-tracked-crabs.md.--boxesto a<clip>_tracks.csv.detect-and-track-maskentry point built on the same pass.Other references:
claude-plans/proposal-mask-tracked-crabs.md— the full proposal, with the measurements behind each decision.timeaxis dense and clip-anchored.scripts/generate_masks_from_bboxes.py— the standalone SAM2 script this supersedes for tracked boxes.How has this PR been tested?
Existing and added tests pass locally and in CI.
The
test_sam2CI job that installs themasksextra (CPU-only torch) and checks thatpredict_masks_on_device(our wrapper of SAM2's private API, which keeps the masks on SAM2's device) matches SAM2's publicpredictWe defer the integration test for now, which would required a small trajectories store plus the clip
.mp4s it names added to the GIN registry. The unit tests cover the same contract locally, including the end-to-end entry point run, but against a fake predictor rather than real SAM2.Does this PR require an update to the documentation?
Yes, it is included as part of the work in:
guides/MaskTrackedCrabsHPC.md, andcrabs/tracker/README.md.Checklist:
🤖 Generated with Claude Code and edited by @sfmig