nms python代码_soft-nms python代码

自己写的soft nms的python代码，这里只是调整评分的功能。

def soft_nms(dicts):
    for label, boxes in dicts.items():
        boxesArray = np.asarray(boxes)

        x1 = boxesArray[:, 0]
        y1 = boxesArray[:, 1]
        x2 = boxesArray[:, 2]
        y2 = boxesArray[:, 3]
        scores = boxesArray[:,4]
        areas = (x2-x1+1)*(y2-y1+1)
        index = scores.argsort()[::-1]

        while index.size>0:
            box = boxes[index[0]]
            x1tmp = np.maximum(x1[index[0]], x1[index[1:]])
            y1tmp = np.maximum(y1[index[0]], y1[index[1:]])
            x2tmp = np.minimum(x2[index[0]], x2[index[1:]])
            y2tmp = np.minimum(y2[index[0]], y2[index[1:]])
            interArea = np.maximum(x2tmp - x1tmp,0.0)*np.maximum(y2tmp - y1tmp, 0.0)
            iou = interArea/(areas[index[0]] + areas[index[1:]] -interArea)
            for i in range(1,len(iou)+1):
                if iou[i-1]>0.2:
                    #boxes[index[i]][4] *= 1-iou[i-1]
                    ov = 1-iou[i-1]
                    sigma = 0.5
                    boxes[index[i]][4] *= np.exp(-(ov*ov)/sigma)
            index = np.delete(index,[0])

        dicts[label] = boxes
        return dicts

if __name__ == "__main__":
    predicts_dict = {"cup": [[2, 3, 30, 60, 0.12], [1, 4, 31, 62, 0.44], [1, 4, 31, 61, 0.11],[105, 107, 111, 109, 0.18], [123, 111, 167, 190, 0.23]]}
    print(soft_nms(predicts_dict))