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)
NMS代码实现
最新推荐文章于 2024-07-10 10:23:09 发布