时间戳

首页 > 解决方案 > 如何获得连接组件的最大值和坐标？

python - 如何获得连接组件的最大值和坐标？

问题描述

例如，给定一个预测的概率图，例如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]])

如何在 中找到两个连通分量的两个最大概率 ( 0.9, 0.9) 和坐标 ( ) ？(1,1), (2,5)a

标签： pythonnumpyopencv-python

解决方案

这是我的答案，但可能太复杂了。

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

推荐阅读