NMS代码实现

def iou(the_first_score_location, the_rest_score_location):
    if the_first_score_location[0] > the_rest_score_location[0] and the_first_score_location[0] > \
            the_rest_score_location[2]:
        return 0

    elif the_first_score_location[2] < the_rest_score_location[0] and the_first_score_location[2] < \
            the_rest_score_location[2]:
        return 0

    elif the_first_score_location[1] > the_rest_score_location[1] and the_first_score_location[1] > \
            the_rest_score_location[3]:
        return 0

    elif the_first_score_location[3] < the_rest_score_location[1] and the_first_score_location[3] < \
            the_rest_score_location[3]:
        return 0

    x_location_past = []
    y_location_past = []
    x_location_past.append(the_first_score_location[0])
    x_location_past.append(the_first_score_location[2])
    x_location_past.append(the_rest_score_location[0])
    x_location_past.append(the_rest_score_location[2])
    y_location_past.append(the_first_score_location[1])
    y_location_past.append(the_first_score_location[3])
    y_location_past.append(the_rest_score_location[1])
    y_location_past.append(the_rest_score_location[3])
    x_location_new = sorted(x_location_past)
    y_location_new = sorted(y_location_past)
    w = x_location_new[2] - x_location_new[1]
    h = y_location_new[2] - y_location_new[1]
    union_set = w * h
    the_first_score_location_area = (the_first_score_location[2] - the_first_score_location[0]) * (
            the_first_score_location[3] - the_first_score_location[1])
    the_rest_score_location_area = (the_rest_score_location[2] - the_rest_score_location[0]) * (
            the_rest_score_location[3] - the_rest_score_location[1])
    intersection = the_first_score_location_area + the_rest_score_location_area - union_set
    iou_area = union_set / intersection
    return iou_area

animal = [(0, 0, 100, 100, 0.92, 0), (10, 10, 90, 90, 0.94, 0), (20, 20, 110, 110, 0.89, 0),
          (20, 20, 120, 120, 0.82, 0),
          (150, 20, 250, 120, 0.83, 0), (160, 30, 260, 130, 0.73, 0), (200, 200, 300, 300, 0.91, 1),
          (210, 210, 310, 310, 0.93, 1)]
threshold = 0.3
the_choose_location = []

while True:
    animal_index_and_value_outside = [one for one in enumerate(animal)]
    the_first_location = []
    animal_index_and_value_sort = sorted(animal_index_and_value_outside, key=lambda x: x[1][4], reverse=True)
    the_choose_location.append(animal_index_and_value_sort[0][1])
    the_first_location.append(animal_index_and_value_sort[0][1])
    animal_index_and_value_outside.remove(animal_index_and_value_sort[0])
    animal.remove(animal_index_and_value_sort[0][1])
    the_rest_location = []
    for i in animal_index_and_value_outside:
        if i[1][5] == animal_index_and_value_sort[0][1][5]:
            the_rest_location.append(i[1])
    for i in the_rest_location:
        if iou(the_first_location[0][0:4],i[0:4])>threshold:
            animal.remove(i)
    if len(animal)==0:
        break

print(the_choose_location)