Python(Numpy)实现非极大值抑制

1.Numpy的几个骚操作

(1).np.maximum的使用

import numpy as np
box = [3,5,7,9]  # A single box with a first coordinate of 3
boxes = np.array([[1, 4], 
                  [5, 2], 
                  [2, 6]])  # An array of multiple boxes
#把box把3拿出来和boxes的第一列，逐个比较，返回最大值list
result = np.maximum(box[0], boxes[:, 0])
#result=[3 5 3]
print(result)

(2). np.argsort的使用：

import numpy as np
scores = np.array([0.5, 0.2, 0.8])
list1 = np.argsort(scores)
list2=np.argsort(scores)[::-1]
#list1=[1 0 2]
print(f"list1={list1}")
#list2=[2 0 1]
print(f"list2={list2}")
#scores=[0.5 0.2 0.8],原来的顺序没改
print(f"scores={scores}")

(3).向量和矩阵的乘法和加法

import numpy as np
box = [3,5]  # A single box with a first coordinate of 3
boxes = np.array([[1, 4], 
                  [5, 2], 
                  [2, 6]])  # An array of multiple boxes

"""
矩阵乘法是"@"不是"*"
*不是矩阵乘法，把boxes的每一行都对应乘以box
box*boxes=[[ 3 20]
           [15 10]
           [ 6 30]]
"""
print(box*boxes)


"""
boxes的每一列都对应乘以box
box+boxes=[[ 4  9]
           [ 8  7]
           [ 5 11]]
"""
print(box+boxes)

2.非极大值抑制

def non_max_suppression(boxes, scores, threshold):
    """
    Apply non-maximum suppression to eliminate redundant bounding boxes.

    Parameters:
    - boxes: A list of bounding boxes in the format (x1, y1, x2, y2).
    - scores: A list of confidence scores for each bounding box.
    - threshold: The IoU (Intersection over Union) threshold for suppression.

    Returns:
    - selected_indices: A list of indices for the selected bounding boxes.
    """
    if len(boxes) == 0:
        return []

    # Convert the bounding boxes to the format (x1, y1, x2, y2)
    boxes = np.array(boxes)

    # Sort indices based on confidence scores (in descending order)
    #sorted_indices=[2,0,1]
    sorted_indices = np.argsort(scores)[::-1]

    selected_indices = []

    while len(sorted_indices) > 0:
        # Pick the bounding box with the highest confidence score
        i = sorted_indices[0]
        selected_indices.append(i)


        # Calculate IoU with the remaining bounding boxes
        ious = calculate_iou(boxes[i], boxes[sorted_indices[1:]])

        # Filter out bounding boxes with IoU higher than the threshold
        #np.where(ious <= threshold) is ([])
        #type(remaining_indices)=ndarray
        remaining_indices = np.where(ious <= threshold)[0]
        print(f"remaining_indices={remaining_indices}")
        print(f"remaining_indices+1={remaining_indices+1}")
        #i=2的哪个框已经进入selected
        sorted_indices = sorted_indices[remaining_indices + 1]
        print("sorted_indices=",sorted_indices)

    return selected_indices

def calculate_iou(box, boxes):
    """
    Calculate the Intersection over Union (IoU) between a bounding box and a list of bounding boxes.

    Parameters:
    - box: The bounding box in the format (x1, y1, x2, y2).
    - boxes: A list of bounding boxes.

    Returns:
    - ious: A numpy array of IoU values.
    """
    #计算intersection的4个坐标
    x1 = np.maximum(box[0], boxes[:, 0])
    y1 = np.maximum(box[1], boxes[:, 1])
    x2 = np.minimum(box[2], boxes[:, 2])
    y2 = np.minimum(box[3], boxes[:, 3])
    #没有交集的框的IoU就是0
    intersection = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    area_box = (box[2] - box[0]) * (box[3] - box[1])
    area_boxes = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    union = area_box + area_boxes - intersection

    ious = intersection / union

    return ious

# Example usage:
import numpy as np
if __name__=="__main__":
    # Example bounding boxes and confidence scores
    #the 4 number in each box is the coordinate of left-top corner
    #and  right-bottom corner
    boxes = [(10, 20, 50, 60), (15, 25, 55, 65), (30, 40, 70, 80)]
    scores = [0.9, 0.75, 0.95]

    # Set the IoU threshold for non-maximum suppression
    iou_threshold = 0.5

    # Apply non-maximum suppression
    selected_indices = non_max_suppression(boxes, scores, iou_threshold)

    # Display the selected bounding boxes
    for i in selected_indices:
        print(f"Selected Box: {boxes[i]}, Score: {scores[i]}")

Selected Box: (30, 40, 70, 80), Score: 0.95
Selected Box: (10, 20, 50, 60), Score: 0.9