import cv2
import numpy as np
import onnxruntime
import xlsxwriter
import os
# coco80类别
CLASSES = ['card']
class YOLOV5():
def __init__(self, onnxpath):
self.onnx_session = onnxruntime.InferenceSession(onnxpath)
self.input_name = self.get_input_name()
self.output_name = self.get_output_name()
def get_input_name(self):
input_name = []
for node in self.onnx_session.get_inputs():
input_name.append(node.name)
return input_name
def get_output_name(self):
output_name = []
for node in self.onnx_session.get_outputs():
output_name.append(node.name)
return output_name
def get_input_feed(self, img_tensor):
input_feed = {}
for name in self.input_name:
input_feed[name] = img_tensor
return input_feed
def inference(self, img):
or_img = cv2.resize(img, (640, 640))
img = or_img[:, :, ::-1].transpose(2, 0, 1) # BGR2RGB和HWC2CHW
img = img.astype(dtype=np.float32)
img /= 255.0
img = np.expand_dims(img, axis=0)
input_feed = self.get_input_feed(img)
pred = self.onnx_session.run(None, input_feed)[0]
return pred, or_img
# dets: array [x,6] 6个值分别为x1,y1,x2,y2,score,class
# thresh: 阈值
def nms(dets, thresh):
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
scores = dets[:, 4]
keep = []
index = scores.argsort()[::-1]
while index.size > 0:
i = index[0]
keep.append(i)
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1)
h = np.maximum(0, y22 - y11 + 1)
overlaps = w * h
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
idx = np.where(ious <= thresh)[0]
index = index[idx + 1]
return keep
def xywh2xyxy(x):
y = np.copy(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2
y[:, 1] = x[:, 1] - x[:, 3] / 2
y[:, 2] = x[:, 0] + x[:, 2] / 2
y[:, 3] = x[:, 1] + x[:, 3] / 2
return y
def filter_box(org_box, conf_thres, iou_thres):
org_box = np.squeeze(org_box)
conf = org_box[..., 4] > conf_thres
box = org_box[conf == True]
cls_cinf = box[..., 5:]
cls = []
for i in range(len(cls_cinf)):
cls.append(int(np.argmax(cls_cinf[i])))
all_cls = list(set(cls))
output = []
for i in range(len(all_cls)):
curr_cls = all_cls[i]
curr_cls_box = []
curr_out_box = []
for j in range(len(cls)):
if cls[j] = = curr_cls:
box[j][5] = curr_cls
curr_cls_box.append(box[j][:6])
curr_cls_box = np.array(curr_cls_box)
curr_cls_box = xywh2xyxy(curr_cls_box)
curr_out_box = nms(curr_cls_box, iou_thres)
for k in curr_out_box:
output.append(curr_cls_box[k])
output = np.array(output)
return output
def draw(image, box_data):
boxes = box_data[..., :4].astype(np.int32)
scores = box_data[..., 4]
classes = box_data[..., 5].astype(np.int32)
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
print('class: {}, score: {}'.format(CLASSES[cl], score))
print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
(top, left),
cv2.FONT_HERSHEY_SIMPLEX,
0.6, (0, 0, 255), 2)
def write_to_excel(data, filename):
with xlsxwriter.Workbook(filename) as workbook:
worksheet = workbook.add_worksheet()
worksheet.write('A1', '序号')
worksheet.write('B1', '类别')
worksheet.write('C1', '得分')
worksheet.write('D1', 'x')
worksheet.write('D1', 'y')
worksheet.write('E1', 'w')
worksheet.write('F1', 'h')
for i, box in enumerate(data):
x1, y1, x2, y2, score, cls = box
x_center = (x1 + x2) / 2
y_center = (y1 + y2) / 2
width = x2 - x1
height = y2 - y1
worksheet.write(i + 2, 0, i) # 序号
worksheet.write(i + 2, 1, CLASSES[cls]) # 类别
worksheet.write(i + 2, 2, score) # 得分
worksheet.write(i + 2, 3, x_center) # x中心点
worksheet.write(i + 2, 4, y_center) # y中心点
worksheet.write(i + 2, 5, width) # 宽度
worksheet.write(i + 2, 6, height) # 高度
if __name__ == "__main__":
onnx_path = 'C:\\yolov5-7.0\\runs\\train\\exp42\\weights\\best.onnx'
model = YOLOV5(onnx_path)
output_folder = 'oyt'
excel_filename = 'result.xlsx'
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# 加载预训练的人脸识别分类器
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
cap = cv2.VideoCapture(0)
excel_data = []
try:
while True:
ret, frame = cap.read()
if not ret:
break
output, or_img = model.inference(frame)
outbox = filter_box(output, 0.5, 0.5)
if len(outbox) > 0:
draw(or_img, outbox)
cv2.imshow('YOLOv5 Detection', or_img)
excel_data.extend(outbox)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
except Exception as e:
print(f"程序出现错误: {e}")
finally:
cap.release()
cv2.destroyAllWindows()
write_to_excel(excel_data, excel_filename)
标签:box,调用,name,img,write,opencv,output,摄像头,cls
From: https://blog.csdn.net/jiebinchen/article/details/142641573