基于yolov7的危险物品检测代码

这里是一个基于YOLOv5的危险物品检测代码示例，你可以根据自己的需求进行修改。

import cv2
import numpy as np
import os

# 加载模型
weights_path = "path/to/weights"
config_path = "path/to/config"
net = cv2.dnn.readNetFromDarknet(config_path, weights_path)

# 加载类别标签
labels = ["dangerous_object"]

# 设置阈值和NMS抑制阈值
conf_threshold = 0.5
nms_threshold = 0.4

# 获取输入图像的尺寸
input_shape = net.getLayer(0).outputShape[1:3]

# 读取测试图像
img_path = "path/to/image"
img = cv2.imread(img_path)

# 将图像缩放到网络输入尺寸
resized_img = cv2.resize(img, input_shape)

# 将图像转换为blob格式
blob = cv2.dnn.blobFromImage(resized_img, 1 / 255.0, input_shape, swapRB=True, crop=False)

# 将blob输入到神经网络中
net.setInput(blob)

# 进行前向传递
layer_names = net.getLayerNames()
output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
outputs = net.forward(output_layers)

# 解析输出并绘制边界框
boxes = []
confidences = []
class_ids = []

for output in outputs:
    for detection in output:
        scores = detection[5:]
        class_id = np.argmax(scores)
        confidence = scores[class_id]
        if confidence > conf_threshold and class_id == 0:
            # 将边界框坐标缩放回原图尺寸
            center_x = int(detection[0] * img.shape[1])
            center_y = int(detection[1] * img.shape[0])
            width = int(detection[2] * img.shape[1])
            height = int(detection[3] * img.shape[0])
            left = int(center_x - width / 2)
            top = int(center_y - height / 2)

            boxes.append([left, top, width, height])
            confidences.append(float(confidence))
            class_ids.append(class_id)

# 应用NMS抑制
indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

# 绘制边界框和类别标签
font = cv2.FONT_HERSHEY_PLAIN
for i in indices:
    i = i[0]
    box = boxes[i]
    left = box[0]
    top = box[1]
    width = box[2]
    height = box[3]

    cv2.rectangle(img, (left, top), (left + width, top + height), (0, 255, 0), 2)
    cv2.putText(img, labels[class_ids[i]], (left, top - 5), font, 1, (0, 255, 0), 2)

# 显示结果图像
cv2.imshow("Result", img)
cv2.waitKey(0)

在这个例子中，我们使用了OpenCV的`dnn`模块来加载YOLOv5模型并进行危险物品检测。我们还使用了NMS抑制来消除重复的检测结果。你需要根据自己的需求修改模型路径、类别标签、阈值等参数。