Skip to content

mlvbench.tasks.binary_segmentation

Binary segmentation task.

BinarySegmentationEvaluator

Bases: Evaluator

Evaluator for binary segmentation tasks.

__init__(patch_size, save_predictions=8)

Initialize the evaluator.

Parameters:

Name Type Description Default
patch_size int

The patch size of the model, used to reconstruct the spatial layout of patches for visualization.

 required
save_predictions int

The maximum number of samples for which to keep raw predictions. Samples are taken in the order they are passed to update. Set to 0 to disable.

 8

finalize()

Finalize the evaluation and return summary, results and predictions.

update(batch, prediction, target, weight=None)

Update the evaluator with predictions and targets for a single batch.

Parameters:

Name Type Description Default
batch dict[str, Any]

A batch from the dataset. This evaluator expects a list of sample keys as batch["__key__"].

 required
prediction Tensor

The prediction from the probe with shape (B, N, 1) and dtype float32. Values are logits.

 required
target Tensor

The target for the task with shape (B, N, 1) and dtype float32 with binary values 0.0 and 1.0.

 required
weight Tensor | None

Optional per-patch weight tensor with shape (B, N, 1) and dtype float32. Patches with weight 0.0 are excluded from the metrics.

 None

BinarySegmentationTask

Bases: Task

Binary segmentation task.

Predict whether each patch belongs to the foreground (segmentation > 0).

__init__(patch_size=None)

Initialize the task.

Parameters:

Name Type Description Default
patch_size int | None

The patch size of the model. Alternatively, this can be specified later by calling configure. The Trainer always calls configure, so the the patch size doesn't have to be specified for probe training.

 None

configure(model)

Configure the task for the given model.

criterion(prediction, target, weight=None)

Loss function for the task.

Parameters:

Name Type Description Default
prediction Tensor

Probe output logits with shape (B, N, 1) and dtype float32.

 required
target Tensor

Binary target labels with shape (B, N, 1) and dtype float32.

 required
weight Tensor | None

Optional per-patch weights with shape (B, N, 1) and dtype float32. Patches with weight 0.0 are excluded from the loss.

 None

Returns:

Type Description
Tensor

The scalar binary cross-entropy loss, a float32 tensor of shape ().

get_evaluator(save_predictions=8)

Return an evaluator for the task.

Parameters:

Name Type Description Default
save_predictions int

The maximum number of samples for which the evaluator keeps raw predictions. Set to 0 to disable.

 8

prepare_target(batch)

Prepare the target for the probe.

Parameters:

Name Type Description Default
batch dict[str, Tensor]

A batch from the dataset.

 required

Returns:

Type Description
tuple[Tensor, Tensor]

A tuple of (target, weights) where target has shape (B, N, 1) and dtype float32 with 1.0 where the patch belongs to the foreground and 0.0 otherwise, and weight has shape (B, N, 1) and dtype float32 with 0.0 for patches that contain more than one segment label and 1.0 otherwise.