cellmap_flow.finetune.data.sampler

Where a finetune patch is centred: the dense, sparse and rehearsal pools.

Two-pool stratified sampling. The annotated voxels of a session’s volume are split by its imported_crops boxes (recorded in the volume’s attrs when YAML crops are imported):

• dense pool: foreground voxels inside any imported crop (abundant ground truth);

• sparse pool: annotated voxels outside every crop, background as well as foreground (painted scribbles: by construction sparse and informative, since the user paints where the base model failed).

A patch picks a pool by dense_to_sparse_ratio (0.5/0.5 when both pools have voxels, the non-empty one when only one has), a random voxel of that pool, and a jitter of the centre. Without the split, voxel-uniform sampling buries the scribbles: a typical session has ~40M dense voxels against ~10K painted, so 999 patches in 1000 would be dense and the corrections barely move the gradient. The split gives them a set share of every epoch, whatever the voxel counts.

A third pool, the good regions, is drawn from directly (see PatchSampler).

The index reads only populated chunks, the z.y.x files under annotation/s0/: an empty volume costs nothing, a painted region its annotated voxels.

A one-slice patch (a 2D network’s) can be cut in the plane its paint was drawn in rather than always across z (slice_planes): paint drawn in a z–y slice is one voxel thick in x, and an x–y slice through it holds only a line, a voxel wide, which teaches a 2D network next to nothing (Cellpose learned flows along y alone from it, and segmented strips).

Attributes

logger
PLANE_NAMES
POOL_DTYPE
FOREGROUND_SHARE
ANCHOR_SHARE

Classes

PatchSampler

The patch centres of one annotation volume, and the draw among them.

Functions

paint_depth_axes(voxels) → numpy.ndarray

The axis each painted voxel's paint is thinnest across, as int8: 0 (z,

Module Contents

cellmap_flow.finetune.data.sampler.logger

cellmap_flow.finetune.data.sampler.PLANE_NAMES = ('x–y', 'z–x', 'z–y')

cellmap_flow.finetune.data.sampler.paint_depth_axes(voxels: numpy.ndarray) → numpy.ndarray

The axis each painted voxel’s paint is thinnest across, as int8: 0 (z, an x–y slice) unless y or x is clearly thinner.

voxels (N, 3) are the painted voxels, each once. A voxel’s paint is measured by the run of painted voxels through it along each axis: paint drawn in a z–y slice has runs of 1 along x and long ones along z and y. y or x is taken when its run is at most half z’s, so a blob painted in 3D, or a lone voxel, stays an x–y slice as every patch was before, and a line along z, which lies in both side planes, goes to y.

cellmap_flow.finetune.data.sampler.POOL_DTYPE

cellmap_flow.finetune.data.sampler.FOREGROUND_SHARE = 0.3333333333333333

cellmap_flow.finetune.data.sampler.ANCHOR_SHARE = 0.25

class cellmap_flow.finetune.data.sampler.PatchSampler(volume_zarr_path: str, output_voxel_size_nm, jitter_voxels, dense_to_sparse_ratio: float | None = None, good_regions: list | None = None, rehearsal_fraction: float | None = None, anchor_fraction: float | None = None, slice_planes: str | None = None)

The patch centres of one annotation volume, and the draw among them.

Rehearsal (“good”) regions are boxes the user looked at and certified the model already handles. The supervised loss only touches voxels that were labelled, so nothing stops the adapter drifting everywhere else; a rehearsal patch pins the student to the teacher inside a region the user vouched for, chosen deliberately rather than “wherever happens to sit next to a scribble”, which is where the model is least likely to be right. A good region carries no annotations, so neither of the other pools, which index annotated voxels, can reach it: it is a pool of its own, of box centres.

Built once, where the dataset is made; the pools travel to each loader worker in the dataset’s pickle.

dense, sparse

(N, 3) annotation-voxel indices of the two pools (POOL_DTYPE), in the order the chunk files are listed and scanned. Either may be empty, never both.

sparse_chunk_starts

where each chunk’s voxels begin in sparse, plus its length at the end: chunk i holds rows sparse_chunk_starts[i]:sparse_chunk_starts[i + 1].

sparse_foreground_rows

the rows of sparse that are foreground; the rest are painted background.

effective_dense_ratio

the share of non-rehearsal patches taken from dense; dense_to_sparse_ratio as asked (None: auto), clamped away from an empty pool.

rehearsal_centres

(M, 3) good-region centres in annotation voxels, or None when no region is usable.

effective_rehearsal_fraction

the share of patches centred on one; 0 without a usable region.

effective_anchor_fraction

the share of the other patches centred at a random point of the volume, held to the teacher where unpainted; ANCHOR_SHARE by default for a painted session, else 0.

annotated_chunks

how many chunks hold a voxel of a pool, the default epoch length.

class_voxels

(background, foreground), how many voxels of the volume are annotated as each, the crops’ and the paint’s alike.

corner_nm, shape_voxels

where annotation voxel 0’s lower corner is (see volume_corner_nm), and the volume’s shape in voxels.

paint_planes

how many painted voxels (of sparse) lie in an x–y, a z–x and a z–y slice (paint_depth_axes), or None when slice_planes is None.

sparse_depth

the axis each row of sparse is cut across, when slice_planes is “any” and some paint is in a side plane; else None, and every patch is cut across z.

slice_planes is for a one-slice patch: None, a 3D patch, which takes paint drawn in any plane as it is; “xy”, one slice that is always an x–y one, whose paint is only counted (paint_planes) so that paint in the side planes can be warned about; “any”, one slice in whichever plane the paint around its centre lies in. jitter_voxels is then in the patch’s own order (reader.slice_axes), its first entry across the slice.

volume_zarr_path

output_voxel_size

jitter

dense_to_sparse_ratio

good_regions = []

rehearsal_fraction

draw(rng: numpy.random.Generator) → Tuple[numpy.ndarray, bool, int]

A patch centre in annotation voxels, whether it is a rehearsal patch, and the axis the patch is cut across (0, z, unless slice_planes is “any”).

The centre is not whole voxels yet (see PatchReader.snap). What is drawn from rng, in order: rehearsal or not (only when there is a usable good region); then which region, or else a random anchor or not (only when anchors are on), where, and its plane (only when some paint is in a side plane); or else which pool (only when both can be chosen); in the painted pool, whether to centre on foreground (only when some is painted), then which chunk, or which foreground voxel; which voxel, and the jitter.

log_rehearsal_status() → None

Say what is happening with the good regions, if there are any, and with the random anchors.

Rehearsal switched off – a deliberate choice – and good regions that all fell outside the volume are different states, and are told apart.