def get_segmentation_statistics(
    outputs: torch.Tensor,
    targets: torch.Tensor,
    class_dim: int = 1,
    threshold: float = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Computes true positive, false positive, false negative
    for a multilabel segmentation problem.

    Args:
        outputs: [N; K; ...] tensor that for each of the N examples
            indicates the probability of the example belonging to each of
            the K classes, according to the model.
        targets:  binary [N; K; ...] tensor that encodes which of the K
            classes are associated with the N-th input
        class_dim: indicates class dimention (K) for
            ``outputs`` and ``targets`` tensors (default = 1)
        threshold: threshold for outputs binarization

    Returns:
        Segmentation stats

    Example:

    .. code-block:: python

        import torch
        from catalyst import metrics

        size = 4
        half_size = size // 2
        shape = (1, 1, size, size)
        empty = torch.zeros(shape)
        full = torch.ones(shape)
        left = torch.ones(shape)
        left[:, :, :, half_size:] = 0
        right = torch.ones(shape)
        right[:, :, :, :half_size] = 0
        top_left = torch.zeros(shape)
        top_left[:, :, :half_size, :half_size] = 1
        pred = torch.cat([empty, left, empty, full, left, top_left], dim=1)
        targets = torch.cat([full, right, empty, full, left, left], dim=1)

        metrics.get_segmentation_statistics(
            outputs=pred,
            targets=targets,
            class_dim=1,
            threshold=0.5,
        )
        # (
        #     tensor([ 0.,  0.,  0., 16.,  8.,  4.]),  # per class TP
        #     tensor([0., 8., 0., 0., 0., 0.]),        # per class FP
        #     tensor([16.,  8.,  0.,  0.,  0.,  4.]),  # per class TN
        # )
    """
    assert outputs.shape == targets.shape, (
        f"targets(shape {targets.shape})"
        f" and outputs(shape {outputs.shape}) must have the same shape"
    )
    if threshold is not None:
        outputs = (outputs > threshold).float()

    n_dims = len(outputs.shape)
    dims = list(range(n_dims))
    # support negative index
    if class_dim < 0:
        class_dim = n_dims + class_dim
    dims.pop(class_dim)

    sum_per_class = partial(torch.sum, dim=dims)

    tp = sum_per_class(outputs * targets)
    class_union = sum_per_class(outputs) + sum_per_class(targets)
    class_union -= tp
    fp = sum_per_class(outputs * (1 - targets))
    fn = sum_per_class(targets * (1 - outputs))
    return tp, fp, fn

    Computes true positive, false positive, false negative
    for a multilabel segmentation problem.

    Args:
        outputs: [N; K; ...] tensor that for each of the N examples
            indicates the probability of the example belonging to each of
            the K classes, according to the model.
        targets:  binary [N; K; ...] tensor that encodes which of the K
            classes are associated with the N-th input
        class_dim: indicates class dimention (K) for
            ``outputs`` and ``targets`` tensors (default = 1)
        threshold: threshold for outputs binarization

    Returns:
        Segmentation stats

    Example:

    .. code-block:: python

        import torch
        from catalyst import metrics

        size = 4
        half_size = size // 2
        shape = (1, 1, size, size)
        empty = torch.zeros(shape)
        full = torch.ones(shape)
        left = torch.ones(shape)
        left[:, :, :, half_size:] = 0
        right = torch.ones(shape)
        right[:, :, :, :half_size] = 0
        top_left = torch.zeros(shape)
        top_left[:, :, :half_size, :half_size] = 1
        pred = torch.cat([empty, left, empty, full, left, top_left], dim=1)
        targets = torch.cat([full, right, empty, full, left, left], dim=1)

        metrics.get_segmentation_statistics(
            outputs=pred,
            targets=targets,
            class_dim=1,
            threshold=0.5,
        )
        # (
        #     tensor([ 0.,  0.,  0., 16.,  8.,  4.]),  # per class TP
        #     tensor([0., 8., 0., 0., 0., 0.]),        # per class FP
        #     tensor([16.,  8.,  0.,  0.,  0.,  4.]),  # per class TN
        # )
    
catalyst/metrics/functional/_segmentation.py