选择性搜索算法+非极大值抑制Python实现

以下是选择性搜索算法和非极大值抑制的Python实现：

选择性搜索算法：

import cv2
import selectivesearch

# 读取图像
image = cv2.imread('image.jpg')

# 运行选择性搜索算法
selective_search = selectivesearch.selective_search(image, scale=500, sigma=0.9, min_size=10)

# 获取候选框
boxes = selective_search[0]

# 打印候选框数量
print("候选框数量：", len(boxes))

# 显示候选框
for box in boxes:
    x, y, w, h = box
    cv2.rectangle(image, (x, y), (x+w, y+h), (0, 255, 0), 2)

# 显示图像
cv2.imshow("Selective Search", image)
cv2.waitKey(0)
cv2.destroyAllWindows()

非极大值抑制算法：

import numpy as np

def non_max_suppression(boxes, overlap_thresh):
    if len(boxes) == 0:
        return []

    # 提取候选框的坐标和概率
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    scores = boxes[:, 4]

    # 计算候选框的面积
    area = (x2 - x1 + 1) * (y2 - y1 + 1)

    # 根据概率降序排序
    idxs = np.argsort(scores)[::-1]

    # 保留的候选框列表
    keep = []

    while len(idxs) > 0:
        # 保留当前概率最大的候选框
        i = idxs[0]
        keep.append(i)

        # 计算当前候选框与其他候选框的重叠度
        xx1 = np.maximum(x1[i], x1[idxs[1:]])
        yy1 = np.maximum(y1[i], y1[idxs[1:]])
        xx2 = np.minimum(x2[i], x2[idxs[1:]])
        yy2 = np.minimum(y2[i], y2[idxs[1:]])

        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        overlap = (w * h) / area[idxs[1:]]

        # 删除重叠度大于阈值的候选框
        idxs = np.delete(idxs, np.concatenate(([0], np.where(overlap > overlap_thresh)[0] + 1)))

    # 返回保留的候选框
    return boxes[keep]

# 假设候选框列表为boxes，重叠度阈值为0.5
boxes = np.array([[10, 10, 50, 50, 0.9], [20, 20, 60, 60, 0.8], [30, 30, 70, 70, 0.7]])
overlap_thresh = 0.5

# 应用非极大值抑制算法
selected_boxes = non_max_suppression(boxes, overlap_thresh)

# 打印保留的候选框
print("保留的候选框：", selected_boxes)