cellmap_flow.finetune.losses

The finetune losses, and the masked means they are all made of.

Every loss here is averaged over the voxels a mask supervises, not over the patch: a correction labels a few voxels of it, and the rest must not count. masked_mean is that average, balanced_mean gives the foreground and the background half the weight each, and the losses and the distillation term below are built from the two. IntervalLoss, a distance model’s loss on scribbles, works in logits instead of probabilities (as_logits). A flow model’s loss, instance_flows.FlowLoss, is built from masked_mean too, and lives beside the target it compares with.

The float operations are kept in the order the trainer always used, so a run computes the same numbers bit for bit (tests/finetune/test_training_snapshot.py).

Classes

`DiceLoss`	Dice Loss for segmentation tasks.
`CombinedLoss`	Combined Dice + BCE loss for better convergence.
`MarginLoss`	Margin-based loss for sparse/scribble annotations.
`IntervalLoss`	A distance model's loss on scribbles: stay within the bounds the paint implies.

Functions

`masked_mean`(x, weight)	Mean of `x` weighted by `weight` (a 0/1 mask, usually); 0 when nothing is weighted.
`balanced_mean`(x, hard_target, mask)	The foreground's and the background's masked means, averaged.
`as_probabilities`(pred, model_has_sigmoid)	The model's output as probabilities.
`as_logits`(pred, model_has_sigmoid)	The model's output as logits: `as_probabilities`' inverse.
`soft_target_entropy`(target[, eps])	Per-voxel BCE that a perfectly calibrated prediction still pays.
`distillation_loss`(student, teacher, scope[, ...])	Mean squared distance of the student's output from the teacher's, over `scope`.

Module Contents

cellmap_flow.finetune.losses.masked_mean(x, weight): Mean of x weighted by weight (a 0/1 mask, usually); 0 when nothing is weighted.

cellmap_flow.finetune.losses.balanced_mean(x, hard_target, mask)

The foreground’s and the background’s masked means, averaged.

Each class counts the same however many voxels it has. The split is by the target as annotated (0/1), before any label smoothing: split by the smoothed target, every background voxel carried s/2 of foreground weight.

cellmap_flow.finetune.losses.as_probabilities(pred, model_has_sigmoid)

The model’s output as probabilities.

Left alone when the model already ends in a sigmoid (the cellmap _distance_ UNets do); a second sigmoid would squash [0, 1] into [0.5, 0.73] and make a well-fitting prediction look like a constant.

cellmap_flow.finetune.losses.as_logits(pred, model_has_sigmoid)

The model’s output as logits: as_probabilities’ inverse.

A model ending in a sigmoid is taken back through it. The clamp keeps log(0) out; it costs no gradient short of |logit| ~16, far past where any distance bound sits (SATURATION_SIGMAS is a logit of 6).

cellmap_flow.finetune.losses.soft_target_entropy(target, eps=1e-07)

Per-voxel BCE that a perfectly calibrated prediction still pays.

-(t log t + (1-t) log(1-t)): zero for hard 0/1 targets, log 2 at t = 0.5. On soft targets (distance, smoothed labels) this is the floor of the BCE curve, and it moves with the batch, so a “flat” BCE can be a model sitting on its floor. Report the loss minus this instead.

cellmap_flow.finetune.losses.distillation_loss(student, teacher, scope, anchor_mask=None, unlabeled_mask=None)

Mean squared distance of the student’s output from the teacher’s, over scope.

"anchor": the good regions the user vouched for, anchor_mask. It is about location, so one channel, broadcast over the output’s.
"unlabeled": the voxels the supervised loss leaves out, unlabeled_mask (1 - the supervised mask).
"all": every voxel.

The two are compared as the model emits them: logits, or probabilities for a model that ends in a sigmoid.

class cellmap_flow.finetune.losses.DiceLoss(smooth: float = 1.0)

Dice Loss for segmentation tasks.

Dice loss is effective for imbalanced datasets where the target class may be sparse (e.g., mitochondria in EM images).

Formula: 1 - (2 * |X ∩ Y| + smooth) / (|X| + |Y| + smooth)

smooth = 1.0

apply_sigmoid = True

forward(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor: Dice loss of (B, C, Z, Y, X) logits (or probabilities) against 0/1 targets, over mask.

class cellmap_flow.finetune.losses.CombinedLoss(dice_weight: float = 0.5, bce_weight: float = 0.5)

Combined Dice + BCE loss for better convergence.

Uses both Dice loss (for overlap) and BCE loss (for pixel-wise accuracy).

dice_loss

bce_loss

dice_weight = 0.5

bce_weight = 0.5

forward(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor: Weighted Dice + BCE of (B, C, Z, Y, X) logits against 0/1 targets, over mask.

class cellmap_flow.finetune.losses.MarginLoss(margin: float = 0.3, balance_classes: bool = False)

Margin-based loss for sparse/scribble annotations.

Only penalizes predictions on the wrong side of a margin threshold. For post-sigmoid outputs in [0, 1]: - Foreground (target=1): loss = relu(threshold - pred)^2, threshold = 1 - margin - Background (target=0): loss = relu(pred - margin)^2 - No loss when prediction is already correct with sufficient confidence.

margin = 0.3

balance_classes = False

apply_sigmoid = True

forward(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor

class cellmap_flow.finetune.losses.IntervalLoss(sigma_nm, voxel_size_nm, slope_weight=1.0, balance_classes=False, slack_voxels=1.0, sign_weight=5.0)

A distance model’s loss on scribbles: stay within the bounds the paint implies.

After iSDF (Ortiz et al., RSS 2022). IntervalTargetTransform gives each painted voxel a lower and an upper bound on its logit z (the distance target is z = 2d/sigma). Per painted voxel the loss is

relu(lower - slack - z) + relu(z - upper - slack)

sign_weight * relu(-s * z - slack)

with s the painted side (+1 foreground, -1 background): zero inside the bounds, linear outside them, and much steeper on the wrong side of the boundary. The slack, a voxel of distance by default, forgives a stroke that strays a voxel over an edge. Unpainted voxels get none of it: distillation and the anchor patches hold them.

Bounds alone let the field collapse: a step at the boundary meets every lower bound and most upper ones, and is what margin loss made of a distance model. So a second term limits the slope, everywhere in the patch: the mean of relu(|grad z| / max_slope - 1), from central differences scaled by the voxel size in nm, with max_slope = SLOPE_TOLERANCE * 2/sigma per nm, the steepest a distance field gets in logits. One-sided: the field may be flatter. z is clamped to the saturation (SATURATION_SIGMAS) first, so a model confident far from any boundary is not asked to bring its far field in.

Parameters:

sigma_nm – from the IntervalTargetTransform that made the bounds.
voxel_size_nm – from the IntervalTargetTransform that made the bounds.
slope_weight – the slope term’s weight against the bounds’ (default 1).
balance_classes – average foreground and background voxels separately.
slack_voxels – the slack, in voxels of the finest axis (default 1).
sign_weight – the extra slope of the wrong-side penalty (default 5).

voxel_size_nm

slack

max_slope

saturation = 6.0

slope_weight = 1.0

balance_classes = False

sign_weight = 5.0

bounds_term(z, bounds, mask): The mean penalty for leaving the bounds, over the painted voxels.

slope_term(z): The mean excess of |grad z| over max_slope, as a fraction of it, over the patch’s interior.

forward(z: torch.Tensor, bounds: torch.Tensor, mask: torch.Tensor) → torch.Tensor: The loss of (B, C, Z, Y, X) logits z against IntervalTargetTransform’s bounds and mask.