nx上darknet的使用-目标检测-在python中的使用

1  内置的代码

在darknet中已经内置了两个py文件

darknet_video.py与darknet_images.py用法类似，都是改一改给的参数就行了，我们说一下几个关键的参数

• input 要预测哪张图像
• weights 要使用哪个权重
• config_file 要使用哪个cfg文件
• data_file 要使用哪个data文件
• thresh 置信度给多少

改完上面这些参数就可以直接运行了

用鼠标点一下图片然后按q可以关闭图片

调用视频有时可以正常用，有时不可以，会报下面的错，可能是opencv的版本问题

这两个代码都很好，想的很周道，但是看起来太麻烦，我通常用下面两个简化版的

2  opencv使用模型

优点：只要有python，有opencv-python，有numpy就能用（不需要编译darknet，只要有names、cfg、weights就行）

缺点：只能搞CPU版，用摄像头搞效果较差

2.1  识别图像

你需要一张预测图像，训练时的classes.names yolov4-tiny-cfg yolov4-tiny_final.weights

yolo_opencv.py的内容如下

import cv2
import numpy as npLABELS = open("classes.names").read().strip().split("\n")
net = cv2.dnn.readNetFromDarknet('yolov4-tiny.cfg', 'yolov4-tiny_final.weights')
layer = net.getUnconnectedOutLayersNames()frame = cv2.imread('000003.jpg')
(H, W) = frame.shape[:2]
blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416),swapRB=True, crop=False)
net.setInput(blob)
layerOutputs = net.forward(layer)
boxes = []
confidences = []
classIDs = []for output in layerOutputs:for detection in output:scores = detection[5:]classID = np.argmax(scores)confidence = scores[classID]box = detection[0:4] * np.array([W, H, W, H])(centerX, centerY, width, height) = box.astype("int")x = int(centerX - (width / 2))y = int(centerY - (height / 2))boxes.append([x, y, int(width), int(height)])confidences.append(float(confidence))classIDs.append(classID)idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.3)
if len(idxs) > 0:for i in idxs.flatten():(x, y) = (boxes[i][0], boxes[i][1])(w, h) = (boxes[i][2], boxes[i][3])cv2.rectangle(frame, (x, y), (x + w, y + h), (0,255,0), 1, lineType=cv2.LINE_AA)text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i])cv2.putText(frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,0.5, (255,0,0), 1, lineType=cv2.LINE_AA)cv2.imshow('frame',frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

运行后可以显示预测结果

2.2  flask起服务

我们可以搞一个接口来处理识别的功能

服务端

import numpy as np
import cv2
from flask import Flask,request
import base64LABELS = open("classes.names").read().strip().split("\n")
net = cv2.dnn.readNetFromDarknet('yolov4-tiny.cfg', 'yolov4-tiny_final.weights')
layer = net.getUnconnectedOutLayersNames()app = Flask(__name__)
@app.route('/predict',methods=["POST"])
def predict():if request.method == 'POST':image_base64 = request.json['img_b64']img = base64.b64decode(image_base64)img = np.fromstring(img,np.uint8)frame = cv2.imdecode(img,cv2.IMREAD_COLOR)(H, W) = frame.shape[:2]blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416),swapRB=True, crop=False)net.setInput(blob)layerOutputs = net.forward(layer)boxes = []confidences = []classIDs = []for output in layerOutputs:for detection in output:scores = detection[5:]classID = np.argmax(scores)confidence = scores[classID]box = detection[0:4] * np.array([W, H, W, H])(centerX, centerY, width, height) = box.astype("int")x = int(centerX - (width / 2))y = int(centerY - (height / 2))boxes.append([x, y, int(width), int(height)])confidences.append(float(confidence))classIDs.append(classID)idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.3)result = []if len(idxs) > 0:for i in idxs.flatten():(x, y) = (boxes[i][0], boxes[i][1])(w, h) = (boxes[i][2], boxes[i][3])result_obj = {}result_obj['x'] = xresult_obj['y'] = yresult_obj['w'] = wresult_obj['h'] = hresult_obj['name'] = LABELS[classIDs[i]]result_obj['confidence'] = confidences[i]result.append(result_obj)return {'result':result}if __name__ == '__main__':app.run(host='192.168.0.105')

客户端

import cv2
import requests
import base64def get_result(url,frame):retval, buffer = cv2.imencode('.jpg', frame)image = str(base64.b64encode(buffer), 'utf-8')json_data = {'img_b64': image}response = eval(requests.post(url, json=json_data).text).get('result')return response
if __name__ == '__main__':url = 'http://192.168.0.105:5000/predict'frame = cv2.imread('000003.jpg')response = get_result(url,frame)for result in response:confidence = result.get('confidence')x = int(result.get('x'))y = int(result.get('y'))w = int(result.get('w'))h = int(result.get('h'))name = result.get('name')frame = cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)frame = cv2.putText(frame, name, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)cv2.imshow('img',frame)cv2.waitKey(0)cv2.destroyAllWindows()

请求结果

2.3  摄像头识别

原理是图像识别+多线程

import cv2
import numpy as np
import queue
import threading
import timeLABELS = open("classes.names").read().strip().split("\n")
net = cv2.dnn.readNetFromDarknet('yolov4-tiny.cfg', 'yolov4-tiny_final.weights')
layer = net.getUnconnectedOutLayersNames()cap = cv2.VideoCapture(0)frame_queue = queue.Queue()
detection_result_queue = queue.Queue(maxsize=1)def video_capture():while cap.isOpened():ret, frame = cap.read()if ret:frame_queue.put(frame)cap.release()def predict():while cap.isOpened():start_time = time.time()predict_frame = frame_queue.get()(H, W) = predict_frame.shape[:2]blob = cv2.dnn.blobFromImage(predict_frame, 1 / 255.0, (416, 416), swapRB=True, crop=False)net.setInput(blob)layerOutputs = net.forward(layer)boxes = []confidences = []classIDs = []for output in layerOutputs:for detection in output:scores = detection[5:]classID = np.argmax(scores)confidence = scores[classID]box = detection[0:4] * np.array([W, H, W, H])(centerX, centerY, width, height) = box.astype("int")x = int(centerX - (width / 2))y = int(centerY - (height / 2))boxes.append([x, y, int(width), int(height)])confidences.append(float(confidence))classIDs.append(classID)idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.3)if len(idxs) > 0:result_list = []for i in idxs.flatten():result_dic = {}(x, y) = (boxes[i][0], boxes[i][1])(w, h) = (boxes[i][2], boxes[i][3])label = LABELS[classIDs[i]]confidence = confidences[i]result_dic.__setitem__('x',x)result_dic.__setitem__('y',y)result_dic.__setitem__('w',w)result_dic.__setitem__('h',h)result_dic.__setitem__('label',label)result_dic.__setitem__('confidence',confidence)result_list.append(result_dic)detection_result_queue.put(result_list)print(time.time()-start_time)cap.release()def draw():while cap.isOpened():draw_frame = frame_queue.get()try:predict_results = detection_result_queue.get(block=False)for predict_result in predict_results:x = predict_result.get('x')y = predict_result.get('y')w = predict_result.get('w')h = predict_result.get('h')label = predict_result.get('label')confidence = predict_result.get('confidence')cv2.rectangle(draw_frame, (x, y), (x + w, y + h), (0,255,0), 1, lineType=cv2.LINE_AA)text = "{}: {:.4f}".format(label, confidence)cv2.putText(draw_frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,0.5, (255,0,0), 1, lineType=cv2.LINE_AA)except:passcv2.imshow('draw_frame',draw_frame)cv2.waitKey(1)cap.release()threading.Thread(target=video_capture).start()
threading.Thread(target=predict).start()
threading.Thread(target=draw).start()

由于是CPU运行，效果根据CPU的性能会有差异，我是用NX上的6核CPU运行yolov4-tiny，大约0.5s一张。这个速度如果在摄像头上的连续识别就非常慢了，至少要1秒30帧才能有流畅的感觉，也就是0.03秒需要预测一张

• 0.5/0.03 = 16.6 或许搞个17线程能够达成流畅的效果，没尝试过

3  darknet使用模型

优点：可以使用GPU

缺点：需要编译darknet才能使用

下面在代码中import darknet指的是import下面这个文件。下面这个py文件依赖了darknet文件夹下的其他文件，所以建议把下面的代码放在darknet的根目录下使用

3.1  识别图像

在模型读取的时候是通过data文件找names文件，而不是直接找names文件

import cv2
import darknet
import timenetwork, class_names, class_colors = darknet.load_network('/home/suyu/darknet/yolo_opencv/yolov7-tiny.cfg','/home/suyu/darknet/custom_training/custom_training.data','/home/suyu/darknet/yolo_opencv/yolov7-tiny_final_origin.weights',batch_size=1
)width = darknet.network_width(network)
height = darknet.network_height(network)
darknet_image = darknet.make_image(width, height, 3)image = cv2.imread('/home/suyu/darknet/yolo_opencv/280.jpg')
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_resized = cv2.resize(image_rgb, (width, height),interpolation=cv2.INTER_LINEAR)
darknet.copy_image_from_bytes(darknet_image, image_resized.tobytes())
start_time = time.time()
detections = darknet.detect_image(network, class_names, darknet_image, thresh=0.8)
print(time.time()-start_time)darknet.free_image(darknet_image)
image = darknet.draw_boxes(detections, image_resized, class_colors)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)cv2.imshow('image',image)
cv2.waitKey(0)
cv2.destroyAllWindows()

在终端上可以看到对模型进行了读取

画框的颜色每次都是随机的

在预测中，如果出现了下面这种情况，我们就需要减少预测的结果，可以通过标签后面的置信度减少，但是我们可以看到有两个框的置信度都在0.99以上，这个时候我们就需要用到其他减少框的办法

darknet.detect_image()这个方法，除了通过置信度减少预测结果，还可以通过非最大值抑制消除冗余，也就是最后一个的nms

• hier_thresh也是消除冗余的一个参数，叫控制层次性阈值(与nms类似)，没用过，如果nms不行的化再尝试使用它

nms数值越小，消除冗余的效果就越好，我这里直接改成了0.05

改完之后可以得到还不错的效果

3.2  摄像头识别

原理是识别图像+多线程，我用yolov7-tiny大概能到0.03秒预测一张，这个预测速度不加多线程应该也可以

import cv2
import numpy as np
import queue
import threading
import time
import darknetnetwork, class_names, class_colors = darknet.load_network('/home/suyu/darknet/yolo_opencv/yolov7-tiny.cfg','/home/suyu/darknet/custom_training/custom_training.data','/home/suyu/darknet/yolo_opencv/yolov7-tiny_final_origin.weights',batch_size=1
)width = darknet.network_width(network)
height = darknet.network_height(network)
darknet_image = darknet.make_image(width, height, 3)cap = cv2.VideoCapture(0)frame_queue = queue.Queue()
detection_result_queue = queue.Queue(maxsize=1)def video_capture():while cap.isOpened():ret, frame = cap.read()if ret:image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)image_resized = cv2.resize(image_rgb, (width, height), interpolation=cv2.INTER_LINEAR)frame_queue.put(image_resized)cap.release()def predict():while cap.isOpened():start_time = time.time()image_resized = frame_queue.get()darknet.copy_image_from_bytes(darknet_image, image_resized.tobytes())detections = darknet.detect_image(network, class_names, darknet_image, thresh=0.8)# darknet.free_image(darknet_image)print(time.time()-start_time)detection_result_queue.put(detections)cap.release()def draw():while cap.isOpened():draw_frame = frame_queue.get()# print(draw_frame)try:detections = detection_result_queue.get(block=False)draw_frame = darknet.draw_boxes(detections, draw_frame, class_colors)except:passdraw_frame = cv2.cvtColor(draw_frame, cv2.COLOR_BGR2RGB)cv2.imshow('draw_frame',draw_frame)cv2.waitKey(1)cap.release()threading.Thread(target=video_capture).start()
threading.Thread(target=predict).start()
threading.Thread(target=draw).start()

如果你想对预测的结果进行别的操作你可以用到detection_result_queue中的detections，打印出来是这样的

列表套元组，元组的第一个值是label，第二个值是置信

操作的时候建议多开一个线程，然后将detection_result_queue置为2或者更高，避免两个线程抢数据的情况

4  onnx使用模型

onnx(Open Neural Network eXchange) 开放式神经网络交换，好多模型都可以转换为onnx类型的模型，相当于是人工智能模型界的docker了。以高适配性而著名。在性能上并没有优于其他模型。

4.1  用到的库

onnx，这个是将其他类型的模型转换为onnx类型的库，在arm端上安装可能会有些麻烦，但在amd上直接用pip就可以安装了

onnxruntime是跑onnx模型用的，onnxruntime在arm端可以直接用pip安装

4.2  darknet的weight转onnx

参考 https://zhuanlan.zhihu.com/p/543345367

源码 GitHub - Tianxiaomo/pytorch-YOLOv4: PyTorch ,ONNX and TensorRT implementation of YOLOv4

安装完onnx与onnxruntime后运行demo_darknet2onnx.py，第一个参数是cfg，第二个参数是names，第三个参数是weights，第四个参数是图，第五个参数是batch_size，直接写1就行了

python demo_darknet2onnx.py /home/suyu/darknet/cfg/yolov4.cfg /home/suyu/darknet/data/coco.names /home/suyu/darknet/yolov4.weights /home/suyu/darknet/data/dog.jpg 1

如果opencv版本过高会爆出下面两个问题，是cv2.rectangele()与cv2.putText()的参数需为int

我们需要更改 /pytorch-YOLOv4-master/tool/utils.py

把画红线的地方改成int

成功使用会显示下面这些东西

在文件夹中可以找到通过onnx预测成功的图像和onnx模型

4.3  识别图像

识别的代码是用的上面提供的源码中的东西，有的代码直接从utils中复制过来了，用的时候不需要引入其他py文件了，只需要.onnx文件与.names文件

import sys
import onnx
import os
import argparse
import numpy as np
import cv2
import onnxruntime
import time
import mathdef load_class_names(namesfile):class_names = []with open(namesfile, 'r') as fp:lines = fp.readlines()for line in lines:line = line.rstrip()class_names.append(line)return class_namesdef nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False):# print(boxes.shape)x1 = boxes[:, 0]y1 = boxes[:, 1]x2 = boxes[:, 2]y2 = boxes[:, 3]areas = (x2 - x1) * (y2 - y1)order = confs.argsort()[::-1]keep = []while order.size > 0:idx_self = order[0]idx_other = order[1:]keep.append(idx_self)xx1 = np.maximum(x1[idx_self], x1[idx_other])yy1 = np.maximum(y1[idx_self], y1[idx_other])xx2 = np.minimum(x2[idx_self], x2[idx_other])yy2 = np.minimum(y2[idx_self], y2[idx_other])w = np.maximum(0.0, xx2 - xx1)h = np.maximum(0.0, yy2 - yy1)inter = w * hif min_mode:over = inter / np.minimum(areas[order[0]], areas[order[1:]])else:over = inter / (areas[order[0]] + areas[order[1:]] - inter)inds = np.where(over <= nms_thresh)[0]order = order[inds + 1]return np.array(keep)def post_processing(img, conf_thresh, nms_thresh, output):# anchors = [12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401]# num_anchors = 9# anchor_masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]# strides = [8, 16, 32]# anchor_step = len(anchors) // num_anchors# [batch, num, 1, 4]box_array = output[0]# [batch, num, num_classes]confs = output[1]t1 = time.time()if type(box_array).__name__ != 'ndarray':box_array = box_array.cpu().detach().numpy()confs = confs.cpu().detach().numpy()num_classes = confs.shape[2]# [batch, num, 4]box_array = box_array[:, :, 0]# [batch, num, num_classes] --> [batch, num]max_conf = np.max(confs, axis=2)max_id = np.argmax(confs, axis=2)t2 = time.time()bboxes_batch = []for i in range(box_array.shape[0]):argwhere = max_conf[i] > conf_threshl_box_array = box_array[i, argwhere, :]l_max_conf = max_conf[i, argwhere]l_max_id = max_id[i, argwhere]bboxes = []# nms for each classfor j in range(num_classes):cls_argwhere = l_max_id == jll_box_array = l_box_array[cls_argwhere, :]ll_max_conf = l_max_conf[cls_argwhere]ll_max_id = l_max_id[cls_argwhere]keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)if (keep.size > 0):ll_box_array = ll_box_array[keep, :]ll_max_conf = ll_max_conf[keep]ll_max_id = ll_max_id[keep]for k in range(ll_box_array.shape[0]):bboxes.append([ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k], ll_max_conf[k], ll_max_id[k]])bboxes_batch.append(bboxes)t3 = time.time()print('-----------------------------------')print('       max and argmax : %f' % (t2 - t1))print('                  nms : %f' % (t3 - t2))print('Post processing total : %f' % (t3 - t1))print('-----------------------------------')return bboxes_batchdef plot_boxes_cv2(img, boxes, savename=None, class_names=None, color=None):import cv2img = np.copy(img)colors = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]], dtype=np.float32)def get_color(c, x, max_val):ratio = float(x) / max_val * 5i = int(math.floor(ratio))j = int(math.ceil(ratio))ratio = ratio - ir = (1 - ratio) * colors[i][c] + ratio * colors[j][c]return int(r * 255)width = img.shape[1]height = img.shape[0]for i in range(len(boxes)):box = boxes[i]x1 = int(box[0] * width)y1 = int(box[1] * height)x2 = int(box[2] * width)y2 = int(box[3] * height)bbox_thick = int(0.6 * (height + width) / 600)if color:rgb = colorelse:rgb = (255, 0, 0)if len(box) >= 7 and class_names:cls_conf = box[5]cls_id = box[6]print('%s: %f' % (class_names[cls_id], cls_conf))classes = len(class_names)offset = cls_id * 123457 % classesred = get_color(2, offset, classes)green = get_color(1, offset, classes)blue = get_color(0, offset, classes)if color is None:rgb = (red, green, blue)msg = str(class_names[cls_id])+" "+str(round(cls_conf,3))t_size = cv2.getTextSize(msg, 0, 0.7, thickness=bbox_thick // 2)[0]c1, c2 = (x1,y1), (x2, y2)c3 = (c1[0] + t_size[0], c1[1] - t_size[1] - 3)cv2.rectangle(img, (int(x1),int(y1)), (int(np.float32(c3[0])), int(np.float32(c3[1]))), rgb, -1)img = cv2.putText(img, msg, (int(c1[0]), int(np.float32(c1[1] - 2))), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,0,0), bbox_thick//2,lineType=cv2.LINE_AA)img = cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), rgb, bbox_thick)if savename:print("save plot results to %s" % savename)cv2.imwrite(savename, img)return imgdef detect(session, image_src, namesfile):IN_IMAGE_H = session.get_inputs()[0].shape[2]IN_IMAGE_W = session.get_inputs()[0].shape[3]# Inputresized = cv2.resize(image_src, (IN_IMAGE_W, IN_IMAGE_H), interpolation=cv2.INTER_LINEAR)img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)img_in = np.expand_dims(img_in, axis=0)img_in /= 255.0#print("Shape of the network input: ", img_in.shape)# Computeinput_name = session.get_inputs()[0].nameoutputs = session.run(None, {input_name: img_in})boxes = post_processing(img_in, 0.4, 0.6, outputs)class_names = load_class_names(namesfile)return plot_boxes_cv2(image_src, boxes[0], class_names=class_names)if __name__ == '__main__':session = onnxruntime.InferenceSession('yolov4_1_3_640_640_static.onnx')namesfile = 'classes.names'image_src = cv2.imread('3.png')detected_img = detect(session, image_src, namesfile)cv2.imshow('detected_img',detected_img)cv2.waitKey(0)cv2.destroyAllWindows()

4.4  摄像头识别

用的采集、预测、画 三线程。模型输入width与height都为320，两分类，在树莓派4B中大概是0.5s一张，大概是1秒2个识别帧

import cv2
import numpy as np
import queue
import threading
import timeimport sys
import onnx
import os
import argparse
import onnxruntime
import mathdef load_class_names(namesfile):class_names = []with open(namesfile, 'r') as fp:lines = fp.readlines()for line in lines:line = line.rstrip()class_names.append(line)return class_namesdef nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False):x1 = boxes[:, 0]y1 = boxes[:, 1]x2 = boxes[:, 2]y2 = boxes[:, 3]areas = (x2 - x1) * (y2 - y1)order = confs.argsort()[::-1]keep = []while order.size > 0:idx_self = order[0]idx_other = order[1:]keep.append(idx_self)xx1 = np.maximum(x1[idx_self], x1[idx_other])yy1 = np.maximum(y1[idx_self], y1[idx_other])xx2 = np.minimum(x2[idx_self], x2[idx_other])yy2 = np.minimum(y2[idx_self], y2[idx_other])w = np.maximum(0.0, xx2 - xx1)h = np.maximum(0.0, yy2 - yy1)inter = w * hif min_mode:over = inter / np.minimum(areas[order[0]], areas[order[1:]])else:over = inter / (areas[order[0]] + areas[order[1:]] - inter)inds = np.where(over <= nms_thresh)[0]order = order[inds + 1]return np.array(keep)def post_processing(img, conf_thresh, nms_thresh, output):box_array = output[0]confs = output[1]if type(box_array).__name__ != 'ndarray':box_array = box_array.cpu().detach().numpy()confs = confs.cpu().detach().numpy()num_classes = confs.shape[2]box_array = box_array[:, :, 0]max_conf = np.max(confs, axis=2)max_id = np.argmax(confs, axis=2)bboxes_batch = []for i in range(box_array.shape[0]):argwhere = max_conf[i] > conf_threshl_box_array = box_array[i, argwhere, :]l_max_conf = max_conf[i, argwhere]l_max_id = max_id[i, argwhere]bboxes = []for j in range(num_classes):cls_argwhere = l_max_id == jll_box_array = l_box_array[cls_argwhere, :]ll_max_conf = l_max_conf[cls_argwhere]ll_max_id = l_max_id[cls_argwhere]keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)if (keep.size > 0):ll_box_array = ll_box_array[keep, :]ll_max_conf = ll_max_conf[keep]ll_max_id = ll_max_id[keep]for k in range(ll_box_array.shape[0]):bboxes.append([ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k], ll_max_conf[k], ll_max_id[k]])bboxes_batch.append(bboxes)return bboxes_batchdef get_color(c, x, max_val):colors = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]], dtype=np.float32)ratio = float(x) / max_val * 5i = int(math.floor(ratio))j = int(math.ceil(ratio))ratio = ratio - ir = (1 - ratio) * colors[i][c] + ratio * colors[j][c]return int(r * 255)def detect(session, image_src, namesfile):IN_IMAGE_H = session.get_inputs()[0].shape[2]IN_IMAGE_W = session.get_inputs()[0].shape[3]# Inputresized = cv2.resize(image_src, (IN_IMAGE_W, IN_IMAGE_H), interpolation=cv2.INTER_LINEAR)img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)img_in = np.expand_dims(img_in, axis=0)img_in /= 255.0# Computeinput_name = session.get_inputs()[0].nameoutputs = session.run(None, {input_name: img_in})boxes = post_processing(img_in, 0.4, 0.6, outputs)return boxes[0]session = onnxruntime.InferenceSession('yolov4_1_3_640_640_static.onnx')
namesfile = 'classes.names'
class_names = load_class_names(namesfile)cap = cv2.VideoCapture(0)frame_queue = queue.Queue()
detection_result_queue = queue.Queue(maxsize=1)def video_capture():while cap.isOpened():ret, frame = cap.read()if ret:frame_queue.put(frame)cap.release()def predict():while cap.isOpened():start_time = time.time()predict_frame = frame_queue.get()result_list = detect(session, predict_frame, namesfile)detection_result_queue.put(result_list)print(time.time()-start_time)cap.release()def draw():while cap.isOpened():draw_frame = frame_queue.get()try:boxes = detection_result_queue.get(block=False)img = draw_framewidth = img.shape[1]height = img.shape[0]for i in range(len(boxes)):box = boxes[i]x1 = int(box[0] * width)y1 = int(box[1] * height)x2 = int(box[2] * width)y2 = int(box[3] * height)bbox_thick = int(0.6 * (height + width) / 600)rgb = (255, 0, 0)if len(box) >= 7 and class_names:cls_conf = box[5]cls_id = box[6]print('%s: %f' % (class_names[cls_id], cls_conf))classes = len(class_names)offset = cls_id * 123457 % classesred = get_color(2, offset, classes)green = get_color(1, offset, classes)blue = get_color(0, offset, classes)rgb = (red, green, blue)msg = str(class_names[cls_id])+" "+str(round(cls_conf,3))t_size = cv2.getTextSize(msg, 0, 0.7, thickness=bbox_thick // 2)[0]c1, c2 = (x1,y1), (x2, y2)c3 = (c1[0] + t_size[0], c1[1] - t_size[1] - 3)cv2.rectangle(img, (int(x1),int(y1)), (int(np.float32(c3[0])), int(np.float32(c3[1]))), rgb, -1)img = cv2.putText(img, msg, (int(c1[0]), int(np.float32(c1[1] - 2))), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,0,0), bbox_thick//2,lineType=cv2.LINE_AA)img = cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), rgb, bbox_thick)draw_frame = img#except Exception as e:#print(e)except:passcv2.imshow('draw_frame',draw_frame)cv2.waitKey(1)cap.release()threading.Thread(target=video_capture).start()
threading.Thread(target=predict).start()
threading.Thread(target=draw).start()

在树莓派4B上是下面这个效果