所用图片:
赋予照片框数据
0 187 90 586 607 0.94
0 120 116 600 370 0.68
1 511 185 961 418 0.58
0 340 145 568 478 0.47
1 524 70 920 565 0.92
代码如下:
import cv2
import numpy as np
from utilsW.utils import cvShow
def draw_box(image, bboxes_list):
overlay_color = {
'0': (0, 255, 0),
'1': (255, 0, 0)
}
for coord in bboxes_list:
#类别
class_name = coord[0]
start_point = (int(coord[1]), int(coord[2]))
end_point = (int(coord[3]), int(coord[4]))
prob = float(coord[5])
text_start_point = (int(coord[1]), int(coord[2]) - 10)
image = cv2.rectangle(image, start_point, end_point, overlay_color[class_name], 2)
image = cv2.putText(image, str(prob), text_start_point, None, 0.8, overlay_color[class_name], 2)
cv2.imshow("im", image)
cv2.waitKey(0)
cv2.destroyAllWindows()
def IOU(bboxes1, bboxes2):
bboxes1 = [int(i) for i in bboxes1]
bboxes2 = [int(i) for i in bboxes2]
xA = max(bboxes1[0], bboxes2[0])
yA = max(bboxes1[1], bboxes2[1])
xB = min(bboxes1[2], bboxes2[2])
yB = min(bboxes1[3], bboxes2[3])
intersection_area = max(0, xB - xA + 1) * max(0, yB - yA + 1)
box1_area = (bboxes1[2] - bboxes1[0] + 1) * (bboxes1[3] - bboxes1[1] + 1)
box2_area = (bboxes2[2] - bboxes2[0] + 1) * (bboxes2[3] - bboxes2[1] + 1)
iou = intersection_area / float(box1_area + box2_area - intersection_area)
return iou
def NMS(src_copy, file, conf_threshold, iou_threshold):
req_bboxes, final_boxes = [], []
for coord in file:
prob = float(coord[5])
if prob > conf_threshold:
req_bboxes.append(coord)
# sorting the bounding boxes based on probability score
bboxes_sorted = sorted(req_bboxes, reverse=True, key=lambda x: x[5])
while len(bboxes_sorted) > 0:
# removing the best probability bounding box
box = bboxes_sorted.pop(0)
for b in bboxes_sorted:
# comparing with the same class
if box[0] == b[0]:
iou = IOU(box[1:-1], b[1:-1])
if iou >= iou_threshold:
# if IOU is large then discard the box with lowest probability
bboxes_sorted.remove(b)
print(len(bboxes_sorted))
final_boxes.append(box)
return final_boxes
if __name__ == '__main__':
#1 设置初始参数
conf_threshold = 0.5
iou_threshold = 0.4
#2 读取图片
src = cv2.imread("zoraya.jpg")
src_copy = src.copy()
cvShow("src", src)
#3 读取坐标框信息
with open("coordinates.txt") as F:
file = F.readlines()
file = [i[:-1].split(' ') for i in file]
#现在图片上画框,看一下效果
draw_box(src, file)
#非极大值抑制筛选点
final_boxes = NMS(src_copy, file, conf_threshold, iou_threshold)
draw_box(src_copy, final_boxes)
代码效果: