Source code for blocks.utils.rpn

from typing import Optional, List, Sequence

import numpy as np

from ..datasets import RPNDataset


__all__ = ['decode_rpn_predictions']


def decode_rpn_predictions(dataset: RPNDataset,
                           classifier_probabilities: np.ndarray,
                           regression_predictions: np.ndarray,
                           non_maximum_suppression: Optional[float]=0.3,
                           min_probability: Optional[float]=None,
                           top_k: Optional[int]=None) -> List[Sequence[float]]:
    """
    Decode the given RPN predictions to a list of regions.

    Optionally:
        - use non maximum suppression
        - filter out regions with low probability
        - output only top k regions

    :param dataset: RPN dataset which created the anchors
    :param classifier_probabilities: predicted classifier probabilities (in 0-1 interval)
    :param regression_predictions: predicted diffs regression values
    :param non_maximum_suppression: if specified, use nms with the given threshold
    :param min_probability: if specified, filter out regions with probability lower than the specified threshold
    :param top_k: if specified, output only top k regions
    :return: list of decoded regions
    """
    predicted_labels = np.round(classifier_probabilities)

    # obtain indices, probs and regions with applied diffs
    positives_indices = np.argwhere(predicted_labels == 1)
    positives_probs = [classifier_probabilities[tuple(ix)] for ix in positives_indices]
    positive_regions = [dataset.apply_diff(dataset.get_anchor(ix), regression_predictions[tuple(ix)])
                        for ix in positives_indices]

    # zip and sort the predictions by probs
    positives = zip(positives_indices, positives_probs, positive_regions)
    positives = sorted(positives, key=lambda t: t[1], reverse=True)

    # non-maximum suppression
    if non_maximum_suppression is not None:
        ix = 0
        while True:
            if ix >= len(positives):
                break
            best = positives[ix][2]
            positives = positives[:ix + 1] + [region for region in positives[ix + 1:]
                                              if dataset.anchor_region_overlap(best, region[2]) < non_maximum_suppression]
            ix += 1

    # min_probability filtering
    if min_probability is not None:
        positives = [positive for positive in positives if positive[1] > min_probability]

    # top-k filtering
    if top_k is not None:
        positives = positives[:top_k]

    positive_regions = [positive[2] for positive in positives]
    return positive_regions