要调用已经训练好的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()