spacr.timeflows_model

Timeflows: a time head on Cellpose-SAM, trained on tracked movies.

Cellpose predicts, per pixel, a flow toward the centre of the object the pixel belongs to. Timeflows adds the same idea through TIME: per pixel of frame t, where the centre of that object is in frame t+1, and whether it has a successor at all (cells die, divide, leave the field). Linking objects between frames then reads the model instead of guessing from overlap, which is where the plain IoU stitcher fails: fast cells that move more than their own size between frames (the MuSC movies move 3-4 diameters per frame at p95).

WHAT IS HERE, ALL RUNNABLE ON A CPU:

• track_masks_from_ctc() – full masks carrying track ids, from a Cell Tracking Challenge movie’s silver segmentation relabelled by its tracking markers (the TRA ground truth marks cells with small markers only).

• time_targets() – the training target for one pair: a unit-diameter displacement field toward the successor’s centre, a successor flag, and the pixels both are supervised on.

• augment_pair() – one flip/rotation applied identically to both frames and both label images; the targets are computed AFTER it from the labels, so a flip can never teach a displacement that did not happen.

• pair_sampling_weights() – pairs drawn across the displacement distribution rather than uniformly over mostly still frames.

• TimeflowsNet – a backbone (Cellpose-SAM’s encoder, or any module with the same features contract) shared by both frames, and a head that reads their features together.

• timeflows_loss(), train_timeflows() – masked losses, and the curriculum: the head alone on a frozen backbone, then all of it at a low learning rate.

• link_by_timeflows(), scramble_test() – linking objects from the predictions, and the check that a model recovers identities after the next frame’s labels are shuffled.

Nothing here starts a GPU training run; that is a separate, queued job.

Classes

CellposeSamFeatures

Cellpose-SAM's encoder as a feature extractor, sharing its weights.

Functions

TimeflowsNet(backbone[, channels, ps])	The time head on a shared backbone, as a torch.nn.Module.
augment_pair(→ Tuple[List[numpy.ndarray], ...)	One random flip and quarter-turn, the SAME for every frame and label.
ctc_pairs() → List[_Pair])	Consecutive-frame training pairs from one Cell Tracking Challenge movie.
link_by_timeflows(→ Dict[int, int])	Which object in t+1 each object in t becomes, from the model.
main(→ int)	python -m spacr.timeflows_model --movies DIR [DIR ...] --out FILE.
object_centroids(→ Dict[int, Tuple[float, float, float]])	label -> (y, x, equivalent diameter) for every object.
pair_sampling_weights(→ numpy.ndarray)	Weights that draw pairs evenly across how far their objects move.
predict_pair(→ Dict[str, numpy.ndarray])	The time head's prediction for one pair.
random_window(→ Tuple[List[numpy.ndarray], ...)	One size square, the SAME window of every frame and label.
scramble_test(→ Dict[str, float])	Shuffle frame t+1's labels and ask whether the model finds them.
time_targets(→ Dict[str, numpy.ndarray])	What the time head must predict for one pair of labelled frames.
timeflows_loss(→ object)	Masked loss for one batch.
track_masks_from_ctc(→ numpy.ndarray)	Full object masks labelled with their TRACK ids.
train_timeflows(→ List[float])	Train the time head, then the whole network, on track-labelled pairs.

Module Contents

class spacr.timeflows_model.CellposeSamFeatures(net)

Cellpose-SAM’s encoder as a feature extractor, sharing its weights.

Replays cellpose.vit.CPSAM.forward() up to the neck: 256 feature channels at 1/ps of the image size. The segmentation head stays the network’s own, so what the model already segments is not re-learned.

Parameters:

net – a cellpose.vit.CPSAM (CellposeModel(...).net).

Hold the network.

__call__(x)

Features for x (B, 3, H, W): (B, 256, H/ps, W/ps).

Parameters:

x – a batch of normalised images; cast to the encoder’s own dtype (Cellpose-SAM runs in bfloat16).

Returns:

the neck features, as float32 for the head.

parameters()

The encoder’s parameters, for freezing and the optimiser.

spacr.timeflows_model.TimeflowsNet(backbone, channels: int | None = None, ps: int | None = None)

The time head on a shared backbone, as a torch.nn.Module.

Both frames go through the same backbone; their features are joined and read by a small head that predicts three channels at full resolution: dy, dx (toward the successor’s centre, in diameters) and a successor logit.

Parameters:

• backbone – callable (B, 3, H, W) -> (B, C, H/ps, W/ps) with parameters(); CellposeSamFeatures for the real model.

• channels – C; read from backbone.channels when None.

• ps – the backbone’s downsampling; read from backbone.ps.

Returns:

the module.

spacr.timeflows_model.augment_pair(frames: Sequence[numpy.ndarray], labels: Sequence[numpy.ndarray], rng: numpy.random.Generator) → Tuple[List[numpy.ndarray], List[numpy.ndarray]]

One random flip and quarter-turn, the SAME for every frame and label.

A flip applied to frame t and not t+1 teaches a displacement that never happened, which is the easiest way to get a model that trains beautifully and tracks nothing. Targets are computed from the augmented labels afterwards, so the vectors follow automatically.

Parameters:

• frames – the images, (H, W) or (H, W, C).

• labels – the label images, (H, W).

• rng – the random generator.

Returns:

the augmented frames and labels.

spacr.timeflows_model.ctc_pairs(movie: str, sequence: str = '01', max_pairs: int | None = None, *, segmentation: str = 'ST', gaps: Sequence[int] = (1,)) → List[_Pair]

Consecutive-frame training pairs from one Cell Tracking Challenge movie.

Frames from <movie>/<seq>/t*.tif, full masks from the silver segmentation <movie>/<seq>_ST/SEG/man_seg*.tif relabelled by the tracking markers <movie>/<seq>_GT/TRA/man_track*.tif (track_masks_from_ctc()). A frame missing any of the three is skipped, and a pair is only formed from two consecutive frame numbers. Slice-mask filenames are ignored and duplicate frame numbers are rejected. With segmentation='GT', full masks are read from <seq>_GT/SEG. A source whose next-frame marker lacks an unambiguous full mask is censored for that pair, rather than labelled as a disappearance.

Parameters:

• movie – the movie folder (e.g. .../ctc_dic_hela_timelapse).

• sequence – '01' or '02'.

• max_pairs – at most this many pairs, spaced evenly through the movie and chosen before any file is read. Set this limit to bound the frames loaded from long movies; loading entire collections can require tens of gigabytes of memory.

• segmentation – 'ST' for silver masks (the training default), or 'GT' for supplied ground-truth full masks during validation.

• gaps – frame intervals to pair, (1,) by default. A gap g pairs frame n with n + g, so the same movie supplies larger displacements, the way a sparser acquisition would. Each gap gets its own max_pairs; frames shared between gaps are read once. A track ending inside the gap, including a dividing parent, has no successor.

Returns:

the pairs, gap by gap in the order given, each in time order.

Raises:

ValueError – sequence/limit, duplicate frame identities or annotation arrays are invalid.

spacr.timeflows_model.link_by_timeflows(labels_t: numpy.ndarray, labels_t1: numpy.ndarray, prediction: Dict[str, numpy.ndarray], *, min_successor: float = 0.5, max_distance: float = 1.0) → Dict[int, int]

Which object in t+1 each object in t becomes, from the model.

Each object’s pixels vote, through the predicted vectors, for its next centre. Their mean predictions are assigned jointly to object centres in t+1 with the Hungarian method. Objects below min_successor or with no allowed assignment remain unlinked. Distances are measured in source object diameters, with a minimum diameter of one pixel.

The distance gate is applied before assignment. Allowed distances are scaled to at most one; an unmatched choice costs just above one. A zero distance limit permits only exact centre matches, with zero cost. Thus an impossible edge cannot displace a valid link. The objective minimizes distance plus unmatched costs; it does not maximize the number of links.

Thresholds must be finite, with min_successor in [0, 1] and max_distance non-negative. If either frame has no objects, the result is empty without inspecting predictions. Otherwise prediction shapes must match the source frame, foreground vectors must be finite and foreground successor probabilities must lie in [0, 1]. Background values are ignored. Derived object centres must also be finite.

Parameters:

• labels_t – frame t’s objects, any ids.

• labels_t1 – frame t+1’s objects, any ids.

• prediction – from predict_pair().

• min_successor – the successor probability an object needs.

• max_distance – the furthest link, in the object’s diameters.

Returns:

{id in t: id in t+1}.

Raises:

ValueError – invalid thresholds, prediction shapes, foreground values or derived object centres.

spacr.timeflows_model.main(argv: Sequence[str] | None = None) → int

python -m spacr.timeflows_model --movies DIR [DIR ...] --out FILE.

Trains the time head on Cell Tracking Challenge movies (both sequences of each), starting from --base (a Cellpose checkpoint or cpsam), and saves the head and, after the full stage, the whole network.

Parameters:

argv – arguments; sys.argv[1:] when None.

Returns:

the exit status.

spacr.timeflows_model.object_centroids(labels: numpy.ndarray) → Dict[int, Tuple[float, float, float]]

label -> (y, x, equivalent diameter) for every object.

Parameters:

labels – an instance label image.

Returns:

the centroids and diameters.

spacr.timeflows_model.pair_sampling_weights(label_stack: numpy.ndarray, bins: int = 5) → numpy.ndarray

Weights that draw pairs evenly across how far their objects move.

Most consecutive frames are nearly still, so drawing pairs uniformly shows the head mostly easy pairs. Each pair’s median displacement over diameter falls in one of bins equal-width bins across the range seen, and a pair’s weight is the inverse of its bin’s size. Quantile bins were tried first and collapse when most pairs move alike, which is exactly the movie this exists for.

Parameters:

• label_stack – (T, H, W) track-labelled frames.

• bins – how many displacement bins.

Returns:

one weight per pair (t, t+1), summing to 1.

spacr.timeflows_model.predict_pair(net, frame_t: numpy.ndarray, frame_t1: numpy.ndarray, device: str = 'cpu') → Dict[str, numpy.ndarray]

The time head’s prediction for one pair.

Parameters:

• net – a trained TimeflowsNet().

• frame_t – frame t, normalised.

• frame_t1 – frame t+1, normalised.

• device – where to run.

Returns:

vector (2, H, W) and successor probability (H, W). The frame is read in TILE-pixel tiles with a quarter-tile overlap and the overlaps averaged; the vectors are in each object’s own diameters, so a tile needs no context beyond the object.

spacr.timeflows_model.random_window(frames: Sequence[numpy.ndarray], labels: Sequence[numpy.ndarray], rng: numpy.random.Generator, size: int = TILE) → Tuple[List[numpy.ndarray], List[numpy.ndarray]]

One size square, the SAME window of every frame and label.

Centred on a random object of the first label image, jittered by up to a quarter window, so windows hold cells rather than empty background. A window shared by both frames keeps every displacement true.

Parameters:

• frames – the images.

• labels – the label images.

• rng – the random generator.

• size – the window’s side.

Returns:

the windowed frames and labels.

spacr.timeflows_model.scramble_test(labels_t: numpy.ndarray, labels_t1: numpy.ndarray, predict: Callable[[numpy.ndarray, numpy.ndarray], Dict[str, numpy.ndarray]], frame_t: numpy.ndarray, frame_t1: numpy.ndarray, *, seed: int = 0) → Dict[str, float]

Shuffle frame t+1’s labels and ask whether the model finds them.

Scrambling removes the numeric-ID shortcut, but object positions and overlap remain informative. A good score alone does not prove learned motion. The share of objects linked to their true successor is the score; the plain IoU stitcher on the same pair is the reference the model has to beat.

Parameters:

• labels_t – frame t, track ids.

• labels_t1 – frame t+1, the same track ids.

• predict – fn(frame_t, frame_t1) -> prediction.

• frame_t – frame t.

• frame_t1 – frame t+1.

• seed – the shuffle’s seed.

Returns:

{'model': share, 'iou': share, 'objects': n}.

spacr.timeflows_model.time_targets(labels_t: numpy.ndarray, labels_t1: numpy.ndarray) → Dict[str, numpy.ndarray]

What the time head must predict for one pair of labelled frames.

Inside every object of frame t:

• vector – (dy, dx) from the pixel to its object’s centre in frame t+1, divided by the object’s diameter, so one number means the same thing for a large cell and a small one;

• successor – 1 where the object’s label is present in t+1, 0 where it is not (it died, divided – children take new labels in the movies checked – or left the field).

The vector is supervised only where there IS a successor (vector_weight); the successor flag everywhere inside an object (object_weight). Background is supervised by neither.

Parameters:

• labels_t – frame t’s labels, track ids.

• labels_t1 – frame t+1’s labels, the same ids.

Returns:

vector (2, H, W), successor, vector_weight and object_weight (H, W), all float32.

spacr.timeflows_model.timeflows_loss(output, targets: Dict[str, object]) → object

Masked loss for one batch.

Mean squared error on the vector where a successor exists, binary cross-entropy on the successor flag inside objects, background unsupervised.

Parameters:

• output – the net’s (B, 3, H, W) output.

• targets – tensors from time_targets(), batched.

Returns:

the scalar loss.

spacr.timeflows_model.track_masks_from_ctc(segmentation: numpy.ndarray, markers: numpy.ndarray) → numpy.ndarray

Full object masks labelled with their TRACK ids.

The Cell Tracking Challenge publishes tracking ground truth as MARKERS – a small blob inside each cell, labelled with its track id in every frame – and full outlines separately (the silver ST/SEG masks, labelled per frame). Only one-to-one assignments are retained: the object contains exactly one marker ID and that ID overlaps no other segmented object. Unmarked, merged and split assignments are excluded rather than guessed.

Parameters:

• segmentation – one frame’s instance labels, any ids.

• markers – the same frame’s TRA markers, labelled by track id.

Returns:

int64 segmentation relabelled by track id, 0 elsewhere.

Raises:

ValueError – annotations are not matching 2-D non-negative integer arrays, or marker IDs cannot be represented in int64.

spacr.timeflows_model.train_timeflows(net, pairs: Sequence[_Pair], *, head_steps: int = 100, full_steps: int = 100, lr_head: float = 0.001, lr_full: float = 1e-05, weights: numpy.ndarray | None = None, seed: int = 0, device: str = 'cpu', log: Callable[[str], None] | None = None, validation_pairs: Sequence[_Pair] | None = None, validation_every: int | None = None, on_validation: Callable[[dict], None] | None = None, lr_schedule: str = 'constant', consistency_weight: float = 0.0) → List[float]

Train the time head, then the whole network, on track-labelled pairs.

A two-stage curriculum: the backbone’s segmentation is already paid for, so it is frozen while the new head learns, then everything is trained at a low learning rate. Every pair is augmented identically on both frames before its targets are computed.

A training crop keeps supervision only for complete source masks and, when present in the full target frame, complete successor masks. A successor outside the tile is not a disappearance. Absences in the supplied full-frame labels remain supervised; this does not validate those annotations. Unusable crops are retried up to 32 times per step. Censoring prevents incorrect targets at crop boundaries but removes some fast-motion examples; it does not establish full-motion accuracy.

Parameters:

• net – from TimeflowsNet().

• pairs – the training pairs.

• head_steps – steps with the backbone frozen.

• full_steps – steps with everything trainable.

• lr_head – learning rate for the head stage.

• lr_full – learning rate for the full stage.

• weights – sampling weight per pair (pair_sampling_weights()); uniform when None.

• seed – the random seed.

• device – 'cpu' or 'cuda'.

• log – fn(line) for progress.

• validation_pairs – optional held-out full-frame pairs. Exact input overlap with training is rejected before any optimizer update.

• validation_every – updates between held-out checks; defaults to one epoch of len(pairs) sampled updates. Initial and stage-end checks are always included. Must be positive when supplied.

• on_validation – callback receiving each stratified validation report, including the stage, update count and current training loss. Optional validation never changes the returned loss-list contract.

• lr_schedule – 'constant' keeps each stage’s rate; 'cosine' warms up linearly over the first 5 % of a stage, then decays along a cosine to 5 % of the stage’s rate.

• consistency_weight – weight of the forward-backward term. When positive, every step also predicts the pair in reverse order and penalises, per object present in both crops, the forward and the backward displacement not cancelling (_cycle_consistency()). Zero trains on the forward loss alone.

Returns:

the loss at every step.

Raises:

ValueError – no usable supervision remains after 32 sampled crops for a step; inspect the full masks and motion relative to the tile; an unknown schedule or a negative or non-finite consistency weight.

Nested helpers

TimeflowsNet._Net.__init__(self)

Register the backbone and build the head and upsampler.

The backbone (or its net) is registered as encoder when it is a module, so its weights train and move with the network.

spacr/timeflows_model.py:395

TimeflowsNet._Net.forward(self, frame_t, frame_t1)

Join both frames’ backbone features and predict the time maps.

Parameters:

• frame_t – frame t, (B, 3, H, W).

• frame_t1 – frame t+1, the same shape.

Returns:

(B, 3, H, W): the vector (dy, dx) and the successor logit.

spacr/timeflows_model.py:418

TimeflowsNet._Net.head_parameters(self)

The new layers’ parameters: the head and the upsampler.

spacr/timeflows_model.py:414

augment_pair.apply(array)

Apply this call’s quarter-turns and optional flip to array.

spacr/timeflows_model.py:215

ctc_pairs.indexed(folder, prefix)

Map frame number to path for the prefix*.tif files in folder.

spacr/timeflows_model.py:1117

main.record_validation(report)

Flush each check immediately; completion is recorded only after saving.

spacr/timeflows_model.py:1260

predict_pair.starts(extent)

Tile origins covering extent with a quarter-tile overlap.

spacr/timeflows_model.py:784

train_timeflows.report_validation(stage, step)

Report one held-out check while preserving the training state.

spacr/timeflows_model.py:698

train_timeflows.run(steps, params, lr, frozen)

Train params for steps steps, the backbone frozen or not.

spacr/timeflows_model.py:719