How to get the max value and coordinates of a connected component?-CodePudding

For example, given a predicted probability map, like a

a = np.array([[0.1, 0.2, 0.3, 0.0, 0.0, 0.0],
              [0.1, 0.92, 0.3, 0.0, 0.2, 0.1],
              [0.1, 0.9, 0.3, 0.0, 0.7, 0.89],
              [0.0, 0.0, 0.0, 0.0, 0.4, 0.5],
              [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])

How can I find two max probability (0.9, 0.9) and coordinates ((1,1), (2,5)) of two connected components in a?

CodePudding user response：

Use np.where or np.argwhere

>>> np.unique(a)[-2:]
array([0.89, 0.92])

>>> np.where(np.isin(a, np.unique(a)[-2:]))
(array([1, 2]), array([1, 5]))

# OR

>>> np.argwhere(np.isin(a, np.unique(a)[-2:]))
array([[1, 1],
       [2, 5]])

CodePudding user response：

Here is my answer, but maybe too complicated.

def nms_cls(loc, cls):
    """
    Find the max class and prob point in a mask
    :param loc: binary prediction with 0 and 1 (h, w)
    :param cls: multi-classes prediction with prob (c, h, w)
    :return: list of tuple (class, prob, coordinate)
    """
    prob = np.max(cls, axis=0) # (H, W)
    cls_idx = np.argmax(cls, axis=0)
    point_list = []

    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(loc, connectivity=8, ltype=None)
    for i in range(num_labels):
        # get the mask of connected component i
        label_i = np.copy(labels)
        label_i[label_i != i] = 0
        label_i[label_i > 0] = 1
        prob_mask_i = prob * label_i

        # get max prob's coords and class
        state_i = {}
        state_i['coord'] = np.unravel_index(prob_mask_i.argmax(), prob_mask_i.shape)
        state_i['cls'] = cls_idx[state_i['coord'][0], state_i['coord'][1]]
        state_i['prob'] = prob[state_i['coord'][0], state_i['coord'][1]]
        point_list.append(state_i)
    return point_list