import torch


def categorical_mean_average_precision(preds, targets, pred_name='output',
                                       target_name='map'):
    """
    Computes the mean average precision from predictions given as logits, for a
    categorical segmentation task.
    """
    preds = preds[pred_name]
    gt = targets[target_name]
    # order classes by predicted probability
    top_predictions = torch.argsort(preds, -1, descending=True)
    # find the correct class in there
    top_correct = (top_predictions == gt[..., None])
    # determine the rank of the correct class (requires conversion from bool)
    rank = top_correct.byte().argmax(-1)
    # the average of the inverse ranks is the mean average precision
    # (but averaging over the batch items happens outside of this)
    return (1. / (1 + rank))