def test_pro() -> None:
    """Checks if PRO metric computes the (macro) average of the per-region overlap."""
    labels = torch.Tensor(
        [
            [
                [0, 0, 0, 0, 0],
                [1, 0, 0, 0, 0],
                [0, 0, 0, 0, 0],
                [1, 1, 0, 0, 0],
                [0, 0, 0, 0, 0],
                [1, 1, 1, 0, 0],
                [0, 0, 0, 0, 0],
                [1, 1, 1, 1, 0],
                [0, 0, 0, 0, 0],
                [1, 1, 1, 1, 1],
            ],
        ],
    )
    # ground truth mask is int type
    labels = labels.type(torch.int32)

    preds = (torch.arange(10) / 10) + 0.05
    # metrics receive squeezed predictions (N, H, W)
    preds = preds.unsqueeze(1).repeat(1, 5).view(1, 10, 5)

    thresholds = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
    targets = [1.0, 0.8, 0.6, 0.4, 0.2, 0.0]
    for threshold, target in zip(thresholds, targets, strict=True):
        pro = PRO(threshold=threshold)
        pro.update(preds, labels)
        assert pro.compute() == target
Checks if PRO metric computes the (macro) average of the per-region overlap.
tests/unit/metrics/test_pro.py