import torch


def top_n_predictions_maskrcnn(predictions, n):
    """
    Select top n predictions based on score

    Arguments:
        predictions (dictionary): the result of the computation by the model.
            It should contain the field `scores`.

    Returns:
        prediction (dictionary): the detected objects. Additional information
            of the detection properties can be found in the fields of
            the BoxList via `prediction.fields()`
    """
    scores = predictions["scores"]
    n = min(n, len(scores))
    keep = torch.zeros(scores.shape).bool()
    topk, indices = torch.topk(scores, n)
    keep[indices] = True
    predictions = {key:predictions[key][keep] for key in predictions.keys()}
    # scores = predictions["scores"]
    # _, idx = scores.sort(0, descending=True)
    return predictions