StarDist

StarDist 2-D on a z-stack: each slice is segmented independently. StarDist needs more overlap than Cellpose — 32 voxels for the default 2D_versatile_fluo model.

Code

import numpy as np
from stardist.models import StarDist2D
from patchworks import tile_process

IMAGE = "image.zarr"
OUTPUT = "labels_sd.zarr"

# Load model once (not inside fn — would re-download every call)
model = StarDist2D.from_pretrained("2D_versatile_fluo")


def stardist_fn(tile: np.ndarray) -> np.ndarray:
    squeeze = tile.ndim == 3 and tile.shape[0] == 1
    img = tile[0] if squeeze else tile

    # StarDist expects normalised float32
    norm = img.astype("float32")
    if norm.max() > 0:
        norm = norm / norm.max()

    labels, _ = model.predict_instances(norm)
    labels = labels.astype("int32")
    return labels[np.newaxis] if squeeze else labels


tile_process(
    IMAGE,
    stardist_fn,
    channel=0,
    tile_shape=(1, 1024, 1024),
    overlap=32,  # StarDist receptive field is larger than Cellpose
    write_to=OUTPUT,
    progress=True,
)

Load the model once, outside the function: loaded inside, it is re-initialised for every tile.