要调用已经训练好的YOLOv3模型对图片进行检测,需要完成以下几个步骤:
- 加载预训练模型:从预训练的权重文件中加载模型。
- 准备输入图片:将图片转换为模型所需的格式。
- 进行推理:使用模型对图片进行推理,得到检测结果。
- 处理输出结果:解析模型的输出,得到检测框、类别和置信度。
以下是一个使用PyTorch和OpenCV的示例代码,展示如何调用已经训练好的YOLOv3模型对图片进行检测:
1. 安装必要的库
确保已经安装了以下库:
pip install torch torchvision opencv-python
2. 加载预训练模型
假设已经有一个预训练的YOLOv3模型权重文件 yolov3.weights
和对应的配置文件 yolov3.cfg
。
import torch
import cv2
import numpy as np# 加载预训练模型
model = cv2.dnn.readNetFromDarknet("yolov3.cfg", "yolov3.weights")
model.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
3. 准备输入图片
读取图片并将其转换为模型所需的格式。
# 读取图片
image = cv2.imread("test.jpg")
blob = cv2.dnn.blobFromImage(image, 1/255.0, (416, 416), swapRB=True, crop=False)
model.setInput(blob)
4. 进行推理
使用模型对图片进行推理,得到检测结果。
# 获取输出层的名称
layer_names = model.getLayerNames()
output_layers = [layer_names[i[0] - 1] for i in model.getUnconnectedOutLayers()]# 进行推理
outputs = model.forward(output_layers)
5. 处理输出结果
解析模型的输出,得到检测框、类别和置信度,并绘制检测结果。
class_ids = []
confidences = []
boxes = []
conf_threshold = 0.5
nms_threshold = 0.4# 解析输出
for output in outputs:for detection in output:scores = detection[5:]class_id = np.argmax(scores)confidence = scores[class_id]if confidence > conf_threshold:center_x = int(detection[0] * image.shape[1])center_y = int(detection[1] * image.shape[0])width = int(detection[2] * image.shape[1])height = int(detection[3] * image.shape[0])left = int(center_x - width / 2)top = int(center_y - height / 2)class_ids.append(class_id)confidences.append(float(confidence))boxes.append([left, top, width, height])# 非极大值抑制
indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)# 绘制检测结果
for i in indices:i = i[0]box = boxes[i]left = box[0]top = box[1]width = box[2]height = box[3]cv2.rectangle(image, (left, top), (left + width, top + height), (0, 255, 0), 2)label = f"{class_ids[i]} {confidences[i]:.2f}"cv2.putText(image, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)# 显示结果图片
cv2.imshow("Detection", image)
cv2.waitKey(0)
cv2.destroyAllWindows()