Support all Cellpose v4 models (SAM + DINO) in GPU-resident segmentation by alxndrkalinin · Pull Request #51 · alxndrkalinin/cubic

alxndrkalinin · 2026-06-15T21:00:25Z

Summary

Extends cubic's device-resident Cellpose segmentation to support all four Cellpose v4 models, not just the original SAM weights:

Model	Backbone	Default tile
`cpsam`	sam_vitl	256
`cpsam_v2`	sam_vitl	256
`cpdino`	dino_vitl	384
`cpdino-vitb`	dino_vitb	384

Both CPSAM and CPDINO forwards return the identical (flow[B,3,bsize,bsize], style[B,256]) contract; the only backbone-specific behavior is the tile size, which cubic already mirrored. So the resident path needed only to expose bsize and centralize the backbone→tile-size rule.

Changes

segment_cellpose is now the canonical GPU-resident entry point (covers SAM + DINO). segment_cpsam is kept as a deprecated alias (emits DeprecationWarning, slated for removal after 0.8). Both are exported.
Added a bsize parameter, threaded through the network forward, with cellpose's sam_vitl → 256 guard. DINO tile size is tunable.
Centralized the backbone→tile-size rule into a single _resolve_bsize(backbone, bsize) helper (default + validation), removing the now-unused backbone arg from _run_net_gpu.
Generalized the non-resident cellpose_eval/cellpose_segment docstrings (they already accept any model name).
Packaging: base cellpose pin bumped to >=4.2.0 (the release that ships cpsam_v2/cpdino) and a new cellpose-dino extra (cubic[cellpose] + cellpose[dino]>=4.2.0) for the dinov3 dependency. README install note added.
Bumped version to 0.8.0a1.

Tests

Parametrized 2D parity over cpsam_v2, cpdino, cpdino-vitb; a DINO 3D parity test; a DINO bsize-override parity test (all GPU-gated, vs stock CellposeModel.eval, AP@0.5 ≥ 0.95).
A fast pure-function unit test for _resolve_bsize (no GPU).
A deprecation-alias test verifying segment_cpsam warns and forwards verbatim.
A module-scoped fixture that caches one CellposeModel per backbone, so the GPU suite reloads each weight set only once.

Verification

ruff check, ruff format --check, mypy clean.
Full cellpose-resident GPU suite: 21 passed (AP ≥ 0.95 parity for all four models, 2D + 3D, on real GPU weights).
Full project suite: 364 passed (pre-existing unrelated warnings only).

🤖 Generated with Claude Code

Summary by Sourcery

Extend GPU-resident Cellpose segmentation to support all Cellpose v4 models and rename the main entry point, while preserving compatibility via a deprecated alias.

New Features:

Add a unified segment_cellpose GPU-resident entry point that supports SAM and DINO Cellpose v4 backbones, with configurable tile size.
Introduce a backbone-aware tile-size resolver to mirror Cellpose v4 defaults and validation across SAM and DINO models.

Enhancements:

Generalize existing Cellpose helpers and docstrings from SAM-only to any Cellpose v4 model.
Cache CellposeModel instances per backbone in tests to avoid reloading weights for each test run.
Update project version to 0.8.0a1.

Build:

Bump the Cellpose dependency to >=4.2.0 and introduce a cellpose-dino extra, wiring it into the all extra.

Documentation:

Document the new Cellpose DINO extra and clarify installation instructions and supported models in the README.

Tests:

Expand GPU parity tests to cover new Cellpose v4 backbones, 3D DINO usage, custom tile sizes, and the deprecated alias behavior while adding a unit test for tile-size resolution.

Rename the GPU-resident entry point to segment_cellpose and cover every Cellpose v4 model -- the SAM backbone (cpsam, cpsam_v2) and the DINOv3 backbones (cpdino, cpdino-vitb). The backbones differ only in the default tile size (256 for sam_vitl, 384 for dino_*) and share the (flow, style) forward contract, so the resident path needs only to expose bsize and thread it through _run_net_gpu, with cellpose's sam_vitl->256 guard. segment_cpsam is kept as a deprecated alias (warns, slated for removal after 0.8). Parity tests confirm resident masks match stock CellposeModel.eval (AP >= 0.95) for all four models in 2D and 3D. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The DINO models (cpdino, cpdino-vitb) and cpsam_v2 first ship in cellpose 4.2.0, and the DINO backbones additionally need the dinov3 dependency from cellpose's `dino` extra. Bump the base cellpose extra to >=4.2.0 and add a cellpose-dino extra (cubic[cellpose] + cellpose[dino]>=4.2.0), documenting both install paths in the README. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The default tile size (256 for SAM, 384 for DINO) and the SAM-pinned-to-256 guard were encoded separately in _run_net_gpu and segment_cellpose -- two sources of truth for the same backbone->tile-size rule. Extract _resolve_bsize to own both; segment_cellpose resolves bsize once and passes the concrete value down, so _run_net_gpu no longer needs the backbone argument. Replace the GPU-only bsize-guard test with a fast pure-function unit test that covers the full resolution matrix (defaults per backbone + the SAM guard), removing a heavyweight model load for a string-compare assertion. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The GPU parity tests each constructed their own CellposeModel, reloading the same weights repeatedly (cpsam alone was built ~8 times). Add a module-scoped factory fixture that caches one model per backbone and route every GPU test through it. eval/segment_cellpose are read-only inference, so a shared instance is safe; the fixture is lazy, so no-GPU/no-cellpose runs never construct a model and free VRAM on teardown. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

First alpha of the 0.8 cycle, which adds DINO-model support to the GPU-resident segmentation path and deprecates segment_cpsam. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

sourcery-ai · 2026-06-15T21:00:36Z

Reviewer's Guide

Extends the GPU-resident Cellpose segmentation path to a unified segment_cellpose entry point that supports all Cellpose v4 backbones (SAM and DINO), centralizes tile-size handling via a _resolve_bsize helper, adds a deprecated segment_cpsam alias, and updates tests, packaging, and docs accordingly.

Sequence diagram for unified GPU-resident Cellpose v4 segmentation

sequenceDiagram
    actor User
    participant segment_cpsam
    participant segment_cellpose
    participant _resolve_bsize
    participant _run_net_gpu
    participant CellposeModel_net as CellposeModel.net

    alt [User calls segment_cpsam]
        User->>segment_cpsam: segment_cpsam(model, x, ...)
        segment_cpsam->>segment_cpsam: warnings.warn(...)
        segment_cpsam->>segment_cellpose: segment_cellpose(model, x, ...)
    else [User calls segment_cellpose]
        User->>segment_cellpose: segment_cellpose(model, x, bsize, ...)
    end

    segment_cellpose->>segment_cellpose: _check_gpu_precondition(model)
    segment_cellpose->>_resolve_bsize: _resolve_bsize(model.backbone, bsize)
    _resolve_bsize-->>segment_cellpose: bsize

    segment_cellpose->>_run_net_gpu: _run_net_gpu(model.net, x, bsize=bsize, ...)
    _run_net_gpu->>CellposeModel_net: _forward_gpu(net, ...)
    CellposeModel_net-->>_run_net_gpu: flow, style
    _run_net_gpu-->>segment_cellpose: dP, cellprob, styles

    segment_cellpose-->>User: masks, [None, dP, cellprob], styles

File-Level Changes

Change	Details	Files
Introduce unified GPU-resident Cellpose v4 entry point and deprecate the SAM-only API.	Rename the main GPU-resident segmentation function to segment_cellpose, keeping segment_cpsam as a deprecated alias that forwards with a DeprecationWarning. Update the module docstring and error messages to reference segment_cellpose and generic Cellpose v4 (SAM or DINO) models rather than SAM-only. Export both segment_cellpose and segment_cpsam from cubic.segmentation.init.	`cubic/segmentation/cellpose_sam_gpu.py` `cubic/segmentation/__init__.py`
Centralize backbone-specific tile size (bsize) handling and simplify the GPU net runner.	Add _resolve_bsize(backbone, bsize) to enforce SAM’s fixed 256 tile and DINO’s default 384 with flexible overrides. Change _run_net_gpu to accept a concrete bsize int (no backbone arg) and rely on the caller to resolve it, threading bsize through all call sites. Thread an optional bsize parameter through segment_cellpose, resolving it with _resolve_bsize(model.backbone, bsize) before per-image recursion and _run_net_gpu invocation. Clarify comments around the unused style output in run_net_gpu to cover both SAM and DINO semantics.	`cubic/segmentation/cellpose_sam_gpu.py`
Broaden high-level Cellpose helpers to all v4 models and keep behavior consistent.	Generalize cellpose_eval docstring to describe support for any Cellpose v4 pretrained_model, including SAM and DINO backbones and custom paths. Update cellpose_segment docstring to say it runs a generic Cellpose v4 model (SAM or DINO), not just SAM.	`cubic/segmentation/cellpose.py`
Extend and refactor tests to cover new models, new API, and performance improvements.	Add a module-scoped cellpose_model fixture that caches CellposeModel instances per pretrained_model and clears CUDA cache after tests. Switch existing GPU-resident tests from segment_cpsam to segment_cellpose, while keeping coverage of residency, list inputs, and parity vs stock eval. Add a deprecation test ensuring segment_cpsam forwards to segment_cellpose verbatim and emits DeprecationWarning. Add parity tests for new v4 models (cpsam_v2, cpdino, cpdino-vitb) in 2D, a DINO 3D parity test, a _resolve_bsize unit test, and a DINO bsize-override parity test.	`tests/segmentation/test_cellpose_sam_gpu.py`
Update packaging and documentation to expose DINO support and new version.	Bump the core cellpose extra to require cellpose>=4.2.0. Add a new cellpose-dino extra that layers on cellpose[dino]>=4.2.0 and include it in the all extra. Document separate installation instructions for SAM-only vs SAM+DINO models in README. Bump project version from 0.7.0 to 0.8.0a1.	`pyproject.toml` `README.md` `cubic/__init__.py`

Tips and commands

Interacting with Sourcery

Trigger a new review: Comment @sourcery-ai review on the pull request.
Continue discussions: Reply directly to Sourcery's review comments.
Generate a GitHub issue from a review comment: Ask Sourcery to create an
issue from a review comment by replying to it. You can also reply to a
review comment with @sourcery-ai issue to create an issue from it.
Generate a pull request title: Write @sourcery-ai anywhere in the pull
request title to generate a title at any time. You can also comment
@sourcery-ai title on the pull request to (re-)generate the title at any time.
Generate a pull request summary: Write @sourcery-ai summary anywhere in
the pull request body to generate a PR summary at any time exactly where you
want it. You can also comment @sourcery-ai summary on the pull request to
(re-)generate the summary at any time.
Generate reviewer's guide: Comment @sourcery-ai guide on the pull
request to (re-)generate the reviewer's guide at any time.
Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
pull request to resolve all Sourcery comments. Useful if you've already
addressed all the comments and don't want to see them anymore.
Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
request to dismiss all existing Sourcery reviews. Especially useful if you
want to start fresh with a new review - don't forget to comment
@sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

Enable or disable review features such as the Sourcery-generated pull request
summary, the reviewer's guide, and others.
Change the review language.
Add, remove or edit custom review instructions.
Adjust other review settings.

Getting Help

Contact our support team for questions or feedback.
Visit our documentation for detailed guides and information.
Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

sourcery-ai

Hey - I've left some high level feedback:

In _resolve_bsize, any non-sam_vitl backbone implicitly falls through to the DINO default (384 or the provided bsize); consider making the set of supported backbones explicit (and raising for unknown ones) so that future Cellpose backbones don’t silently inherit DINO behavior.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In `_resolve_bsize`, any non-`sam_vitl` backbone implicitly falls through to the DINO default (384 or the provided `bsize`); consider making the set of supported backbones explicit (and raising for unknown ones) so that future Cellpose backbones don’t silently inherit DINO behavior.

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

Copilot

Pull request overview

This PR generalizes cubic’s GPU-resident Cellpose (v4) segmentation path to support both SAM- and DINO-backed models, making segment_cellpose the canonical entry point while retaining segment_cpsam as a deprecated alias for backward compatibility.

Changes:

Introduces segment_cellpose (GPU-resident) with backbone-aware tile-size resolution via _resolve_bsize, and keeps segment_cpsam as a deprecated forwarding alias.
Expands the GPU parity test suite to cover additional v4 models (SAM v2 + DINO variants), including 3D and bsize override coverage, and adds a unit test for _resolve_bsize.
Updates packaging/docs: bumps Cellpose minimum version, adds a cellpose-dino extra, documents install options, and bumps package version to 0.8.0a1.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/segmentation/test_cellpose_sam_gpu.py	Adds cached model fixture, deprecation-alias test, and parity coverage for new v4 backbones (incl. 3D + `bsize` override).
README.md	Documents new `cellpose-dino` extra and clarifies Cellpose extras install commands.
pyproject.toml	Bumps Cellpose minimum version and adds `cellpose-dino` optional dependency + wiring into `all`.
cubic/segmentation/cellpose.py	Generalizes docstrings from SAM-only wording to “Cellpose v4 (SAM or DINO)”.
cubic/segmentation/cellpose_sam_gpu.py	Implements `segment_cellpose`, `_resolve_bsize`, removes backbone arg from `_run_net_gpu`, and adds deprecated `segment_cpsam` alias.
cubic/segmentation/init.py	Exports `segment_cellpose` alongside deprecated `segment_cpsam`.
cubic/init.py	Bumps package version to `0.8.0a1`.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

The ImportError raised when cellpose is missing still said "(SAM)" only, while _check_gpu_precondition was already updated to "(SAM or DINO)". Align the message and point DINO users at the dinov3 install note in the README. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Published cellpose (4.2.1.1) does not expose a `dino` extra, and dinov3 is git-only (pip install git+https://github.com/facebookresearch/dinov3), so it cannot be declared as a dependency of a PyPI package. The cellpose-dino extra therefore pulled no dinov3 (uv warned the extra did not exist). Remove it and document the manual dinov3 install in the README for the DINO backbones. Regenerate uv.lock, which was stale against the cellpose>=4.2.0 bump (it still recorded >=4.0 and resolved cellpose 4.1.1, predating cpsam_v2/cpdino). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

bsize is now part of the public segment_cellpose API, but _resolve_bsize passed any explicit value straight through for DINO backbones, so bsize<=0 reached make_tiles/run_net_gpu and caused divide-by-zero / invalid tiling. Validate it as a positive integer and fail fast with a clear error. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

alxndrkalinin and others added 5 commits June 15, 2026 13:14

chore: bump version to 0.8.0a1

9bb97e2

First alpha of the 0.8 cycle, which adds DINO-model support to the GPU-resident segmentation path and deprecates segment_cpsam. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

sourcery-ai Bot reviewed Jun 15, 2026

View reviewed changes

alxndrkalinin requested a review from Copilot June 15, 2026 21:02

Copilot started reviewing on behalf of alxndrkalinin June 15, 2026 21:02 View session

Copilot AI reviewed Jun 15, 2026

View reviewed changes

Comment thread cubic/segmentation/cellpose_sam_gpu.py Outdated

alxndrkalinin and others added 3 commits June 15, 2026 14:33

alxndrkalinin merged commit 6dbf90e into main Jun 15, 2026
9 checks passed

alxndrkalinin deleted the feat/cellpose-dino-models branch June 15, 2026 21:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support all Cellpose v4 models (SAM + DINO) in GPU-resident segmentation#51

Support all Cellpose v4 models (SAM + DINO) in GPU-resident segmentation#51
alxndrkalinin merged 8 commits into
mainfrom
feat/cellpose-dino-models

alxndrkalinin commented Jun 15, 2026 •

edited by sourcery-ai Bot

Loading

Uh oh!

sourcery-ai Bot commented Jun 15, 2026 •

edited

Loading

Interacting with Sourcery

Customizing Your Experience

Getting Help

Uh oh!

sourcery-ai Bot left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

alxndrkalinin commented Jun 15, 2026 • edited by sourcery-ai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Tests

Verification

Summary by Sourcery

Uh oh!

sourcery-ai Bot commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviewer's Guide

Sequence diagram for unified GPU-resident Cellpose v4 segmentation

File-Level Changes

Interacting with Sourcery

Customizing Your Experience

Getting Help

Uh oh!

sourcery-ai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

alxndrkalinin commented Jun 15, 2026 •

edited by sourcery-ai Bot

Loading

sourcery-ai Bot commented Jun 15, 2026 •

edited

Loading